From 05b31092f4d92ace27c8c3f47a5d74802d40be6f Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 21 Oct 2021 05:59:28 -0700 Subject: [PATCH] moving options policy to policies, commenting options image, and making the calls to train more flexible --- .../etienne/intersimple/gail_options_image.py | 162 ++++++++++++------ src/policies/__init__.py | 1 + src/policies/options.py | 59 +++++++ 3 files changed, 174 insertions(+), 48 deletions(-) create mode 100644 src/policies/options.py diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py index fb2329f..150afe3 100644 --- a/scratch/etienne/intersimple/gail_options_image.py +++ b/scratch/etienne/intersimple/gail_options_image.py @@ -1,55 +1,38 @@ # %% -from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction +from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction +from src.policies import OptionsCnnPolicy + from imitation.algorithms import adversarial -import stable_baselines3 -import torch.utils.data -import numpy as np -from intersim.envs.intersimple import NRasterized -from intersim.collisions import state_to_polygon -import itertools -from torch.distributions import Categorical -import gym -import torch -import pickle +from imitation.util import logger import imitation.data.rollout as rollout + +import stable_baselines3 +from stable_baselines3.common.env_util import make_vec_env + +import torch +import torch.utils.data +from torch.distributions import Categorical +import numpy as np +import itertools +import gym +import pickle import tempfile import pathlib -from imitation.util import logger -from stable_baselines3.common.env_util import make_vec_env from tqdm import tqdm -model_name = 'gail_options_image' -env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2} +from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent, NRasterizedIncrementingAgent +from intersim.collisions import state_to_polygon 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, eps=1e-6): - s, m = obs['obs'], obs['mask'] - values, prior = self._prior_distribution(s) - posterior = Categorical((prior.probs + eps) * m) - ch = posterior.sample() - return ch, values, posterior.log_prob(ch) - - def evaluate_actions(self, obs, ch, eps=1e-6): - s, m = obs['obs'], obs['mask'] - values, prior = self._prior_distribution(s) - posterior = Categorical((prior.probs + eps) * m) - return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train - class OptionsEnv(gym.Wrapper): - + """ + Wrap an intersimple environment with an options generator + """ def __init__(self, env, *args, **kwargs): + """ + Initialize wrapped environment and set high-level action and observation spaces + """ super().__init__(env, *args, **kwargs) num_hl_options = len(ALL_OPTIONS) self.action_space = gym.spaces.Discrete(num_hl_options) @@ -68,6 +51,13 @@ class OptionsEnv(gym.Wrapper): raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.') def sample(self, generator): + """ + yield transitions using a generator + Args: + generator (sb3.PPO) + Yields: + + """ self.done = True while True: self.episode_start = False @@ -104,13 +94,23 @@ class LLOptions(OptionsEnv): """Sample low-level (state, action) tuples for discriminator training.""" def __init__(self, *args, **kwargs): + """ + LLOption uses the true LL observations + """ super().__init__(*args, **kwargs) + # overwrite observation space to just output obs directly self.observation_space = self.observation_space['obs'] def _after_choice(self): + """ + After each option choice, initialize/reset the transition buffer + """ self._transition_buffer = [] def _after_step(self): + """ + After each ll action, append s, s', a, done to transition buffer + """ self._transition_buffer.append({ 'obs': self.s, 'next_obs': self.nexts, @@ -119,9 +119,18 @@ class LLOptions(OptionsEnv): }) def _transitions(self): + """ + Yield from the transition buffer + """ yield from self._transition_buffer def sample_ll(self, policy): + """ + Args: + policy + Returns: + gen: iterable which samples low-level transitions from the environment + """ return self.sample(policy) class HLOptions(OptionsEnv): @@ -131,11 +140,17 @@ class HLOptions(OptionsEnv): super().__init__(*args, **kwargs) def _after_choice(self): + """ + After an option selection, initialize total reward and number of steps + """ self.obs = {'obs': np.copy(self.s), 'mask': np.copy(self.m)} self.r = 0 self.steps = 0 def _after_step(self): + """ + After each low-level action, add the discounted discriminated reward score (given a discriminator) + """ self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train( state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()), @@ -145,6 +160,18 @@ class HLOptions(OptionsEnv): self.steps += 1 def _transitions(self): + """ + Yield a single dictionary per high-level selected action + Fields: + obs: high-level state and mask at selection + action: chosen high-level action + reward: accumulated option reward + episode_start: whether the action was chosen at the episode start + value: the value estimate from the starting state + log_prob: the log_prob of the selected action from the starting state + done: whether the episode has ended + + """ yield { 'obs': self.obs, 'action': self.ch, @@ -156,16 +183,29 @@ class HLOptions(OptionsEnv): } def sample_hl(self, policy, discriminator): + """ + Args: + policy + discriminator: function with which to score rewards + Returns: + gen: iterable which samples high-level transitions from the environment + """ self.discriminator = discriminator return self.sample(policy) class RenderOptions(LLOptions): def _after_step(self): + """ + Render the environment after each low-level step + """ super()._after_step() self.env.render() def close(self, *args, **kwargs): + """ + On 'close', close the environment + """ self.env.close(*args, **kwargs) def available_actions(env): @@ -326,8 +366,20 @@ def train_generator(env, generator, discriminator, num_samples): generator.train() -def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99): - env = NRasterized(**env_settings) +def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99): + """ + Args: + expert_data: list of transitions + env_class: environment class + env_settings: environment settings + epochs: number of epochs to train for + discrim_batch_size: discriminator batch size + generator_steps: number of steps taken in generator + discount: discount factor + Returns: + generator (stable_baselines3.PPO): options policy + """ + env = env_class(**env_settings) env.discount = discount tempdir = tempfile.TemporaryDirectory(prefix="quickstart") @@ -335,10 +387,10 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di 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) + venv = make_vec_env(env_class, n_envs=1, env_kwargs=env_settings) discriminator = adversarial.GAIL( expert_data=expert_data, - expert_batch_size=expert_batch_size, + expert_batch_size=discrim_batch_size, discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)}, #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)}, venv=venv, # unused @@ -360,7 +412,7 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di ) for _ in tqdm(range(epochs)): - train_discriminator(LLOptions(env), generator, discriminator, num_samples=expert_batch_size) + train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size) train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps) return generator @@ -368,18 +420,32 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di # %% if __name__ == '__main__': # %% - + model_name = 'gail_options_image' + env_class = NRasterizedRandomAgent + env_settings = {'width': 36, 'height': 36, 'm_per_px': 2} + + #env_class = NRasterized + #env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2} + with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f: trajectories = pickle.load(f) transitions = rollout.flatten_trajectories(trajectories) - generator = train(transitions) + generator = train( + transitions, + env_class=env_class, + env_settings=env_settings, + epochs=10, + discrim_batch_size=32, + generator_steps=2048, + discount=0.99 + )) generator.save(model_name) # %% model = stable_baselines3.PPO.load(model_name) - env = RenderOptions(NRasterized(**env_settings)) + env = RenderOptions(NRasterizedRandomAgent(**env_args)) for s in env.sample_ll(model): if s['dones']: diff --git a/src/policies/__init__.py b/src/policies/__init__.py index cff41ff..0679815 100644 --- a/src/policies/__init__.py +++ b/src/policies/__init__.py @@ -1 +1,2 @@ from src.policies.policy import IntersimPolicy, IntersimStateNet, IntersimStateActionNet, generate_transforms +from src.policies.options import OptionsCnnPolicy \ No newline at end of file diff --git a/src/policies/options.py b/src/policies/options.py new file mode 100644 index 0000000..9bf71c9 --- /dev/null +++ b/src/policies/options.py @@ -0,0 +1,59 @@ +import stable_baselines3 +from torch.distributions import Categorical + +class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy): + """ + Class for high-level options policy (generator) + """ + def __init__(self, observation_space, *args, **kwargs): + super().__init__(observation_space['obs'], *args, **kwargs) + + def _prior_distribution(self, s): + """ + Return prior distribution over high-level options (before masking) + Args: + s (torch.tensor): observation + Returns: + values (torch.tensor): values from critic + dist (torch.distributions): prior distribution over actions + """ + 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, eps=1e-6): + """ + Will mask invalid states before making action selections + Args: + obs: dict with keys: + obs (torch.tensor): (*,o) true observations + mask (torch.tensor): (*,m) mask over valid actions + Returns: + ch (torch.tensor): (*,a) sampled actions + values (torch.tensor): (*,) predicted value at observation + log_probs (torch.tensor): (*,) log probabilities of selected actions + """ + s, m = obs['obs'], obs['mask'] + values, prior = self._prior_distribution(s) + posterior = Categorical((prior.probs + eps) * m) + ch = posterior.sample() + return ch, values, posterior.log_prob(ch) + + def evaluate_actions(self, obs, ch, eps=1e-6): + """ + Evaluate particular actions + Args: + obs: dict with keys: + obs (torch.tensor): (*,o) true observations + mask (torch.tensor): (*,m) masks over valid actions + ch (torch.tensor): (*,a) selected actions + Returns: + values (torch.tensor): (*,) predicted value at observation + log_probs (torch.tensor): (*,) log probabilities of selected actions + ent (torch.tensor): (*,) entropy of each distribution over actions + """ + s, m = obs['obs'], obs['mask'] + values, prior = self._prior_distribution(s) + posterior = Categorical((prior.probs + eps) * m) + return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train