Fix failing Travis build for tvt/run.sh

* Restrict the script to a single episode to avoid timeouts
* Ensure that `wheel` is up to date (we may need to build our own wheels for some of the requirements)

PiperOrigin-RevId: 367980522

.travis.yml

@@ -8,18 +8,19 @@ python:
   - "3.6"
 env:
   matrix:
-    - PROJECT="tvt"
+    - PROJECT="avae"
     - PROJECT="cs_gan"
     - PROJECT="iodine"
-    - PROJECT="transporter"
-    - PROJECT="polygen"
-    - PROJECT="learning_to_simulate"
-    - PROJECT="avae"
     - PROJECT="kfac_ferminet_alpha"
+    - PROJECT="learning_to_simulate"
+    - PROJECT="polygen"
+    - PROJECT="transporter"
+    - PROJECT="tvt"
 before_script:
   - sudo apt-get update -qq
   - pip install --upgrade pip
   - pip install --upgrade virtualenv
+  - pip install --upgrade wheel
 script:
   - ./$PROJECT/run.sh
 

tvt/run.sh

@@ -13,8 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-python3 -m venv tvt_venv
-source tvt_venv/bin/activate
+python3 -m venv /tmp/tvt_venv
+source /tmp/tvt_venv/bin/activate
+pip install -U pip
 pip install -r tvt/requirements.txt
 
-python3 -m tvt.main
+python3 -m tvt.main --num_episodes=1