Sample expert data from same distribution

scratch/etienne/intersimple/data/expert.py

@@ -2,10 +2,10 @@ from intersim.envs.intersimple import Intersimple, InfoFilter
 from stable_baselines3.common.policies import BasePolicy
 import gym
 from intersim.envs.intersimple import *
+from gail.envs import *
 import imitation.data.rollout as rollout
 from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
 from imitation.data.wrappers import RolloutInfoWrapper
-from gail.envs import NRasterizedRouteSpeedRandomAgentLocation
 
 class IntersimExpert(BasePolicy):
     

scratch/etienne/intersimple/data/generate.sh

@@ -10,4 +10,5 @@
 #python -m expert --env=NRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_NRasterizedRouteRandomAgentLocationw70h70mppx1.pkl'
 #python -m expert --env=NRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1,map_color:128}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_NRasterizedRouteRandomAgentLocationw70h70mppx1mapc128.pkl'
 #python -m expert --env=NRasterizedRouteSpeedRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,map_color:128,mu:0.001}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001.pkl'
-python -m data.expert --env=NRasterizedRouteSpeedRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,map_color:128,mu:0.001,skip_frames:5}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001skip5.pkl'
+#python -m data.expert --env=NRasterizedRouteSpeedRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,map_color:128,mu:0.001,skip_frames:5}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001skip5.pkl'
+python -m data.expert --env=TLNRasterizedRouteRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,mu:0.001,random_skip:True,max_episode_steps:50}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl'

scratch/etienne/intersimple/gail/envs.py

@@ -1,6 +1,7 @@
 import gym
+from gym.wrappers.time_limit import TimeLimit
 import numpy as np
-from intersim.envs.intersimple import RandomLocation, RandomAgent, RewardVisualization, Reward, \
+from intersim.envs.intersimple import NRasterizedRouteRandomAgentLocation, RandomLocation, RandomAgent, RewardVisualization, Reward, \
     ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedObservation, \
     NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple
 
@@ -31,4 +32,12 @@ class RasterizedSpeed:
 class NRasterizedRouteSpeedRandomAgentLocation(RandomLocation, RandomAgent, RewardVisualization,
         Reward, ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedSpeed, RasterizedObservation,
         NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple):
-    pass
+    pass
+
+class NoBSTimeLimit(TimeLimit):
+
+    def __getattr__(self, name):
+        return getattr(self.env, name)
+
+def TLNRasterizedRouteRandomAgentLocation(max_episode_steps, *args, **kwargs):
+    return NoBSTimeLimit(NRasterizedRouteRandomAgentLocation(*args, **kwargs), max_episode_steps=max_episode_steps)

scratch/etienne/intersimple/gail_options_image_random_location.py

@@ -16,31 +16,32 @@ 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
+from gail.envs import TLNRasterizedRouteRandomAgentLocation
+from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
 
 model_name = 'gail_options_image_random_location'
-env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True}
+env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 50}
+env = TLNRasterizedRouteRandomAgentLocation(**env_settings)
 
 ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
 
 def train(
         expert_data,
-        expert_batch_size=4069,
+        expert_batch_size=2048,
         discriminator_updates_per_round=10,
         generator_steps=256,
-        generator_total_steps=2048,
+        generator_total_steps=1024,
         generator_updates_per_round=10,
         discount=0.99,
         epochs=100,
     ):
-    env = NRasterizedRouteRandomAgentLocation(**env_settings)
 
     tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
     tempdir_path = pathlib.Path(tempdir.name)
     logger.configure(tempdir_path / "GAIL/")
     print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
 
-    venv = make_vec_env(NRasterizedRouteRandomAgentLocation, n_envs=1, env_kwargs=env_settings)
+    venv = DummyVecEnv([lambda: env])
     discriminator = adversarial.GAIL(
         expert_data=expert_data,
         expert_batch_size=expert_batch_size,
@@ -50,13 +51,13 @@ def train(
         gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
     )
 
-    options_env = TimeLimit(OptionsEnv(
+    options_env = OptionsEnv(
         env,
         options=ALL_OPTIONS,
         discriminator=imitation_discriminator(discriminator),
         discount=discount,
         ll_buffer_capacity=expert_batch_size,
-    ), max_episode_steps=15)
+    )
     generator = stable_baselines3.PPO(
         OptionsCnnPolicy,
         options_env,
@@ -70,9 +71,9 @@ def train(
         generator.learn(total_timesteps=generator_total_steps)
 
         # train discriminator
-        generator_samples = options_env.sample_ll(expert_batch_size)
-        generator_samples = flatten_transitions(generator_samples)
         for _ in range(discriminator_updates_per_round):
+            generator_samples = options_env.sample_ll(expert_batch_size)
+            generator_samples = flatten_transitions(generator_samples)
             discriminator.train_disc(gen_samples=generator_samples)
 
         generator.save(model_name)
@@ -94,13 +95,12 @@ def video(model_name, env):
 def evaluate():
     video(
         model_name=model_name,
-        env=NRasterizedRouteRandomAgentLocation(**env_settings)
+        env=env
     )
 
 # %%
 if __name__ == '__main__':
-
+    with open("data/NormalizedIntersimpleExpertMu.001N10000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
-    with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001skip5.pkl", "rb") as f:
         trajectories = pickle.load(f)
     transitions = rollout.flatten_trajectories(trajectories)
     train(transitions)