From 6416fceb60067af6c4ad9f73bb8bffe95a4ad232 Mon Sep 17 00:00:00 2001
From: ebuehrle <43623224+ebuehrle@users.noreply.github.com>
Date: Tue, 2 Nov 2021 10:34:57 +0100
Subject: [PATCH] Fix horizon

---
 .../intersimple/gail_options_image_random_location.py     | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/scratch/etienne/intersimple/gail_options_image_random_location.py b/scratch/etienne/intersimple/gail_options_image_random_location.py
index fc4909b..f1d550d 100644
--- a/scratch/etienne/intersimple/gail_options_image_random_location.py
+++ b/scratch/etienne/intersimple/gail_options_image_random_location.py
@@ -15,8 +15,8 @@ from stable_baselines3.common.env_util import make_vec_env
 from tqdm import tqdm
 from src.policies.options import OptionsCnnPolicy
 from src.gail.train import flatten_transitions
-
 from gail.options2 import OptionsEnv, RenderOptions, imitation_discriminator
+from gym.wrappers import TimeLimit
 
 model_name = 'gail_options_image_random_location'
 env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'map_color': 128, 'mu': 0.001}
@@ -51,7 +51,11 @@ def train(
         gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
     )
 
-    options_env = OptionsEnv(env, discriminator=imitation_discriminator(discriminator), options=ALL_OPTIONS, ll_buffer_capacity=expert_batch_size)
+    options_env = TimeLimit(OptionsEnv(env,
+        discriminator=imitation_discriminator(discriminator),
+        options=ALL_OPTIONS,
+        ll_buffer_capacity=expert_batch_size
+    ), max_episode_steps=10)
     generator = stable_baselines3.PPO(
         OptionsCnnPolicy,
         options_env,