Merge branch 'options-env' into dev

scratch/etienne/intersimple/gail_options_image_random_location.py

@@ -5,13 +5,7 @@ sys.path.append('../../../')
 from src.discriminator import CnnDiscriminatorFlatAction
 from imitation.algorithms import adversarial
 import stable_baselines3
-import torch.utils.data
-import numpy as np
 from intersim.envs import NRasterizedRouteSpeedRandomAgentLocation
-import itertools
-from torch.distributions import Categorical
-import gym
-import torch
 import pickle
 import imitation.data.rollout as rollout
 import tempfile
@@ -20,8 +14,9 @@ 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.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
-from src.gail.train import train_discriminator, train_generator
+from src.gail.train import flatten_transitions
+
+from gail.options2 import OptionsEnv, RenderOptions, imitation_discriminator
 
 model_name = 'gail_options_image_random_location'
 env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'map_color': 128, 'mu': 0.001}
@@ -30,12 +25,13 @@ ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is sa
 
 def train(
         expert_data,
-        epochs=200,
         expert_batch_size=1024,
-        generator_steps=1024,
+        discriminator_updates_per_round=10,
+        generator_steps=256,
+        generator_total_steps=1024,
+        generator_updates_per_round=10,
         discount=0.99,
-        n_disc_updates_per_round=10,
-        n_gen_updates_per_round=10,
+        epochs=200,
     ):
     env = NRasterizedRouteSpeedRandomAgentLocation(**env_settings)
     env.discount = discount
@@ -55,24 +51,25 @@ 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)
     generator = stable_baselines3.PPO(
         OptionsCnnPolicy,
-        OptionsEnv(env, options=ALL_OPTIONS),
+        options_env,
         verbose=1,
         n_steps=generator_steps,
-        n_epochs=n_gen_updates_per_round,
-    )
-
-    # PPO.train requires logger as set up in
-    # PPO._setup_learn (called by PPO.learn)
-    generator._logger = stable_baselines3.common.utils.configure_logger(
-        generator.verbose,
-        generator.tensorboard_log,
+        n_epochs=generator_updates_per_round,
     )
 
     for _ in tqdm(range(epochs)):
-        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size, n_updates=n_disc_updates_per_round)
-        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
+        # train generator
+        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):
+            discriminator.train_disc(gen_samples=generator_samples)
+
         generator.save(model_name)
     
     return generator