From 071c73192195e514740c39e5a410aa308064e84a Mon Sep 17 00:00:00 2001
From: ebuehrle <43623224+ebuehrle@users.noreply.github.com>
Date: Thu, 4 Nov 2021 19:54:40 +0100
Subject: [PATCH] More expert data

---
 scratch/etienne/intersimple/data/generate.sh                | 2 +-
 .../intersimple/gail_options_image_random_location.py       | 6 ++----
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/scratch/etienne/intersimple/data/generate.sh b/scratch/etienne/intersimple/data/generate.sh
index 307d5a9..b17f59b 100755
--- a/scratch/etienne/intersimple/data/generate.sh
+++ b/scratch/etienne/intersimple/data/generate.sh
@@ -11,4 +11,4 @@
 #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=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'
+python -m data.expert --env=TLNRasterizedRouteRandomAgentLocation --min_timesteps=100000 --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.001N100000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl'
diff --git a/scratch/etienne/intersimple/gail_options_image_random_location.py b/scratch/etienne/intersimple/gail_options_image_random_location.py
index 22a4fe5..9131ba3 100644
--- a/scratch/etienne/intersimple/gail_options_image_random_location.py
+++ b/scratch/etienne/intersimple/gail_options_image_random_location.py
@@ -5,13 +5,11 @@ sys.path.append('../../../')
 from src.discriminator import CnnDiscriminatorFlatAction
 from imitation.algorithms import adversarial
 import stable_baselines3
-from intersim.envs import NRasterizedRouteRandomAgentLocation
 import pickle
 import imitation.data.rollout as rollout
 import tempfile
 import pathlib
 from imitation.util import logger
-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
@@ -56,7 +54,7 @@ def train(
         options=ALL_OPTIONS,
         discriminator=imitation_discriminator(discriminator),
         discount=discount,
-        ll_buffer_capacity=expert_batch_size,
+        ll_buffer_capacity=generator_total_steps*10,
     )
     generator = stable_baselines3.PPO(
         OptionsCnnPolicy,
@@ -100,7 +98,7 @@ def evaluate():
 
 # %%
 if __name__ == '__main__':
-    with open("data/NormalizedIntersimpleExpertMu.001N10000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
+    with open("data/NormalizedIntersimpleExpertMu.001N100000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
         trajectories = pickle.load(f)
     transitions = rollout.flatten_trajectories(trajectories)
     train(transitions)