From 45a99978e421d3828bf96c91286e480e082b70a5 Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 21 Oct 2021 05:33:05 -0700 Subject: [PATCH] adding discriminators to main folder, utilities to render a video from a saved model --- src/discriminator/__init__.py | 1 + src/discriminator/discriminator.py | 101 +++++++++++++++++++++++++++++ src/util/__init__.py | 1 + src/util/render_env.py | 58 +++++++++++++++++ tests/test_discriminator.py | 45 +++++++++++++ 5 files changed, 206 insertions(+) create mode 100644 src/discriminator/__init__.py create mode 100644 src/discriminator/discriminator.py create mode 100644 src/util/render_env.py create mode 100644 tests/test_discriminator.py diff --git a/src/discriminator/__init__.py b/src/discriminator/__init__.py new file mode 100644 index 0000000..4ad66b2 --- /dev/null +++ b/src/discriminator/__init__.py @@ -0,0 +1 @@ +from src.discriminator.discriminator import * \ No newline at end of file diff --git a/src/discriminator/discriminator.py b/src/discriminator/discriminator.py new file mode 100644 index 0000000..1c9664c --- /dev/null +++ b/src/discriminator/discriminator.py @@ -0,0 +1,101 @@ +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.""" + + def __init__(self, env): + super().__init__() + + obs_channels, _, _ = env.observation_space.shape + (action_size,) = env.action_space.shape + in_channels = obs_channels + action_size + + self.cnn = torch.nn.Sequential( + torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # 5+1 -> 32 + torch.nn.ReLU(), + torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64 + torch.nn.ReLU(), + torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64 + torch.nn.ReLU(), + torch.nn.Flatten(start_dim=1, end_dim=-1), + torch.nn.LazyLinear(512), # 28224 -> 512 + torch.nn.ReLU(), + torch.nn.LazyLinear(1), # 512 -> 1 + ) + + @staticmethod + def _concatenate(state, action): + b, _, h, w = state.shape + _, a = action.shape + act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w)) + sa = torch.cat((state, act), -3) + return sa + + def forward(self, state, action): + sa = self._concatenate(state, action) + assert sa.ndim == 4 + return self.cnn(sa).squeeze(1) + +class CnnDiscriminatorFlatAction(torch.nn.Module): + """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy.""" + + def __init__(self, env): + super().__init__() + + obs_channels, _, _ = env.observation_space.shape + (action_size,) = env.action_space.shape + in_channels = obs_channels + + self.cnn = torch.nn.Sequential( + torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # in_channels -> 32 + torch.nn.ReLU(), + torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64 + torch.nn.ReLU(), + torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64 + torch.nn.ReLU(), + torch.nn.Flatten(start_dim=1, end_dim=-1), + torch.nn.LazyLinear(128), # 28224 -> 128 + ) + self.decoder = torch.nn.Sequential( + torch.nn.LazyLinear(64), #128 + 2 -> 64 + torch.nn.ReLU(), + torch.nn.LazyLinear(64), #64 -> 64 + torch.nn.ReLU(), + torch.nn.LazyLinear(1) #64 -> 1 + ) + + @staticmethod + def _concatenate(state, action): + b, s= state.shape + b, a = action.shape + sa = torch.cat((state, action), -1) + return sa + + def forward(self, state, action): + s = self.cnn(state.float()) + sa = self._concatenate(s, action) + assert sa.ndim == 2 + return self.decoder(sa).squeeze(1) + +class MlpDiscriminator(torch.nn.Module): + """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy.""" + + def __init__(self, env=None): + super().__init__() + self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1) + self.mlp = torch.nn.Sequential( + torch.nn.LazyLinear(64), # 42 -> 64 + torch.nn.Tanh(), + torch.nn.LazyLinear(64), # 64 -> 64 + torch.nn.Tanh(), + torch.nn.LazyLinear(1), # 64 -> 1 + ) + + def forward(self, state, action): + flat = self.flatten(state) + sa = torch.cat((action, flat), -1) + assert sa.ndim == 2 + return self.mlp(sa).squeeze(1) diff --git a/src/util/__init__.py b/src/util/__init__.py index e69de29..000365b 100644 --- a/src/util/__init__.py +++ b/src/util/__init__.py @@ -0,0 +1 @@ +from src.util.render_env import * \ No newline at end of file diff --git a/src/util/render_env.py b/src/util/render_env.py new file mode 100644 index 0000000..8223c73 --- /dev/null +++ b/src/util/render_env.py @@ -0,0 +1,58 @@ + +import stable_baselines3 as sb3 +from intersim.envs.intersimple import NRasterized + + +def render_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized): + """ + Render a video from an model, agent, and environment + Args: + model_name (str): name of the model + agent (int): agent to start the video from + environment (gym.Env): gym environment class to render environment on + """ + + model = sb3.PPO.load(model_name) + + env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent) + + obs = env.reset() + i=0 + while True and i < 600: + i+=1 + 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+'_agent%i'%(agent)) + +def render_options_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized): + """ + Render a video from an model, agent, and environment + Args: + model_name (str): name of the model + agent (int): agent to start the video from + environment (gym.Env): gym environment class to render environment on + """ + + model = sb3.PPO.load(model_name) + + env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent) + + obs = env.reset() + i=0 + while True and i < 600: + i+=1 + 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+'_agent%i'%(agent)) + +if __name__ == '__main__': + import fire + fire.Fire(render_env) \ No newline at end of file diff --git a/tests/test_discriminator.py b/tests/test_discriminator.py new file mode 100644 index 0000000..1172aa3 --- /dev/null +++ b/tests/test_discriminator.py @@ -0,0 +1,45 @@ +from intersim.envs.intersimple import NRasterized +from src.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()