From 284d4af63789ff6a1d8b3e50bd9e511834312c3c Mon Sep 17 00:00:00 2001
From: ebuehrle <43623224+ebuehrle@users.noreply.github.com>
Date: Wed, 10 Nov 2021 11:54:41 +0100
Subject: [PATCH] Train discriminator

---
 scratch/etienne/intersimple/train_discrim.py | 71 ++++++++++++++++++++
 1 file changed, 71 insertions(+)
 create mode 100644 scratch/etienne/intersimple/train_discrim.py

diff --git a/scratch/etienne/intersimple/train_discrim.py b/scratch/etienne/intersimple/train_discrim.py
new file mode 100644
index 0000000..71a3140
--- /dev/null
+++ b/scratch/etienne/intersimple/train_discrim.py
@@ -0,0 +1,71 @@
+# %%
+import sys
+sys.path.append('../../../')
+
+import pickle
+import imitation.data.rollout as rollout
+import imitation.data.types as types
+import torch
+from gail.envs import TLNRasterizedRouteRandomAgentLocation
+import tempfile
+import pathlib
+from imitation.util import logger
+from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
+from imitation.algorithms import adversarial
+from src.discriminator import CnnDiscriminator
+import stable_baselines3
+from tqdm import tqdm
+
+with open("data/NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
+    trajectories = pickle.load(f)
+transitions = rollout.flatten_trajectories(trajectories)
+
+# %%
+env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 200}
+env = TLNRasterizedRouteRandomAgentLocation(**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}/.")
+
+expert_batch_size = 4096
+
+venv = DummyVecEnv([lambda: env])
+discriminator = adversarial.GAIL(
+    expert_data=transitions,
+    expert_batch_size=expert_batch_size,
+    #discrim_kwargs={'discrim_net': NoisyDiscriminator(venv, std=0.25)},
+    disc_opt_cls=torch.optim.RMSprop,
+    disc_opt_kwargs={'lr': 0.0001, 'weight_decay': 0.003},
+    discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+    venv=venv, # unused
+    gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+)
+
+expert_data_loader = torch.utils.data.DataLoader(
+    transitions,
+    batch_size=expert_batch_size,
+    collate_fn=types.transitions_collate_fn,
+    shuffle=True,
+    drop_last=True,
+)
+
+gen_data_loader = torch.utils.data.DataLoader(
+    transitions,
+    batch_size=expert_batch_size,
+    collate_fn=types.transitions_collate_fn,
+    shuffle=True,
+    drop_last=True,
+)
+
+# %%
+epochs = 1000
+for i in tqdm(range(epochs)):
+    for expert_samples, gen_samples in zip(expert_data_loader, gen_data_loader):
+        # randomly corrupt actions
+        gen_samples['acts'] = -1 + 2 * torch.rand(*gen_samples['acts'].shape)
+
+        discriminator.train_disc(expert_samples=expert_samples, gen_samples=gen_samples)
+
+    torch.save(discriminator.discrim_net.state_dict(), 'train_discrim.pt')