511 lines
19 KiB
Python
511 lines
19 KiB
Python
# %%
|
|
from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
|
|
from imitation.algorithms import adversarial
|
|
import stable_baselines3
|
|
import torch.utils.data
|
|
import numpy as np
|
|
from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent
|
|
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
|
|
from tqdm import tqdm
|
|
|
|
import logging
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
|
|
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):
|
|
"""
|
|
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):
|
|
"""
|
|
Will mask invalid states before making action selections
|
|
Args:
|
|
obs: dict with keys:
|
|
obs (torch.tensor): (B,o) true observations
|
|
mask (torch.tensor): (B,m) mask over valid actions
|
|
Returns:
|
|
ch (torch.tensor): (B,a) sampled actions
|
|
values (torch.tensor): (B,) predicted value at observation
|
|
log_probs (torch.tensor): (B,) log probabilities of selected actions
|
|
"""
|
|
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):
|
|
"""
|
|
Evaluate particular actions
|
|
Args:
|
|
obs: dict with keys:
|
|
obs (torch.tensor): (B,o) true observations
|
|
mask (torch.tensor): (B,m) masks over valid actions
|
|
ch (torch.tensor): (B,a) selected actions
|
|
Returns:
|
|
values (torch.tensor): (B,) predicted value at observation
|
|
log_probs (torch.tensor): (B,) log probabilities of selected actions
|
|
ent (torch.tensor): (B,) entropy of each distribution over actions
|
|
"""
|
|
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
|
|
|
|
class OptionsEnv(gym.Wrapper):
|
|
"""
|
|
Wrap an intersimple environment with an options generator
|
|
"""
|
|
def __init__(self, env, render=False, *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)
|
|
self.observation_space = gym.spaces.Dict({
|
|
'obs': env.observation_space,
|
|
'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
|
|
})
|
|
self._hl_transition_buffer = []
|
|
self._ll_transition_buffer = []
|
|
self.render=render
|
|
|
|
def _after_option_choice(self):
|
|
"""
|
|
After initial option choice,
|
|
"""
|
|
self._hl_r = 0
|
|
self._hl_steps = 0
|
|
|
|
def _after_step(self):
|
|
"""
|
|
After each step, add the ll transition to the appropriate buffer, add to reward, add to steps, and possibly render
|
|
"""
|
|
|
|
self._ll_transition_buffer.append({
|
|
'obs': self.s,
|
|
'next_obs': self.nexts,
|
|
'acts': np.array((self.a,)),
|
|
'dones': np.array(self.done),
|
|
})
|
|
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()),
|
|
next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
|
|
done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
|
|
)
|
|
self.steps += 1
|
|
if self.render:
|
|
self.env.render()
|
|
|
|
def _after_option(self):
|
|
"""
|
|
After each low-level action, add the discounted discriminated reward score (given a discriminator)
|
|
"""
|
|
self._hl_transition_buffer.append({
|
|
'obs': {'obs': self.os, 'mask': self.m},
|
|
'action': self.ch,
|
|
'reward': self.r.detach(),
|
|
'episode_start': self.episode_start,
|
|
'value': self.value.detach(),
|
|
'log_prob': self.log_prob.detach(),
|
|
'done': self.done,
|
|
})
|
|
|
|
def close(self, *args, **kwargs):
|
|
"""
|
|
On 'close', close the environment
|
|
"""
|
|
self.env.close(*args, **kwargs)
|
|
|
|
def sample(self, generator, controller):
|
|
"""
|
|
yield transitions using a generator
|
|
Args:
|
|
generator (sb3.PPO)
|
|
controller (str): 'high' or 'low' to yield from proper buffer
|
|
Yields:
|
|
|
|
"""
|
|
self.done = True
|
|
# DO I WANT TO EMPTY THE BUFFERS??? Probs naw
|
|
while True:
|
|
|
|
# yield from buffers to empty what was stored previously
|
|
if controller = 'high':
|
|
yield from self._hl_transition_buffer
|
|
elif controller == 'low':
|
|
yield from self._ll_transition_buffer
|
|
else:
|
|
raise('Improper buffer')
|
|
|
|
self.episode_start = False
|
|
if self.done:
|
|
# reset environment
|
|
self.s = self.env.reset()
|
|
self.done = False
|
|
self.episode_start = True
|
|
|
|
self.os = self.s.copy() # option start state
|
|
self.m = available_actions(self.env)
|
|
|
|
# set the action, the value of the start state, and the logprob of the action
|
|
# according to the current environment state and mask
|
|
self.ch, self.value, self.log_prob = generator.policy.predict({
|
|
'obs': torch.tensor(self.os).unsqueeze(0).to(generator.policy.device),
|
|
'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
|
|
})
|
|
|
|
# store a float list of actions to take given the option selected in the environment
|
|
self.plan = list(map(float, generate_plan(self.env, self.ch)))
|
|
|
|
# run whatever _after_choice might dictate in a child class
|
|
self._after_option_choice()
|
|
|
|
# some checks
|
|
assert not self.done
|
|
assert self.plan
|
|
assert feasible(self.env, self.plan, self.ch)
|
|
|
|
# execute the option so long as the episode isn't complete and the plan is still feasible
|
|
while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
|
|
|
|
# pop first action
|
|
self.a, self.plan = self.plan[0], self.plan[1:]
|
|
|
|
# normalize action ??
|
|
self.a = self.env._normalize(self.a)
|
|
|
|
# step through environment
|
|
self.nexts, _, self.done, _ = self.env.step(self.a)
|
|
|
|
# run whatever _after_step might dictate in child class
|
|
self._after_step()
|
|
|
|
# update state and mask to current
|
|
self.s = self.nexts
|
|
|
|
# run whatever to do after option
|
|
self._after_option()
|
|
|
|
def sample_ll(self, policy):
|
|
"""
|
|
Not quite sure how this works????
|
|
Why would you do this over LLOptions.sample(policy)
|
|
"""
|
|
return self.sample(policy, 'low')
|
|
|
|
def sample_hl(self, policy, discriminator):
|
|
"""
|
|
Args:
|
|
policy
|
|
discriminator: function with which to score rewards
|
|
Returns:
|
|
gen: an which samples high-level transitions from the environment
|
|
"""
|
|
self.discriminator = discriminator
|
|
return self.sample(policy)
|
|
|
|
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 check_future_collisions_circles(env, actions, n_circles:int=2):
|
|
"""Checks whether `env._agent` would collide with other agents assuming `actions` as input.
|
|
|
|
Vehicles are (over-)approximated by multiple 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
|
|
"""
|
|
assert n_circles >= 2
|
|
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)
|
|
centers = states[:, :, :, :2]
|
|
psi = states[:, :, :, 3]
|
|
lon = torch.stack([psi.cos(), psi.sin()],dim=-1) # (B, T, nv, 2)
|
|
|
|
# offset between [-env._env.lengths+env._env.widths/2, env._env.lengths/2-env._env.widths/2]
|
|
back = (-env._env._lengths/2+env._env._widths/2).unsqueeze(-1) # (nv, 1)
|
|
length = (env._env._lengths-env._env._widths).unsqueeze(-1) # (nv, 1)
|
|
diff_d = back + length*(torch.arange(n_circles)/(n_circles-1)).unsqueeze(0) # (nv, n_circles)
|
|
assert diff_d.shape == (nv, n_circles)
|
|
|
|
offsets = diff_d[None, None, :, :, None] * lon[:, :, :, None, :]
|
|
assert offsets.shape == (B, T, nv, n_circles, 2)
|
|
|
|
expanded_centers=centers.unsqueeze(-2) + offsets #(B, T, nv, n_circles, 2)
|
|
assert expanded_centers.shape == (B, T, nv, n_circles, 2)
|
|
agent_centers = expanded_centers[:,:,env._agent:env._agent+1,:,:] #(B, T, 1, n_circles, 2)
|
|
ds = expanded_centers.reshape((B, T, nv*n_circles, 1, 2)) - agent_centers #(B, T, nv*nc,1, 2) - (B, T, 1, nc, 2) = (B, T, nv*nc, nc, 2)
|
|
|
|
distance = (ds**2).sum(-1).sqrt().reshape((B, T, nv, n_circles, n_circles)) # (B, T, nv, nc, nc)
|
|
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, n_circles, n_circles)
|
|
|
|
radius = env._env._widths*np.sqrt(2) / 2
|
|
min_distance = radius[env._agent] + radius
|
|
min_distance = min_distance[None, None, :, None, None]
|
|
assert min_distance.shape == (1, 1, nv, 1, 1)
|
|
|
|
return (distance > min_distance).all(-1).all(-1).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
|
|
valid = check_future_collisions_circles(env, [full_plan])
|
|
return ch == 0 or valid.item()
|
|
|
|
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(env.sample_ll(generator), 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(env.sample_hl(generator, discriminator), num_samples+1))
|
|
|
|
generator.rollout_buffer.reset()
|
|
for s in generator_samples[:-1]:
|
|
generator.rollout_buffer.add(
|
|
obs=s['obs'],
|
|
action=s['action'].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()
|
|
|
|
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")
|
|
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(env_class, n_envs=1, env_kwargs=env_settings)
|
|
discriminator = adversarial.GAIL(
|
|
expert_data=expert_data,
|
|
expert_batch_size=discrim_batch_size,
|
|
discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
|
|
#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 tqdm(range(epochs)):
|
|
train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size)
|
|
train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps)
|
|
|
|
return generator
|
|
|
|
# %%
|
|
if __name__ == '__main__':
|
|
# %%
|
|
model_name = 'gail_options_image'
|
|
env_class = NRasterizedRandomAgent
|
|
env_settings = {'width': 36, 'height': 36, 'm_per_px': 2}
|
|
|
|
with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl", "rb") as f:
|
|
trajectories = pickle.load(f)
|
|
#import pdb
|
|
#pdb.set_trace()
|
|
transitions = rollout.flatten_trajectories(trajectories)
|
|
generator = train(
|
|
transitions,
|
|
env_class=env_class,
|
|
env_settings=env_settings,
|
|
epochs=2,
|
|
discrim_batch_size=32,
|
|
generator_steps=2048,
|
|
discount=0.99
|
|
)
|
|
|
|
generator.save(model_name) # save ppo sb3 generator class
|
|
|
|
# %%
|
|
model = stable_baselines3.PPO.load(model_name) # not actually used
|
|
|
|
env = OptionsGail(NRasterizedRandomAgent(**env_settings), render=True)
|
|
for s in env.sample_ll(generator):
|
|
if s['dones']:
|
|
break
|
|
|
|
env.close(filestr='render/'+model_name)
|
|
|
|
# %% Tests
|
|
|
|
def test_ll_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 = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
|
|
|
|
gen_transitions = list(itertools.islice(env.sample_ll(
|
|
policy=stable_baselines3.PPO(
|
|
OptionsCnnPolicy,
|
|
OptionsEnv(env),
|
|
verbose=1,
|
|
)
|
|
), 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_ll_states():
|
|
env = NRasterized()
|
|
policy = stable_baselines3.PPO(
|
|
OptionsCnnPolicy,
|
|
OptionsEnv(env),
|
|
verbose=1,
|
|
)
|
|
llenv = LLOptions(env)
|
|
transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
|
|
|
|
env2 = NRasterized()
|
|
s2 = env2.reset()
|
|
for i, t in enumerate(transitions):
|
|
assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
|
|
assert np.array_equal(t['obs'], s2)
|
|
assert t['acts'].shape == (1,)
|
|
|
|
nexts2, _, done2, _ = env2.step(t['acts'])
|
|
assert np.array_equal(t['next_obs'], nexts2)
|
|
assert np.array_equal(t['dones'], done2)
|
|
|
|
if done2:
|
|
break
|
|
|
|
s2 = nexts2
|
|
|
|
def test_hl_transitions():
|
|
pass
|