Add debug output and original implementation

.gitignore

@@ -1,7 +1,3 @@
-*.pkl
-*.pt
-*.zip
-
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

config/value_dice.json5

@@ -9,12 +9,12 @@
         deepsets: {
             input_dim: 6, // number of relative state vars for others
             phi: {
-                hidden_n: 2,
+                hidden_n: 1,
                 hidden_dim: 20,
             },
             latent_dim: 20,
             rho: {
-                hidden_n: 2,
+                hidden_n: 1,
                 hidden_dim: 10,
             },
             output_dim: 10
@@ -43,12 +43,12 @@
         deepsets: {
             input_dim: 6, // number of relative state vars for others
             phi: {
-                hidden_n: 2,
+                hidden_n: 1,
                 hidden_dim: 20,
             },
             latent_dim: 20,
             rho: {
-                hidden_n: 2,
+                hidden_n: 1,
                 hidden_dim: 10,
             },
             output_dim: 10
@@ -70,13 +70,13 @@
     },
     policy_optim: {
             optimizer: 'adam',
-            lr: 1e-3,
+            lr: 1e-0,
-            weight_decay: 0.1,
+            weight_decay: 0.01,
     },
     value_optim: {
             optimizer: 'adam',
-            lr: 1e-3,
+            lr: 1e-6,
-            weight_decay: 0.1,
+            weight_decay: 0.01,
     },
     train_epochs: 200,
     train_batch_size: 32,

scratch/etienne/intersimple/data/generate.sh

@@ -1,4 +1,3 @@
-#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl'
+#python -m intersimple.expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl'
-#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.005}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.005.pkl'
+#python -m intersimple.expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.005}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.005.pkl'
-#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl'
+python -m intersimple.expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl' --video
-python -m expert --env=NRasterized --min_timesteps=200 --env_args='{agent:51,width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl'

scratch/etienne/intersimple/gail/discriminator.py

@@ -1,8 +1,5 @@
 import torch
 
-# imitation.rewards.discrim_nets.DiscrimNetGAIL is composed of self.discriminator (nn.Module),
-# which gets called with inputs (state, action) when needed.
-
 class CnnDiscriminator(torch.nn.Module):
     """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy."""
 
@@ -26,16 +23,11 @@ class CnnDiscriminator(torch.nn.Module):
             torch.nn.LazyLinear(1), # 512 -> 1
         )
     
-    @staticmethod
+    def forward(self, state, action):
-    def _concatenate(state, action):
         b, _, h, w = state.shape
         _, a = action.shape
-        act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
+        act_layer = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
-        sa = torch.cat((state, act), -3)
+        sa = torch.cat((act_layer, state), -3)
-        return sa
-
-    def forward(self, state, action):
-        sa = self._concatenate(state, action)
         return self.cnn(sa).squeeze()
 
 class MlpDiscriminator(torch.nn.Module):

scratch/etienne/intersimple/gail/test_discriminator.py

@@ -1,45 +0,0 @@
-from intersim.envs.intersimple import NRasterized
-from discriminator import CnnDiscriminator
-import torch
-
-def test_image_concatenation():
-    env = NRasterized()
-    disc = CnnDiscriminator(env)
-    s = torch.tensor(env.reset()).unsqueeze(0)
-    a = torch.tensor([[0.5]])
-    sa = disc._concatenate(s, a)
-
-    assert s.shape == (1, 5, 200, 200)
-    assert a.shape == (1, 1)
-    assert sa.shape == (1, 6, 200, 200)
-    assert torch.allclose(sa[:, :5], 1.0 * s)
-    assert (sa[:, 5] == a.unsqueeze(-1)).all()
-
-def test_image_concatenation3():
-    env = NRasterized()
-    disc = CnnDiscriminator(env)
-
-    s1 = env.reset()
-    a1 = 0.15
-    s2, _, _, _ = env.step(0.9)
-    a2 = 0.25
-    s3, _, _, _ = env.step(-0.9)
-    a3 = 0.35
-    
-    s = torch.stack([
-        torch.tensor(s1),
-        torch.tensor(s2),
-        torch.tensor(s3)
-    ], axis=0)
-    a = torch.tensor([
-        [a1],
-        [a2],
-        [a3],
-    ])
-    sa = disc._concatenate(s, a)
-
-    assert s.shape == (3, 5, 200, 200)
-    assert a.shape == (3, 1)
-    assert sa.shape == (3, 6, 200, 200)
-    assert torch.allclose(sa[:, :5], 1.0 * s)
-    assert (sa[:, 5] == a.unsqueeze(-1)).all()

scratch/etienne/intersimple/gail_flat.py

@@ -10,22 +10,15 @@ from imitation.algorithms import adversarial, bc
 from imitation.data import rollout
 from imitation.util import logger
 
-from intersimple.intersimple import IntersimpleReward, speed_reward
+from intersim.envs.intersimple import IntersimpleReward
 
 from gail.discriminator import MlpDiscriminator
-import numpy as np
-import functools
-from stable_baselines3.common.evaluation import evaluate_policy
-from ray import tune
-import os
-import torch
 
 model_name = 'gail_flat'
 
 # %%
 # Load pickled test demonstrations.
-#with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
+with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
-with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl", "rb") as f:
     # This is a list of `imitation.data.types.Trajectory`, where
     # every instance contains observations and actions for a single expert
     # demonstration.
@@ -43,50 +36,20 @@ tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
 tempdir_path = pathlib.Path(tempdir.name)
 print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
 
-def training_function(config, checkpoint_dir=None):
+# Train GAIL on expert data.
+# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
+# iterates over dictionaries containing observations, actions, and next_observations.
 logger.configure(tempdir_path / "GAIL/")
-
-    discriminator = MlpDiscriminator()
-    if checkpoint_dir:
-        discriminator.load_state_dict(torch.load(os.path.join(checkpoint_dir, 'disc_checkpoint')))
-        generator = sb3.PPO.load(os.path.join(checkpoint_dir, 'gen_checkpoint'))
-    else:
-        generator = sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=config['n_steps'])
-
 gail_trainer = adversarial.GAIL(
     venv,
     expert_data=transitions,
-        expert_batch_size=config['expert_batch_size'],
+    expert_batch_size=220,
-        n_disc_updates_per_round=config['n_disc_updates_per_round'],
+    #n_disc_updates_per_round=32,
     discrim_kwargs={'discrim_net': MlpDiscriminator()},
-        gen_algo=generator,
+    gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=4096),
 )
-
+gail_trainer.train(total_timesteps=80000)
-    def callback(epoch):
+gail_trainer.gen_algo.save(model_name)
-        eval_env = IntersimpleReward(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.))
-        #sync_envs_normalization(self.training_env, self.eval_env)
-        episode_rewards, episode_lengths = evaluate_policy(generator, eval_env)
-        tune.report(progress=np.mean(episode_rewards))
-
-        with tune.checkpoint_dir(step=epoch) as checkpoint_dir:
-            gail_trainer.gen_algo.save(os.path.join(checkpoint_dir, 'gen_checkpoint'))
-            torch.save(discriminator.state_dict(), os.path.join(checkpoint_dir, 'disc_checkpoint'))
-
-    gail_trainer.train(total_timesteps=400000, callback=callback)
-
-analysis = tune.run(
-    training_function,
-    config = {
-        'expert_batch_size': tune.randint(1, 220), #220,
-        'n_disc_updates_per_round': tune.randint(2, 100), #16,
-        'n_steps': tune.randint(1, 10000), #4096,
-    },
-    resources_per_trial={
-        'gpu': 1,
-    }
-)
-
-print('Best config', analysis.get_best_config(metric='progress', mode='max'))
 
 #del gail_trainer
 

scratch/etienne/intersimple/gail_image.py

@@ -18,7 +18,7 @@ model_name = 'gail_image'
 
 # %%
 # Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+with open("data/NormalizedIntersimpleExpert_NRasterizedAgent51.pkl", "rb") as f:
     # This is a list of `imitation.data.types.Trajectory`, where
     # every instance contains observations and actions for a single expert
     # demonstration.
@@ -30,7 +30,7 @@ with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.
 # (observation, actions, next_observation) transitions.
 transitions = rollout.flatten_trajectories(trajectories)
 
-venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2})
+venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'agent': 51})
 
 tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
 tempdir_path = pathlib.Path(tempdir.name)
@@ -43,11 +43,10 @@ logger.configure(tempdir_path / "GAIL/")
 gail_trainer = adversarial.GAIL(
     venv,
     expert_data=transitions,
-    expert_batch_size=32,
+    expert_batch_size=200,
-    #n_disc_updates_per_round=2048,
+    n_disc_updates_per_round=2048,
     discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-    gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
+    gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=128),
-    allow_variable_horizon=True,
 )
 gail_trainer.train(total_timesteps=100000)
 gail_trainer.gen_algo.save(model_name)
@@ -57,7 +56,7 @@ gail_trainer.gen_algo.save(model_name)
 # %%
 model = sb3.PPO.load(model_name)
 
-env = NRasterized(agent=51, width=36, height=36, m_per_px=2)
+env = NRasterized(agent=51)
 
 obs = env.reset()
 while True:

scratch/etienne/intersimple/gail_options.py

@@ -0,0 +1,116 @@
+from gail.discriminator import MlpDiscriminator
+from imitation.algorithms import adversarial
+import stable_baselines3
+import torch.utils.data
+import numpy as np
+from intersim.envs.intersimple import Intersimple
+import itertools
+from torch.distributions import Categorical
+import gym
+
+class OptionsMlpPolicy:
+
+    def __init__(self, *args, **kwargs):
+        self._policy = stable_baselines3.common.policies.ActorCriticPolicy(
+            *args, **kwargs
+        )
+
+    def _prior_distribution(self, s):
+        latent_pi, _, latent_sde = self._policy._get_latent(s)
+        distribution = self._policy._get_action_dist_from_latent(latent_pi, latent_sde)
+        return distribution.distribution
+
+    def predict(self, obs):
+        s, m = obs
+        prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        ch = posterior.sample()
+        return ch
+
+    def evaluate_actions(self, obs, ch):
+        s, m = obs
+        values = self._policy.value_net(s)
+        prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        return values, posterior.logprob(ch), posterior.entropy() # additional values used by PPO.train
+
+def available_actions(env):
+    """Return mask of available actions given current `env` state."""
+    return np.ones((env.num_hl_actions,))
+
+def generate_plan(env, i):
+    """Generate input profile for high-level action `i`."""
+    return np.zeros((env.num_hl_steps,))
+
+def feasible(env, plan):
+    """Check if input profile is feasible given current `env` state."""
+    return True
+
+def sample_ll(env, generator):
+    """Sample low-level (state, action) pairs for discriminator training."""
+    done = True
+    while True:
+        if done:
+            s = env.reset()
+
+        m = available_actions(env)
+        ch = generator.policy.predict((s, m))
+        plan = list(generate_plan(env, ch))
+
+        while not done and plan and feasible(env, plan):
+            a = plan.pop()
+            yield (s, a)
+            s, _, done, _ = env.step(a)
+
+def train_discriminator(env, expert_data, generator, discriminator, generator_batch_size):
+    expert_samples = next(expert_data)
+    generator_samples = itertools.islice(sample_ll(env, generator), generator_batch_size)
+    discriminator.train_disc(expert_samples, generator_samples)
+
+def sample_hl(env, generator, discriminator):
+    """Sample high-level (state, action, reward) tuples for generator training."""
+    done = True
+    while True:
+        if done:
+            s = env.reset()
+
+        m = available_actions(env)
+        obs = (s, m)
+        ch = generator.policy.predict((s, m))
+        plan = list(generate_plan(env, ch))
+        r = 0
+        discount = 1
+
+        while not done and plan and feasible(env, plan):
+            a = plan.pop()
+            r += discount * discriminator.discrim_net(s, a)
+            discount *= env.discount
+            s, _, done, _ = env.step(a)
+        
+        yield (obs, ch, r)
+
+def train_generator(env, generator, discriminator, generator_batch_size):
+    generator_samples = itertools.islice(sample_hl(env, generator, discriminator), generator_batch_size)
+    generator.rollout_buffer.reset()
+    generator.rollout_buffer.add(generator_samples)
+    generator.train()
+
+class OptionsEnv(gym.Wrapper):
+    
+    def __init__(self, env):
+        super().__init__(env)
+        self.action_space = gym.spaces.Discrete(env.num_hl_options)
+
+def train(expert_data, epochs=10, generator_batch_size=1024, expert_batch_size=1024, num_hl_options=10, num_hl_steps=10, discount=0.99):
+    env = Intersimple()
+    env.num_hl_options = num_hl_options
+    env.num_hl_steps = num_hl_steps
+    env.discount = discount
+
+    discriminator = adversarial.GAIL(discrim_kwargs={'discrim_net': MlpDiscriminator()})
+    generator = stable_baselines3.PPO(OptionsMlpPolicy, OptionsEnv(env))
+    expert_data = torch.utils.data.DataLoader(expert_data, expert_batch_size)
+
+    for _ in range(epochs):
+        train_discriminator(env, expert_data, generator, discriminator, generator_batch_size)
+        train_generator(env, generator, discriminator, generator_batch_size)

scratch/etienne/intersimple/gail_options_image.py

@@ -1,297 +0,0 @@
-# %%
-from gail.discriminator import CnnDiscriminator
-from imitation.algorithms import adversarial
-import stable_baselines3
-import torch.utils.data
-import numpy as np
-from intersim.envs.intersimple import NRasterized
-import itertools
-from torch.distributions import Categorical
-import gym
-import torch
-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
-
-model_name = 'gail_options_image'
-env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
-
-ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
-
-class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
-
-    def __init__(self, observation_space, *args, **kwargs):
-        super().__init__(observation_space['obs'], *args, **kwargs)
-
-    def _prior_distribution(self, s):
-        latent_pi, latent_vf, latent_sde = self._get_latent(s)
-        distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
-        values = self.value_net(latent_vf)
-        return values, distribution.distribution
-
-    def predict(self, obs):
-        s, m = obs['obs'], obs['mask']
-        values, prior = self._prior_distribution(s)
-        posterior = Categorical(prior.probs * m)
-        ch = posterior.sample()
-        return ch, values, posterior.log_prob(ch)
-
-    def evaluate_actions(self, obs, ch):
-        s, m = obs['obs'], obs['mask']
-        values, prior = self._prior_distribution(s)
-        posterior = Categorical(prior.probs * m)
-        return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train
-
-def available_actions(env):
-    """Return mask of available actions given current `env` state."""
-    valid = np.array([feasible(env, generate_plan(env, i), i) for i in range(len(ALL_OPTIONS))])
-    return valid
-
-def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
-    """Smoothly target a velocity in a given number of steps"""
-    # for now, constant acceleration
-    a = (target_v - current_v)  / (t * dt)
-    return a*np.ones((t,))
-
-def generate_plan(env, i):
-    """Generate input profile for high-level action `i`."""
-    assert i < len(ALL_OPTIONS), "Invalid option index {i}"
-    target_v, t = ALL_OPTIONS[i]
-    current_v = env._env.state[env._agent, 1].item() # extract from env
-    plan = target_velocity_plan(current_v, target_v, t, env._env._dt)
-    assert len(plan) == t, "incorrect plan length"
-    return plan
-
-def check_future_collisions_fast(env, actions):
-    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
-    
-    Vehicles are (over-)approximated by single circles.
-
-    Args:
-        env (gym.Env): current environment state
-        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
-    Returns:
-        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
-    """
-    B, (T, nv, _) = len(actions), actions[0].shape
-
-    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
-    assert states.shape == (B, T, nv, 5)
-
-    distance = ((states[:, :, :, :2] - states[:, :, env._agent:env._agent+1, :2])**2).sum(-1).sqrt()
-    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
-    distance[:, :, env._agent] = np.inf # cannot collide with itself
-    assert distance.shape == (B, T, nv)
-
-    radius = (env._env._lengths**2 + env._env._widths**2).sqrt() / 2
-    min_distance = radius[env._agent] + radius
-    min_distance = min_distance.unsqueeze(0).unsqueeze(0)
-    assert min_distance.shape == (1, 1, nv)
-
-    return (distance > min_distance).all(-1).all(-1)
-
-def feasible(env, plan, ch):
-    """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
-    
-    # zero pad plan - Take (T,) np plan and convert it to (T, nv, 1) torch.Tensor
-    full_plan = torch.zeros(len(plan), env._env._nv, 1)
-    full_plan[:, env._agent, 0] = torch.tensor(plan)
-    valid = check_future_collisions_fast(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
-    return ch == 0 or valid.item()
-
-def sample(env, generator, discriminator, level: str):
-    """
-        Sample low-level (state, action, next_state) tuples for discriminator training or 
-        high-level (state, action, reward) tuples for generator training.
-    """
-    done = True
-    while True:
-        episode_start = False
-        if done:
-            s = env.reset()
-            m = available_actions(env)
-            done = False
-            episode_start = True
-
-        obs = {'obs': s, 'mask': m}
-        ch, value, log_prob = generator.policy.predict({
-            'obs': torch.tensor(s).unsqueeze(0).to(generator.policy.device),
-            'mask': torch.tensor(m).unsqueeze(0).to(generator.policy.device),
-        })
-        plan = list(map(float, generate_plan(env, ch)))
-
-        assert not done
-        assert plan
-        assert feasible(env, plan, ch), f'Infeasible hl action {ch}'
-
-        r = 0
-        discount = 1
-        while not done and plan and feasible(env, plan, ch):
-            a, plan = env._normalize(plan[0]), plan[1:]
-            if level == 'high':
-                r += discount * discriminator.discrim_net.discriminator(
-                    torch.tensor(s).unsqueeze(0).to(discriminator.discrim_net.device()),
-                    torch.tensor([[a]]).to(discriminator.discrim_net.device()),
-                )
-                discount *= env.discount
-
-            nexts, _, done, _ = env.step(a)
-            m = available_actions(env)
-
-            if level == 'low':
-                yield {
-                    'obs': s,
-                    'next_obs': nexts,
-                    'acts': np.array((a,)),
-                    'dones': np.array(done),
-                }
-            s = nexts
-
-        if level == 'high':
-            yield {
-                'obs': obs,
-                'option': ch,
-                'reward': r.detach(),
-                'episode_start': episode_start,
-                'value': value.detach(),
-                'log_prob': log_prob.detach(),
-                'done': done,
-            }
-
-def flatten_transitions(transitions):
-    return {
-        'obs': np.stack(list(t['obs'] for t in transitions), axis=0),
-        'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0),
-        'acts': np.stack(list(t['acts'] for t in transitions), axis=0),
-        'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
-    }
-
-def train_discriminator(env, generator, discriminator, num_samples):
-    transitions = list(itertools.islice(sample(env, generator, None, 'low'), num_samples))
-    generator_samples = flatten_transitions(transitions)
-    discriminator.train_disc(gen_samples=generator_samples)
-
-def train_generator(env, generator, discriminator, num_samples):
-    generator_samples = list(itertools.islice(sample(env, generator, discriminator, 'high'), num_samples+1))
-    
-    generator.rollout_buffer.reset()
-    for s in generator_samples[:-1]:
-        generator.rollout_buffer.add(
-            obs=s['obs'],
-            action=s['option'].cpu(),
-            reward=s['reward'].cpu(),
-            episode_start=s['episode_start'],
-            value=s['value'],
-            log_prob=s['log_prob'],
-        )
-    
-    generator.rollout_buffer.compute_returns_and_advantage(
-        last_values=generator_samples[-1]['value'],
-        dones=generator_samples[-1]['done'],
-    )
-    
-    generator.train()
-
-class OptionsEnv(gym.Wrapper):
-    
-    def __init__(self, env):
-        super().__init__(env)
-        num_hl_options = len(ALL_OPTIONS)
-        self.action_space = gym.spaces.Discrete(num_hl_options)
-        self.observation_space = gym.spaces.Dict({
-            'obs': env.observation_space,
-            'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
-        })
-
-def train(expert_data, epochs=10, expert_batch_size=32, generator_steps=2048, discount=0.99):
-    env = NRasterized(**env_settings)
-    env.discount = discount
-
-    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(NRasterized, n_envs=1, env_kwargs=env_settings)
-    discriminator = adversarial.GAIL(
-        expert_data=expert_data,
-        expert_batch_size=expert_batch_size,
-        discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-        venv=venv, # unused
-        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
-    )
-
-    generator = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env),
-        verbose=1,
-        n_steps=generator_steps,
-    )
-
-    # 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,
-    )
-
-    for _ in range(epochs):
-        train_discriminator(env, generator, discriminator, num_samples=expert_batch_size)
-        train_generator(env, generator, discriminator, num_samples=generator_steps)
-    
-    return generator
-
-# %%
-if __name__ == '__main__':
-    # %%
-    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
-        trajectories = pickle.load(f)
-    transitions = rollout.flatten_trajectories(trajectories)
-    generator = train(transitions, generator_steps=200)
-
-    generator.save(model_name)
-
-    # %%
-    model = stable_baselines3.PPO.load(model_name)
-
-    env = NRasterized(**env_settings)
-
-    for transition in sample(env, generator, None, 'low'):
-        env.render()
-        if transition['dones']:
-            break
-
-    env.close(filestr='render/'+model_name)
-
-# %% Tests
-
-def test_ll_transitions_vs_expert_data():
-    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
-        expert_trajectories = pickle.load(f)
-    expert_transitions = rollout.flatten_trajectories(expert_trajectories)
-
-    env = NRasterized(agent=51, width=36, height=36, m_per_px=2)
-
-    gen_transitions = list(itertools.islice(sample(
-        env=NRasterized(**env_settings),
-        generator=stable_baselines3.PPO(
-            OptionsCnnPolicy,
-            OptionsEnv(env),
-            verbose=1,
-        ),
-        discriminator=None,
-        level='low'
-    ), 10))
-    gen_transitions = flatten_transitions(gen_transitions)
-
-    assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape
-    assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape
-    assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape
-    assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape
-
-
-def test_hl_transitions():
-    pass

scratch/etienne/intersimple/ppo_speed_image_lowres.py

@@ -1,49 +0,0 @@
-# %%
-from stable_baselines3 import PPO
-from intersim.envs.intersimple import NRasterized, speed_reward
-import functools
-
-model_name = "ppo_speed_image_lowres"
-
-#def reward(state, action, info):
-#    speed = state[2].item()
-#    r = speed if speed < 10 else (10 - 5 * (speed - 10))
-#    return 0.1 * r
-
-env = NRasterized(
-    agent=51,
-    height=36,
-    width=36,
-    m_per_px=2,
-    reward=functools.partial(
-        speed_reward,
-        collision_penalty=0
-    ),
-)
-
-# %%
-model = PPO(
-    "CnnPolicy", env,
-    verbose=1,
-)
-model.learn(total_timesteps=100000)
-model.save(model_name)
-
-print('Done training.')
-
-del model # remove to demonstrate saving and loading
-
-# %%
-model = PPO.load(model_name)
-
-obs = env.reset()
-while True:
-    action, _states = model.predict(obs)
-    obs, rewards, done, info = env.step(action)
-    env.render(mode='post')
-    if done:
-        break
-
-env.close(filestr='render/'+model_name)
-
-# %%

scratch/etienne/intersimple/ppo_speed_image_lowres_random.py

@@ -1,42 +0,0 @@
-# %%
-from stable_baselines3 import PPO
-from intersim.envs.intersimple import NRasterizedRandomAgent, speed_reward
-import functools
-
-model_name = "ppo_speed_image_lowres_random"
-
-env = NRasterizedRandomAgent(
-    height=36,
-    width=36,
-    m_per_px=2,
-    reward=functools.partial(
-        speed_reward,
-        collision_penalty=0
-    )
-)
-
-# %%
-model = PPO(
-    "CnnPolicy", env,
-    verbose=1,
-    batch_size=2048,
-)
-model.learn(total_timesteps=2e5)
-model.save(model_name)
-
-print('Done training.')
-
-del model # remove to demonstrate saving and loading
-
-# %%
-model = PPO.load(model_name)
-
-obs = env.reset()
-while True:
-    action, _states = model.predict(obs)
-    obs, rewards, done, info = env.step(action)
-    env.render(mode='post')
-    if done:
-        break
-
-env.close(filestr='render/'+model_name)

scratch/johannes/valuedice.py

@@ -0,0 +1,121 @@
+def weighted_softmax(x, weights, axis=0):
+  x = x - tf.reduce_max(x, axis=axis)
+  return weights * tf.exp(x) / tf.reduce_sum(
+      weights * tf.exp(x), axis=axis, keepdims=True)
+
+
+@tf.function
+  def update(self,
+             expert_dataset_iter,
+             policy_dataset_iter,
+             discount,
+             replay_regularization=0.05,
+             nu_reg=10.0):
+    """A function that updates nu network.
+    When replay regularization is non-zero, it learns
+    (d_pi * (1 - replay_regularization) + d_rb * replay_regulazation) /
+    (d_expert * (1 - replay_regularization) + d_rb * replay_regulazation)
+    instead.
+    Args:
+      expert_dataset_iter: An tensorflow graph iteratable over expert data.
+      policy_dataset_iter: An tensorflow graph iteratable over training policy
+        data, used for regularization.
+      discount: An MDP discount.
+      replay_regularization: A fraction of samples to add from a replay buffer.
+      nu_reg: A grad penalty regularization coefficient.
+    """
+
+    (expert_states, expert_actions,
+     expert_next_states) = expert_dataset_iter.get_next()
+
+    expert_initial_states = expert_states
+
+    rb_states, rb_actions, rb_next_states, _, _ = policy_dataset_iter.get_next(
+    )[0]
+
+    with tf.GradientTape(
+        watch_accessed_variables=False, persistent=True) as tape:
+      tape.watch(self.actor.variables)
+      tape.watch(self.nu_net.variables)
+
+      _, policy_next_actions, _ = self.actor(expert_next_states)
+    #   _, rb_next_actions, rb_log_prob = self.actor(rb_next_states)
+
+      _, policy_initial_actions, _ = self.actor(expert_initial_states)
+
+      Inputs for the linear part of DualDICE loss.
+      expert_init_inputs = tf.concat(
+          [expert_initial_states, policy_initial_actions], 1)
+
+      expert_inputs = tf.concat([expert_states, expert_actions], 1)
+      expert_next_inputs = tf.concat([expert_next_states, policy_next_actions],
+                                     1)
+
+      rb_inputs = tf.concat([rb_states, rb_actions], 1)
+      rb_next_inputs = tf.concat([rb_next_states, rb_next_actions], 1)
+
+      expert_nu_0 = self.nu_net(expert_init_inputs)
+      expert_nu = self.nu_net(expert_inputs)
+      expert_nu_next = self.nu_net(expert_next_inputs)
+
+      rb_nu = self.nu_net(rb_inputs)
+      rb_nu_next = self.nu_net(rb_next_inputs)
+
+      expert_diff = expert_nu - discount * expert_nu_next
+      rb_diff = rb_nu - discount * rb_nu_next
+
+      linear_loss_expert = tf.reduce_mean(expert_nu_0 * (1 - discount))
+
+      linear_loss_rb = tf.reduce_mean(rb_diff)
+
+      rb_expert_diff = tf.concat([expert_diff, rb_diff], 0)
+      rb_expert_weights = tf.concat([
+          tf.ones(expert_diff.shape) * (1 - replay_regularization),
+          tf.ones(rb_diff.shape) * replay_regularization
+      ], 0)
+
+      rb_expert_weights /= tf.reduce_sum(rb_expert_weights)
+      non_linear_loss = tf.reduce_sum(
+          tf.stop_gradient(
+              weighted_softmax(rb_expert_diff, rb_expert_weights, axis=0)) *
+          rb_expert_diff)
+
+      linear_loss = (
+          linear_loss_expert * (1 - replay_regularization) +
+          linear_loss_rb * replay_regularization)
+
+      loss = (non_linear_loss - linear_loss)
+
+      alpha = tf.random.uniform(shape=(expert_inputs.shape[0], 1))
+
+      nu_inter = alpha * expert_inputs + (1 - alpha) * rb_inputs
+      nu_next_inter = alpha * expert_next_inputs + (1 - alpha) * rb_next_inputs
+
+      nu_inter = tf.concat([nu_inter, nu_next_inter], 0)
+
+      with tf.GradientTape(watch_accessed_variables=False) as tape2:
+        tape2.watch(nu_inter)
+        nu_output = self.nu_net(nu_inter)
+      nu_grad = tape2.gradient(nu_output, [nu_inter])[0] + EPS
+      nu_grad_penalty = tf.reduce_mean(
+          tf.square(tf.norm(nu_grad, axis=-1, keepdims=True) - 1))
+
+      nu_loss = loss + nu_grad_penalty * nu_reg
+      pi_loss = -loss + keras_utils.orthogonal_regularization(self.actor.trunk)
+
+    nu_grads = tape.gradient(nu_loss, self.nu_net.variables)
+    pi_grads = tape.gradient(pi_loss, self.actor.variables)
+
+    self.nu_optimizer.apply_gradients(zip(nu_grads, self.nu_net.variables))
+    self.actor_optimizer.apply_gradients(zip(pi_grads, self.actor.variables))
+
+    del tape
+
+    self.avg_nu_expert(expert_nu)
+    self.avg_nu_rb(rb_nu)
+
+    self.nu_reg_metric(nu_grad_penalty)
+    self.avg_loss(loss)
+
+    self.avg_actor_loss(pi_loss)
+    self.avg_actor_entropy(-rb_log_prob)

src/value_dice/value_dice.py

@@ -215,9 +215,16 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
 
         # nonlinear loss
         value_diff = value - discount * value_next
-        nonlinear_loss = torch.logsumexp(value_diff, dim=0) - np.log(len(value_diff))
+        # print(value_diff)
+        # nonlinear_loss = torch.logsumexp(value_diff, dim=0) #- np.log(len(value_diff))
+        nonlinear_loss = torch.log(torch.mean(torch.exp(value_diff), dim=0))
 
         loss = nonlinear_loss - linear_loss
+        print("Loss report:")
+        print("Linear: {}".format(linear_loss.item()))
+        print("Nonlinear: {}".format(nonlinear_loss.item()))
+        print("Total: {}".format(loss.item()))
+
         return loss
 
 
@@ -237,7 +244,6 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
         # train
         epoch_loss = 0
         for (batch_idx, batch) in enumerate(training_loader):
-
             loss = f_value_dice_loss(batch)
 
             # In original implementation policy is regularized with orthogonal regularization, 
@@ -260,12 +266,22 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
             if batch_idx % 2 == 0:
                 policy_optimizer.zero_grad()
                 policy_loss.backward()
-                clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
+                # clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
                 policy_optimizer.step()
+
+                grad_list = torch.cat([torch.flatten(p.grad) for p in policy.policy.parameters()])
+                torch.mean(grad_list)
+                print("gradient stats:")
+                print(torch.mean(grad_list))
+                print(torch.std(grad_list))
+                print(torch.min(grad_list))
+                print(torch.max(grad_list))
+                # print(policy.policy.head.layers[0].weight.grad)
+                # print(policy.policy.head.layers[0].bias.grad)
             else:
                 value_optimizer.zero_grad()
                 value_loss.backward()
-                clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
+                # clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
                 value_optimizer.step()
 
             epoch_loss += loss.item() / len(train_dataset)