removed scratch except for expert data generator, which i moved to the top

This commit is contained in:
Arec Jamgochian
2022-03-03 19:37:29 -08:00
parent fc04f8e9ee
commit 4076b0361a
145 changed files with 61 additions and 13708 deletions

View File

@@ -0,0 +1,61 @@
import torch
import functools
from src.core.sampling import rollout_sb3
from intersim.envs import IntersimpleLidarFlatIncrementingAgent
from intersim.envs.intersimple import speed_reward
from intersim.expert import NormalizedIntersimpleExpert
from src.util.wrappers import CollisionPenaltyWrapper, Setobs
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
def main(track:int, loc:int=0):
env = IntersimpleLidarFlatIncrementingAgent(
loc=loc,
track=track,
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
policy = NormalizedIntersimpleExpert(env, mu=0.001)
env = Setobs(TransformObservation(
CollisionPenaltyWrapper(
env,
collision_distance=6, collision_penalty=100
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
))
print(env.nv, 'vehicles')
expert_data = rollout_sb3(env, policy, n_episodes=150, max_steps_per_episode=200)
states, actions, rewards, dones = expert_data
print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
print(f'Observation mean', states[~dones].mean(0))
print(f'Observation std', states[~dones].std(0))
torch.save(expert_data, f'intersimple-expert-data-setobs2-loc{loc}-track{track}.pt')
if __name__=='__main__':
import fire
fire.Fire(main)

View File

@@ -1,22 +0,0 @@
python -m render_options --model_name='gail_options_image_mid_wcollision' --env='NRasterizedRoute' --options=True --width=36 --height=36 --m_per_px=2 --agent=50 --stop_on_collision=False
import torch, os
from src.data import load_experts
folder = 'expert_data/DR_USA_Roundabout_FT/track0000'
single_agent = os.path.join(folder, 'expert.pkl')
multi_agent = os.path.join(folder,'joint_expert_states.pt')
multi_agent_actions = os.path.join(folder,'joint_expert_actions.pt')
demonstrations = load_experts([single_agent], flatten=False)
demonstrations[0].__dict__.keys()
len(demonstrations[0].obs)
single_agent_lengths = [len(demonstration.obs) for demonstration in demonstrations]
states = torch.load(multi_agent)
actions = torch.load(multi_agent_actions)
multi_agent_lengths = [sum(~torch.isnan(states[:,i,0])).item() for i in range(states.shape[1])]
single_agent_actions = [demonstration.acts for demonstration in demonstrations]
multi_agent_actions = [actions[~torch.isnan(actions[:,i,0])] for i in range(actions.shape[1])]
import pickle
with open(single_agent, "rb") as f:
new_trajectories = pickle.load(f)

View File

@@ -1,145 +0,0 @@
from intersim.envs.intersimple import Intersimple
from stable_baselines3.common.policies import BasePolicy
import gym
import intersim.envs.intersimple
import imitation.data.rollout as rollout
from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
from imitation.data.wrappers import RolloutInfoWrapper
class IntersimExpert(BasePolicy):
def __init__(self, intersim_env, mu=0, *args, **kwargs):
super().__init__(
observation_space=gym.spaces.Space(),
action_space=gym.spaces.Space(),
*args, **kwargs
)
self._intersim = intersim_env
self._mu = mu
def forward(self, *args, **kwargs):
raise NotImplementedError()
def _predict(self, *args, **kwargs):
raise NotImplementedError()
def _action(self):
target_t = min(self._intersim._ind + 1, len(self._intersim._svt.simstate) - 1)
target_state = self._intersim._svt.simstate[target_t]
return self._intersim.target_state(target_state, mu=self._mu)
def predict(self, *args, **kwargs):
return self._action(), None
class IntersimpleExpert(BasePolicy):
def __init__(self, intersimple_env, mu=0, *args, **kwargs):
super().__init__(
observation_space=intersimple_env.observation_space,
action_space=intersimple_env.action_space,
*args, **kwargs
)
self._intersimple = intersimple_env
self._intersim_expert = IntersimExpert(intersimple_env._env, mu=mu)
def forward(self, *args, **kwargs):
raise NotImplementedError()
def _predict(self, *args, **kwargs):
raise NotImplementedError()
def _action(self):
return self._intersim_expert._action()[self._intersimple._agent]
def predict(self, *args, **kwargs):
return self._action(), None
class NormalizedIntersimpleExpert(IntersimpleExpert):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def predict(self, *args, **kwargs):
action, _ = super().predict(*args, **kwargs)
return self._intersimple._normalize(action), None
class DummyVecEnvPolicy(BasePolicy):
def __init__(self, experts):
self._experts = [e() for e in experts]
def forward(self, *args, **kwargs):
raise NotImplementedError()
def _predict(self, *args, **kwargs):
raise NotImplementedError()
def predict(self, *args, **kwargs):
predictions = [e.predict() for e in self._experts]
actions = [p[0] for p in predictions]
states = [p[1] for p in predictions]
return actions, states
def forward(self, *args, **kwargs):
raise NotImplementedError()
def _predict(self, *args, **kwargs):
raise NotImplementedError()
def save_video(env, expert):
env.reset()
env.render()
done = False
while not done:
actions, _ = expert.predict()
_, _, done, _ = env.step(actions)
env.render()
env.close()
def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedIncrementingAgent', path=None, min_timesteps=None, min_episodes=None, video=False, env_args={}, policy_args={}):
"""Rollout and save expert demos.
Usage:
python -m intersimple.expert <flags>
Args:
expert (class): class of expert
env (class): class of env intersim.envs.intersimple
path (str): path to store output
min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save
min_episodes (int): min number of episodes for call to rollout.rollout_and_save
video (bool): whether to save a video of the expert until a single environment instantiation stops
env_args (dict): dictionary of kwargs when instantiating environment class
policy_args (dict): dictionary of kwargs when instantiating Expert policy
"""
Env = intersim.envs.intersimple.__dict__[env]
Expert = globals()[expert]
env = Env(**env_args)
info_env = RolloutInfoWrapper(env) # getting rollout info (dictionary) from environment
venv = DummyVecEnv([lambda: info_env]) # making a DummyVecEnv with a list of a function that when called returns the rollout info
policy = Expert(env, **policy_args) # instantiate an expert policy from specified class with instantiated environment and policy kwargs
venv_policy = DummyVecEnvPolicy([lambda: policy]) # make a DummyVecEnvPolicy with a list of a function that when called returns the Expert policy
if min_timesteps is None and min_episodes is None:
min_episodes = env.nv # one episode per vehicle being controlled in environment (hopefully an incrementing agent environment)
if video:
save_video(env, policy)
path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
suntil = rollout.make_sample_until(
min_timesteps=min_timesteps,
min_episodes=min_episodes,
)
rollout.rollout_and_save(
path=path,
policy=venv_policy,
venv=venv,
sample_until=suntil
)
if __name__ == '__main__':
import fire
fire.Fire(demonstrations)

View File

@@ -1,9 +0,0 @@
#python -m 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 expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl'
#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'
# python -m expert --env=NRasterizedRandomAgent --min_timesteps=10000 --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.pkl'
#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'
#python -m expert --env=NRasterized --min_timesteps=3000 --video --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl'
#python -m expert --env=NRasterizedIncrementingAgent --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl'
python -m process_all_experts --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}'

View File

@@ -1,23 +0,0 @@
import pickle
import imitation.data.rollout as rollout
from tqdm import tqdm
def load_experts(expert_files=[]):
"""
Load expert trajectories from files and combine their transitions into a single RB
Args:
expert_files (list): list of expert file strings
Returns:
transitions (list): list of combined expert episode transitions
"""
transitions = []
for file in tqdm(expert_files):
with open(file, "rb") as f:
trajectories = pickle.load(f)
transitions = transitions + rollout.flatten_trajectories(trajectories)
return transitions
if __name__=='__main__':
import fire
fire.Fire(load_experts)

View File

@@ -1,48 +0,0 @@
import tqdm
import expert
import copy
import os
import intersim
from tqdm import tqdm
def process_all_experts(filename='expert.pkl',env_args={}, policy_args={}):
"""
Process all experts in the Interaction Dataset
For now, using NormalizedIntersimpleExpert with NRasterizedIncrementingAgent environment
Args:
filename (str): name for track file
env_args (dict): default environment kwargs
policy_args (dict): default policy kwargs
"""
I, J = len(intersim.LOCATIONS), intersim.MAX_TRACKS
pbar = tqdm(total=I*J)
for loc in range(I):
for track in range(J):
it_env_args = copy.deepcopy(env_args)
it_env_args.update({
'loc':loc,
'track':track,
})
out_folder = os.path.join(intersim.LOCATIONS[loc], 'track%04i'%(track))
if not os.path.isdir(out_folder):
os.makedirs(out_folder)
it_path = os.path.join(out_folder,filename)
expert.demonstrations(
expert='NormalizedIntersimpleExpert',
env='NRasterizedIncrementingAgent',
path=it_path,
env_args=it_env_args,
policy_args=policy_args,
)
pbar.update(1)
pbar.close()
if __name__=='__main__':
import fire
fire.Fire(process_all_experts)

View File

@@ -1,101 +0,0 @@
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)

View File

@@ -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()

View File

@@ -1,70 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import NRasterized
from gail.discriminator import CnnDiscriminatorFlatAction
model_name = 'gail_image_multiagent_nocollision'
# %%
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2})
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# 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/")
gail_trainer = adversarial.GAIL(
venv,
expert_data=transitions,
expert_batch_size=32,
#n_disc_updates_per_round=2048,
discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
allow_variable_horizon=True,
)
gail_trainer.train(total_timesteps=100000)
gail_trainer.gen_algo.save(model_name)
#del gail_trainer
# %%
model = sb3.PPO.load(model_name)
env = NRasterized(stop_on_collision=False, width=36, height=36, m_per_px=2)
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)

View File

@@ -1,70 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import NRasterized
from gail.discriminator import CnnDiscriminator
model_name = 'gail_image_singleagent_nocollision'
# %%
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'agent':51, 'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2})
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# 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/")
gail_trainer = adversarial.GAIL(
venv,
expert_data=transitions,
expert_batch_size=32,
#n_disc_updates_per_round=2048,
discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
allow_variable_horizon=True,
)
gail_trainer.train(total_timesteps=100000)
gail_trainer.gen_algo.save(model_name)
#del gail_trainer
# %%
model = sb3.PPO.load(model_name)
env = NRasterized(agent=51, width=36, height=36, m_per_px=2, stop_on_collision=False)
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)

View File

@@ -1,172 +0,0 @@
# %%
import sys
sys.path.append('../../../')
from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
from src.policies import OptionsCnnPolicy
from src.util import render_env
from src.data import load_experts
from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
from src.gail.train import train_discriminator, train_generator
from imitation.algorithms import adversarial
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
import numpy as np
import itertools
import gym
import pickle
import tempfile
import pathlib
from tqdm import tqdm
from intersim.envs.intersimple import NRasterized, NRasterizedRoute, NRasterizedRandomAgent, NRasterizedIncrementingAgent, NRasterizedRouteRandomAgent
ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
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(expert_data, env_class=NRasterizedRouteRandomAgent, 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, options=ALL_OPTIONS),
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, options=ALL_OPTIONS), generator, discriminator, num_samples=discrim_batch_size)
train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
return generator
# %%
if __name__ == '__main__':
# %%
model_name = 'gail_options_image_mid_wcollision'
env_class = NRasterizedRouteRandomAgent
env_settings = {'width': 36, 'height': 36, 'm_per_px': 2, 'stop_on_collision': False}
#env_class = NRasterized
#env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
files = ['../../../expert_data/DR_USA_Roundabout_FT/track%04i/expert.pkl'%(i) for i in range(5)]
transitions=load_experts(files)
generator = train(
transitions,
env_class=env_class,
env_settings=env_settings,
epochs=2,
discrim_batch_size=256,
generator_steps=10,#256,
discount=0.99
)
generator.save(model_name)
# Render
render_settings = {'width': 36, 'height': 36, 'm_per_px': 2, 'agent':51, 'stop_on_collision': False}
render_env(model_name=model_name, env='NRasterizedRoute', options=True, options_list=ALL_OPTIONS,
**render_settings)
# %% 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

View File

@@ -1,559 +0,0 @@
# %%
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, *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,)),
})
def _after_choice(self):
pass
def _after_step(self):
pass
def _transitions(self):
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
if self.done:
# reset environment
self.s = self.env.reset()
self.m = available_actions(self.env)
self.done = False
self.episode_start = True
# 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.s).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_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)
self.nextm = available_actions(self.env)
# run whatever _after_step might dictate in child class
self._after_step()
# update state and mask to current
self.s = self.nexts
self.m = self.nextm
# transitions yielded from self._transitions() functions specied in child classes
yield from self._transitions()
### NOTE: only yields after a full option has been executed / exited
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,
'acts': np.array((self.a,)),
'dones': np.array(self.done),
})
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):
"""Sample high-level (state, action, reward) tuples for generator training."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def _after_choice(self):
"""
After an option selection, initialize total reward and number of steps
"""
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()),
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
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': {'obs': self.s, '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 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):
"""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 = RenderOptions(NRasterizedRandomAgent(**env_settings))
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

View File

@@ -1,510 +0,0 @@
# %%
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

View File

@@ -1,55 +0,0 @@
Environment
-- each 'environment' follows a single roundabout and track id (recording of that roundabout)
-- on reset, the environment we will use changes the vehicle to control while having the other agents follow their true data (expert controller)
---- Note this can be problematic as it can lead to vehicles behind you crashing into you
TRAINING
---------
1. Load pre-trained massive set of transitions
-- For all roundabouts
-- For all tracks
-- For all vehicles
-- For all valid timesteps
-- Rasterized state (incl. path), action
2. HGAIL
-- For each epoch
-- INSTANTIATE A NEW ENVIRONMENT (Roundabout + Track) w/ randomized agent, from set of all expert environments
-- Train discriminator off training data + yielded low-level transitions in replay buffer
-- Train generator off yielded high-level transitions + summed low-level discriminator rewards
TESTING
----------
1. Save average vehicle velocities for all expert vehicles (loop roundabout + track + vehicle, average over time)
2. Run test suite for: expert, BC, GAIL, RAIL, HGAIL, (and hopefully HRAIL)
-- For all roundabouts, tracks
-- Get expert velocities for track
-- Simulate incrementing agent environment (e.g. on reset, agent +=1)
-- Store low-level true joint states, actions, and controlled vehicle index
-- Per-vehicle statistics (v_all, v_mean, v_shortfall, a_all, jerk_all, n_collisions, T)
-- Aggregate statistics + joint
Problems
-----------
Should train without stopping for collisions, however when doing so, end up with policy that always takes decelerate option
-- It seems safe at the start of each vehicles sim, but actually it isn't since a car will spawn and hit it
Solutions:
-- Hold cars from spawning if their spawn location is full
-- Start simulations a few seconds later (after cars clear their spawn places) <- Preferred
Test could run indefinitely if stop_on_collision is off
Solution:
-- Set maximum episode length in intersimple
Save massive set of transition raw states beforehand (1 from training, but with raw states)
# -- For all roundabouts, tracks
# -- For all vehicles, steps
# -- Raw vehicle state, action

View File

@@ -1,58 +0,0 @@
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)

View File

@@ -1,11 +0,0 @@
import sys
sys.path.append('../../../')
from src.util import render_env
ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]]
def render_wrapper(**kwargs):
render_env(**kwargs, options_list=ALL_OPTIONS)
if __name__=='__main__':
import fire
fire.Fire(render_wrapper)

View File

View File

@@ -1,64 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import IntersimpleReward
model_name = 'airl_flat'
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# Train AIRL 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 / "AIRL/")
airl_trainer = adversarial.AIRL(
venv,
expert_data=transitions,
expert_batch_size=64,
gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024), # n_steps = 2048 ?
)
airl_trainer.train(total_timesteps=100000)
airl_trainer.gen_algo.save(model_name)
del airl_trainer
# %%
model = sb3.PPO.load(model_name)
env = IntersimpleReward(agent=51)
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)

Binary file not shown.

View File

@@ -1,59 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import IntersimpleReward
model_name = 'bc_flat'
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# Train BC on expert data.
# BC also accepts as `expert_data` any PyTorch-style DataLoader that iterates over
# dictionaries containing observations and actions.
logger.configure(tempdir_path / "BC/")
bc_trainer = bc.BC(venv.observation_space, venv.action_space, expert_data=transitions)
bc_trainer.train(n_epochs=1000)
bc_trainer.save_policy(model_name)
del bc_trainer
# %%
model = bc.reconstruct_policy(model_name)
env = IntersimpleReward(agent=51)
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)

View File

@@ -1,138 +0,0 @@
from intersim.envs.intersimple import Intersimple, InfoFilter
from stable_baselines3.common.policies import BasePolicy
import gym
from intersim.envs.intersimple import *
from gail.envs import *
import imitation.data.rollout as rollout
from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
from imitation.data.wrappers import RolloutInfoWrapper
class IntersimExpert(BasePolicy):
def __init__(self, intersim_env, mu=0, *args, **kwargs):
super().__init__(
observation_space=gym.spaces.Space(),
action_space=gym.spaces.Space(),
*args, **kwargs
)
self._intersim = intersim_env
self._mu = mu
def forward(self, *args, **kwargs):
raise NotImplementedError()
def _predict(self, *args, **kwargs):
raise NotImplementedError()
def _action(self):
target_t = min(self._intersim._ind + 1, len(self._intersim._svt.simstate) - 1)
target_state = self._intersim._svt.simstate[target_t]
return self._intersim.target_state(target_state, mu=self._mu)
def predict(self, *args, **kwargs):
return self._action(), None
class IntersimpleExpert(BasePolicy):
def __init__(self, intersimple_env, mu=0, *args, **kwargs):
super().__init__(
observation_space=intersimple_env.observation_space,
action_space=intersimple_env.action_space,
*args, **kwargs
)
self._intersimple = intersimple_env
self._intersim_expert = IntersimExpert(intersimple_env._env, mu=mu)
def forward(self, *args, **kwargs):
raise NotImplementedError()
def _predict(self, *args, **kwargs):
raise NotImplementedError()
def _action(self):
# RandomLocation mixin re-initializes the intersim sub-env
self._intersim_expert._intersim = self._intersimple._env
return self._intersim_expert._action()[self._intersimple._agent]
def predict(self, *args, **kwargs):
return self._action(), None
class NormalizedIntersimpleExpert(IntersimpleExpert):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def predict(self, *args, **kwargs):
action, _ = super().predict(*args, **kwargs)
return self._intersimple._normalize(action), None
class DummyVecEnvPolicy(BasePolicy):
def __init__(self, experts):
self._experts = [e() for e in experts]
def forward(self, *args, **kwargs):
raise NotImplementedError()
def _predict(self, *args, **kwargs):
raise NotImplementedError()
def predict(self, *args, **kwargs):
predictions = [e.predict() for e in self._experts]
actions = [p[0] for p in predictions]
states = [p[1] for p in predictions]
return actions, states
def forward(self, *args, **kwargs):
raise NotImplementedError()
def _predict(self, *args, **kwargs):
raise NotImplementedError()
def save_video(env, expert):
env.reset()
env.render()
done = False
while not done:
actions, _ = expert.predict()
_, _, done, _ = env.step(actions)
env.render()
env.close()
def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomAgent', path=None, min_timesteps=25000, min_episodes=None, video=False, env_args={}, policy_args={}):
"""Rollout and save expert demos.
Usage:
python -m intersimple.expert <flags>
"""
Env = globals()[env]
Expert = globals()[expert]
env = Env(**env_args)
info_env = RolloutInfoWrapper(env)
venv = DummyVecEnv([lambda: info_env])
policy = Expert(env, **policy_args)
venv_policy = DummyVecEnvPolicy([lambda: policy])
if video:
save_video(env, policy)
path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
include_infos = isinstance(env, InfoFilter)
rollout.rollout_and_save(
path=path,
policy=venv_policy,
venv=venv,
sample_until=rollout.make_sample_until(
min_timesteps=min_timesteps,
min_episodes=min_episodes,
),
exclude_infos=not include_infos,
)
if __name__ == '__main__':
import fire
fire.Fire(demonstrations)

View File

@@ -1,15 +0,0 @@
#python -m 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 expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl'
#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'
#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'
#python -m expert --env=NRasterized --min_timesteps=3000 --video --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl'
#python -m expert --env=NRasterizedRandomAgent --min_timesteps=200 --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl'
#python -m expert --env=NRasterizedRandomAgent --min_timesteps=10000 --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.pkl'
#python -m expert --env=NRasterizedRouteRandomAgent --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteRandomAgentw70h70mppx1.pkl'
#python -m expert --env=NRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_NRasterizedRouteRandomAgentLocationw70h70mppx1.pkl'
#python -m expert --env=NRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1,map_color:128}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_NRasterizedRouteRandomAgentLocationw70h70mppx1mapc128.pkl'
#python -m expert --env=NRasterizedRouteSpeedRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,map_color:128,mu:0.001}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001.pkl'
#python -m data.expert --env=NRasterizedRouteSpeedRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,map_color:128,mu:0.001,skip_frames:5}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001skip5.pkl'
#python -m data.expert --env=TLNRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1,mu:0.001,random_skip:True,max_episode_steps:50}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl'
python -m data.expert --env=TLNRasterizedRouteRandomAgentLocation --min_timesteps=50000 --env_args='{width:70,height:70,m_per_px:1,mu:0.001,random_skip:True,max_episode_steps:50}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl'

View File

@@ -1,46 +0,0 @@
import gym
from gym.wrappers.time_limit import TimeLimit
import numpy as np
from intersim.envs.intersimple import NRasterizedRouteRandomAgentLocation, RandomLocation, RandomAgent, RewardVisualization, Reward, \
ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedObservation, \
NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple
class RasterizedSpeed:
def __init__(self, max_speed=12, *args, **kwargs):
super().__init__(*args, **kwargs)
channels, height, width = self.observation_space.shape
self.observation_space = gym.spaces.Box(
low=0,
high=255,
shape=(channels+1, height, width),
dtype=np.uint8
)
self._max_speed = max_speed
def _simple_obs(self, intersim_obs, intersim_info):
img = super()._simple_obs(intersim_obs, intersim_info)
ego_speed = intersim_obs['state'][self._agent, 2]
scaled_speed = (255 * ego_speed) // self._max_speed
speed_layer = scaled_speed * np.ones_like(img[:1], dtype=np.uint8)
speed_layer = speed_layer.clamp(0, 255)
obs = np.concatenate((img, speed_layer), axis=0)
return obs
class NRasterizedRouteSpeedRandomAgentLocation(RandomLocation, RandomAgent, RewardVisualization,
Reward, ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedSpeed, RasterizedObservation,
NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple):
pass
class TransparentTimeLimit(TimeLimit):
def __getattr__(self, name):
return getattr(self.env, name)
def close(self, *args, **kwargs):
return self.env.close(*args, **kwargs)
def TLNRasterizedRouteRandomAgentLocation(max_episode_steps, *args, **kwargs):
return TransparentTimeLimit(NRasterizedRouteRandomAgentLocation(*args, **kwargs), max_episode_steps=max_episode_steps)

View File

@@ -1,127 +0,0 @@
import gym
import torch
from src.util.collisions import feasible
import numpy as np
from collections import deque
def imitation_discriminator(discriminator):
return lambda obs, action, next_obs, done: discriminator.discrim_net.predict_reward_train(
state=torch.tensor(obs).unsqueeze(0).to(discriminator.discrim_net.device()),
action=torch.tensor([[action]]).to(discriminator.discrim_net.device()),
next_state=torch.tensor(next_obs).unsqueeze(0).to(discriminator.discrim_net.device()), # unused
done=torch.tensor(done).unsqueeze(0).to(discriminator.discrim_net.device()), # unused
).item()
class OptionsEnv(gym.Wrapper):
def __init__(self, env, options, discriminator, discount, ll_buffer, *args, **kwargs):
super().__init__(env, *args, **kwargs)
self.options = options
num_hl_options = len(self.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.discriminator = discriminator
self.discount = discount
self.ll_buffer = ll_buffer
@staticmethod
def _hl_observation(obs, mask):
return {
'obs': obs,
'mask': mask,
}
def reset(self):
self.done = False
self.obs = self.env.reset()
self.m = available_actions(self.env, self.options)
return self._hl_observation(self.obs, self.m)
def _ll_step(self, action):
return self.env.step(action)
def step(self, action):
assert self.m[action]
assert not self.done
plan = list(map(float, generate_plan(self.env, action, self.options)))
reward = 0
steps = 0
while not self.done and plan and \
(feasible(self.env, safety_plan(self.env, plan)) or self.m.sum() == 1):
a, plan = plan[0], plan[1:]
a = self.env._normalize(a)
next_obs, _, self.done, info = self._ll_step(a)
reward += self.discount**steps * self.discriminator(self.obs, a, next_obs, self.done)
self.ll_buffer.append({
'obs': self.obs,
'next_obs': next_obs,
'acts': np.array((a,)),
'dones': np.array(self.done),
})
steps += 1
self.obs = next_obs
self.m = available_actions(self.env, self.options)
return self._hl_observation(self.obs, self.m), reward, self.done, info
class RenderOptions(OptionsEnv):
def __init__(self, env, options, *args, **kwargs):
super().__init__(env, options, discriminator=lambda s, a, n, d: 0, discount=1, ll_buffer=deque(maxlen=0), *args, **kwargs)
def _ll_step(self, action):
out = super()._ll_step(action)
self.env.render(mode='post')
return out
def close(self, *args, **kwargs):
self.env.close(*args, **kwargs)
def safety_plan(env, plan):
return np.concatenate((plan, np.array(5 * [env._env._min_acc])), axis=0)
def available_actions(env, options):
"""Return mask of available actions given current `env` state.
Action 0 is considered safe fallback.
"""
plans = [generate_plan(env, i, options) for i, _ in enumerate(options)]
# is emergency braking still possible?
plans = list(map(lambda p: safety_plan(env, p), plans))
T = max(len(p) for p in plans)
plans = [np.pad(p, ((0, T-len(p)),), constant_values=np.nan) for p in plans]
plans = np.stack(plans, axis=0)
valid = feasible(env, plans)
if not valid.any():
valid[0] = True
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, options):
"""Generate input profile for high-level action `i`."""
assert i < len(options), "Invalid option index {i}"
target_v, t = 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

View File

@@ -1,70 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import IntersimpleReward
from gail.discriminator import MlpDiscriminator
model_name = 'gail_flat'
# %%
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# 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/")
gail_trainer = adversarial.GAIL(
venv,
expert_data=transitions,
expert_batch_size=150,
n_disc_updates_per_round=32,
discrim_kwargs={'discrim_net': MlpDiscriminator()},
gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=4530),
allow_variable_horizon=True,
)
gail_trainer.train(total_timesteps=400000)
gail_trainer.gen_algo.save(model_name)
#del gail_trainer
# %%
model = sb3.PPO.load(model_name)
env = IntersimpleReward(agent=51)
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)

View File

@@ -1,115 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import IntersimpleReward, speed_reward
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_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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
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):
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'],
n_disc_updates_per_round=config['n_disc_updates_per_round'],
discrim_kwargs={'discrim_net': MlpDiscriminator()},
gen_algo=generator,
allow_variable_horizon=True,
)
def callback(epoch):
print("callback")
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, return_episode_rewards=True)
tune.report(
reward=np.mean(episode_rewards),
length=np.mean(episode_lengths),
training_iteration=epoch,
)
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=40000, callback=callback)
analysis = tune.run(
training_function,
config = {
'expert_batch_size': tune.randint(1, 22), #220,
'n_disc_updates_per_round': tune.randint(2, 100), #16,
'n_steps': tune.randint(1, 10000), #4096,
},
resources_per_trial={
'cpu': 1,
# 'gpu': 1,
},
local_dir='ray',
num_samples=10,
)
print('Best config', analysis.get_best_config(metric='progress', mode='max'))
#del gail_trainer
# %%
model = sb3.PPO.load(model_name)
env = IntersimpleReward(agent=51)
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)

View File

@@ -1,70 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import NRasterized
from gail.discriminator import CnnDiscriminator
model_name = 'gail_image'
# %%
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (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})
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# 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/")
gail_trainer = adversarial.GAIL(
venv,
expert_data=transitions,
expert_batch_size=32,
#n_disc_updates_per_round=2048,
discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
allow_variable_horizon=True,
)
gail_trainer.train(total_timesteps=100000)
gail_trainer.gen_algo.save(model_name)
#del gail_trainer
# %%
model = sb3.PPO.load(model_name)
env = NRasterized(agent=51, width=36, height=36, m_per_px=2)
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)

View File

@@ -1,70 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import NRasterized
from gail.discriminator import CnnDiscriminatorFlatAction
model_name = 'gail_image_multiagent_nocollision'
# %%
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2})
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# 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/")
gail_trainer = adversarial.GAIL(
venv,
expert_data=transitions,
expert_batch_size=32,
#n_disc_updates_per_round=2048,
discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
allow_variable_horizon=True,
)
gail_trainer.train(total_timesteps=100000)
gail_trainer.gen_algo.save(model_name)
#del gail_trainer
# %%
model = sb3.PPO.load(model_name)
env = NRasterized(stop_on_collision=False, width=36, height=36, m_per_px=2)
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)

View File

@@ -1,79 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import NRasterizedRandomAgent, IntersimpleReward, speed_reward
import functools
from stable_baselines3.common.evaluation import evaluate_policy
from gail.discriminator import CnnDiscriminator
model_name = 'gail_image_random'
# %%
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
env_kwargs = {'width': 36, 'height': 36, 'm_per_px': 2}
venv = make_vec_env(NRasterizedRandomAgent, n_envs=2, env_kwargs=env_kwargs)
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# 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/")
generator = sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024)
gail_trainer = adversarial.GAIL(
venv,
expert_data=transitions,
expert_batch_size=32,
#n_disc_updates_per_round=2048,
discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
gen_algo=generator,
allow_variable_horizon=True,
)
def callback(round):
eval_env = NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs)
#sync_envs_normalization(self.training_env, self.eval_env)
episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
gail_trainer.train(total_timesteps=100000, callback=callback)
gail_trainer.gen_algo.save(model_name)
#del gail_trainer
# %%
model = sb3.PPO.load(model_name)
env = NRasterizedRandomAgent(width=36, height=36, m_per_px=2)
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)

View File

@@ -1,171 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import os
import random
import numpy as np
import torch
# set up ray tune
import ray
from ray import tune
from ray.tune import Analysis, ExperimentAnalysis
from ray.tune.schedulers import ASHAScheduler
from ray.tune.suggest.hyperopt import HyperOptSearch
from ray.tune.suggest import ConcurrencyLimiter
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import NRasterizedRandomAgent, IntersimpleReward, speed_reward, NRasterized, NRasterizedRandomAgentVerbose
import functools
from stable_baselines3.common.evaluation import evaluate_policy
from gym.wrappers import TimeLimit
from gail.discriminator import CnnDiscriminator
model_name = 'gail_image_random_ray'
env_kwargs={'width': 36, 'height': 36, 'm_per_px': 2}
# %%
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--outdir", help="result directory", default='ray')
parser.add_argument("--test", help="test run", default=False, action="store_true")
args = parser.parse_args()
outdir = args.outdir
# %%
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.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.
trajectories = pickle.load(f)
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
# Store transitions in shared ray memory
ray_transitions = ray.put(transitions)
# %%
venv = make_vec_env(NRasterizedRandomAgent, n_envs=2, env_kwargs=env_kwargs)
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
logger.configure(tempdir_path / "GAIL/")
def get_ray_config(test=False):
if test:
return {
'expert_batch_size': 2,
'ppo_n_steps': 2,
'ppo_batch_size': 2,
'ppo_n_epochs': 1,
'total_timesteps': 10,
}
else:
return {
'expert_batch_size': tune.choice([2**x for x in range(6,10)]),
'ppo_n_steps': tune.choice([2048, 3072, 4096]),
'ppo_batch_size': tune.choice([2**x for x in range(9,13)]),
'ppo_n_epochs': tune.choice([6,10]),
'total_timesteps': 400_000,
}
def ray_train(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.
discriminator = CnnDiscriminator(venv)
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(
"CnnPolicy", venv, verbose=0,
n_steps=config["ppo_n_steps"],
batch_size=config["ppo_batch_size"],
n_epochs=config["ppo_n_epochs"]
)
gail_trainer = adversarial.GAIL(
venv,
expert_data=ray.get(ray_transitions),
expert_batch_size=config["expert_batch_size"],
#n_disc_updates_per_round=2048,
discrim_kwargs={'discrim_net': discriminator},
gen_algo=generator,
allow_variable_horizon=True,
)
def callback(round):
# eval_env = NRasterized(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs)
eval_env = TimeLimit(NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs), max_episode_steps=1000)
episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
tune.report(
reward=np.mean(episode_rewards),
length=np.mean(episode_lengths),
training_iteration=round,
)
with tune.checkpoint_dir(step=round) 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=config['total_timesteps'], callback=callback)
ray_config = get_ray_config(args.test)
search = HyperOptSearch(ray_config, metric='length', mode="max",)
search = ConcurrencyLimiter(search, max_concurrent=10)
custom_scheduler = ASHAScheduler(time_attr='training_iteration', metric='length', mode="max", grace_period=15)
analysis = tune.run(
ray_train,
# config=ray_config,
search_alg=search,
scheduler=custom_scheduler,
local_dir=outdir,
resources_per_trial={"cpu":10, "gpu": 0.2},
num_samples=1 if args.test else 100,
)
del analysis
# %%
# outdir = "ray/ray_train_2021-09-20_13-33-50/ray_train_f06785b0_33_expert_batch_size=128,ppo_batch_size=1024,ppo_n_epochs=6,ppo_n_steps=2048,total_timesteps=400000_2021-09-20_15-52-05"
# %%
analysis = Analysis(outdir, default_metric="length", default_mode="max")
filepath = analysis.get_best_logdir()
print("Best ray experiment:", filepath)
config = analysis.get_best_config()
print("Best config:", config)
# %%
model = sb3.PPO.load(os.path.join(analysis.get_last_checkpoint(), 'gen_checkpoint'))
# env = NRasterized(agent=51, **env_kwargs)
env = TimeLimit(NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs), max_episode_steps=1000)
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.env.close(filestr='render/'+model_name)
# %%

View File

@@ -1,70 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import NRasterized
from gail.discriminator import CnnDiscriminator
model_name = 'gail_image_singleagent_nocollision'
# %%
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'agent':51, 'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2})
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# 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/")
gail_trainer = adversarial.GAIL(
venv,
expert_data=transitions,
expert_batch_size=32,
#n_disc_updates_per_round=2048,
discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
allow_variable_horizon=True,
)
gail_trainer.train(total_timesteps=100000)
gail_trainer.gen_algo.save(model_name)
#del gail_trainer
# %%
model = sb3.PPO.load(model_name)
env = NRasterized(agent=51, width=36, height=36, m_per_px=2, stop_on_collision=False)
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)

View File

@@ -1,90 +0,0 @@
# %%
import sys
sys.path.append('../../../')
from src.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
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
from src.policies.options import OptionsCnnPolicy
from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
from src.gail.train import train_discriminator, train_generator
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]] # option 0 is safe fallback
def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, 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': 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, options=ALL_OPTIONS),
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, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
train_generator(HLOptions(env, options=ALL_OPTIONS), 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.save(model_name)
# %%
model = stable_baselines3.PPO.load(model_name)
env = RenderOptions(NRasterized(**env_settings), options=ALL_OPTIONS)
for s in env.sample_ll(model):
if s['dones']:
break
env.close(filestr='render/'+model_name)

View File

@@ -1,410 +0,0 @@
# %%
import sys
sys.path.append('../../../')
from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
from src.policies import OptionsCnnPolicy
from src.util import feasible
from src.data import load_experts
from imitation.algorithms import adversarial
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
import numpy as np
import itertools
import gym
import pickle
import tempfile
import pathlib
from tqdm import tqdm
from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent, NRasterizedIncrementingAgent
ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
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)
self.observation_space = gym.spaces.Dict({
'obs': env.observation_space,
'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
})
def _after_choice(self):
pass
def _after_step(self):
pass
def _transitions(self):
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
if self.done:
self.s = self.env.reset()
self.done = False
self.episode_start = True
self.m = available_actions(self.env)
self.ch, self.value, self.log_prob = generator.policy.predict({
'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
})
self.plan = list(map(float, generate_plan(self.env, self.ch)))
self._after_choice()
assert not self.done
assert self.plan
#assert feasible(self.env, self.plan, self.ch)
while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
self.a, self.plan = self.plan[0], self.plan[1:]
self.a = self.env._normalize(self.a)
self.nexts, _, self.done, _ = self.env.step(self.a)
self._after_step()
self.s = self.nexts
yield from self._transitions()
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,
'acts': np.array((self.a,)),
'dones': np.array(self.done),
})
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):
"""Sample high-level (state, action, reward) tuples for generator training."""
def __init__(self, *args, **kwargs):
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()),
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
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,
'reward': self.r.detach(),
'episode_start': self.episode_start,
'value': self.value.detach(),
'log_prob': self.log_prob.detach(),
'done': self.done,
}
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):
"""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`.
Args:
env (gym.Env): current environment state
i (int): high-level action `i`
Returns:
plan (np.array): length T array of acceleration values
"""
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 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}
#env_class = NRasterized
#env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
files = ['../../../expert_data/DR_USA_Roundabout_FT/track%04i/expert.pkl'%(i) for i in range(5)]
transitions=load_experts(files)
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(NRasterizedRandomAgent(**env_args))
for s in env.sample_ll(model):
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

View File

@@ -1,91 +0,0 @@
# %%
import sys
sys.path.append('../../../')
from src.discriminator import CnnDiscriminatorFlatAction
from imitation.algorithms import adversarial
import stable_baselines3
import torch.utils.data
import numpy as np
from intersim.envs.intersimple import NRasterizedRouteRandomAgent
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
from src.policies.options import OptionsCnnPolicy
from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
from src.gail.train import train_discriminator, train_generator
model_name = 'gail_options_image_random'
env_settings = {'width': 70, 'height': 70, 'm_per_px': 1}
ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
def train(expert_data, epochs=100, expert_batch_size=64, generator_steps=1024, discount=0.99):
env = NRasterizedRouteRandomAgent(**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(NRasterizedRouteRandomAgent, n_envs=1, env_kwargs=env_settings)
discriminator = adversarial.GAIL(
expert_data=expert_data,
expert_batch_size=expert_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, options=ALL_OPTIONS),
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, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
generator.save(model_name)
return generator
def video(model_name, env):
model = stable_baselines3.PPO.load(model_name)
env = RenderOptions(env, options=ALL_OPTIONS)
for s in env.sample_ll(model):
if s['dones']:
break
env.close(filestr='render/'+model_name)
def evaluate():
video(
model_name=model_name,
env=NRasterizedRouteRandomAgent(**env_settings)
)
# %%
if __name__ == '__main__':
with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteRandomAgentw70h70mppx1.pkl", "rb") as f:
trajectories = pickle.load(f)
transitions = rollout.flatten_trajectories(trajectories)
train(transitions)

View File

@@ -1,144 +0,0 @@
# %%
from collections import deque
import sys
sys.path.append('../../../')
from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
from imitation.algorithms import adversarial
import stable_baselines3
import pickle
import imitation.data.rollout as rollout
import tempfile
import pathlib
from imitation.util import logger
from tqdm import tqdm
from src.policies.options import OptionsCnnPolicy
from src.gail.train import flatten_transitions
from gail.options2 import OptionsEnv, RenderOptions, imitation_discriminator
from gail.envs import TLNRasterizedRouteRandomAgentLocation
from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
from stable_baselines3.common.env_util import make_vec_env
import torch
import numpy as np
model_name = 'gail_options_image_random_location'
env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 200}
ALL_OPTIONS = [(v,t) for v in [0,2,4,8,10] for t in [5, 10, 20]] # option 0 is safe fallback
class NoisyDiscriminator(CnnDiscriminatorFlatAction):
def __init__(self, *args, std=0.0, **kwargs):
super().__init__(*args, **kwargs)
self.std = std
def forward(self, state, action):
noise = self.std * torch.randn(*action.shape, device=action.device)
return super().forward(state, action + noise)
class LLBuffer(deque):
def sample(self, n):
assert n <= self.maxlen, f'Sample size of {n} exceeds buffer capacity of {self.maxlen}'
assert n <= len(self), f'Sample size of {n} exceeds buffer size of {len(self)}'
ind = np.random.randint(len(self), size=n)
return list(self[i] for i in ind)
def train(
expert_data,
expert_batch_size=4096,
discriminator_updates_per_round=20,
generator_steps=1024,
generator_batch_size=1024,
generator_total_steps=8192,
generator_updates_per_round=10,
discount=1.0,
epochs=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}/.")
venv = DummyVecEnv([lambda: env])
discriminator = adversarial.GAIL(
expert_data=expert_data,
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
)
ll_buffer = LLBuffer(maxlen=expert_batch_size*10)
options_env = make_vec_env(
OptionsEnv,
n_envs=1,
#vec_env_cls=SubprocVecEnv,
env_kwargs={
'env': env,
'options': ALL_OPTIONS,
'discriminator': imitation_discriminator(discriminator),
'discount': discount,
'll_buffer': ll_buffer,
}
)
generator = stable_baselines3.PPO(
OptionsCnnPolicy,
options_env,
verbose=1,
batch_size=generator_batch_size,
n_steps=generator_steps,
n_epochs=generator_updates_per_round,
gamma=1.0,
learning_rate=1e-4,
)
for _ in tqdm(range(epochs)):
ll_buffer.clear()
# train generator
generator.learn(total_timesteps=generator_total_steps)
# train discriminator
for _ in range(discriminator_updates_per_round):
generator_samples = ll_buffer.sample(expert_batch_size)
generator_samples = flatten_transitions(generator_samples)
discriminator.train_disc(gen_samples=generator_samples)
generator.save(model_name)
return generator
def video(model_name, env):
model = stable_baselines3.PPO.load(model_name)
done = False
obs = env.reset()
while not done:
action, _ = model.predict(obs)
obs, _, done, _ = env.step(action)
env.close(filestr='render/'+model_name)
def evaluate():
video_settings = { **env_settings, 'random_skip': False, 'max_episode_steps': 200 }
env = TLNRasterizedRouteRandomAgentLocation(**video_settings)
env = RenderOptions(env, options=ALL_OPTIONS)
video(
model_name=model_name,
env=env
)
# %%
if __name__ == '__main__':
with open("data/NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
trajectories = pickle.load(f)
transitions = rollout.flatten_trajectories(trajectories)
train(transitions)

View File

@@ -1,63 +0,0 @@
# %%
import pathlib
import pickle
import tempfile
import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env
from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger
from intersim.envs.intersimple import IntersimpleReward
# Load pickled test demonstrations.
with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.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.
trajectories = pickle.load(f)
# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)
venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
# Train BC on expert data.
# BC also accepts as `expert_data` any PyTorch-style DataLoader that iterates over
# dictionaries containing observations and actions.
logger.configure(tempdir_path / "BC/")
bc_trainer = bc.BC(venv.observation_space, venv.action_space, expert_data=transitions)
bc_trainer.train(n_epochs=1)
# 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/")
gail_trainer = adversarial.GAIL(
venv,
expert_data=transitions,
expert_batch_size=32,
gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024),
)
gail_trainer.train(total_timesteps=2048)
# Train AIRL on expert data.
logger.configure(tempdir_path / "AIRL/")
airl_trainer = adversarial.AIRL(
venv,
expert_data=transitions,
expert_batch_size=32,
gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024),
)
airl_trainer.train(total_timesteps=2048)
# %%

View File

@@ -1,36 +0,0 @@
# %%
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from intersim.envs.intersimple import IntersimpleReward, speed_reward
model_name = "ppo_const"
env = IntersimpleReward(
agent=51,
#reward=speed_reward,
)
# %%
model = PPO(
"MlpPolicy", 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)

View File

@@ -1,41 +0,0 @@
# %%
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from intersim.envs.intersimple import ConstCollisionReward, IntersimpleFlatAgent
model_name = "ppo_const_collision"
class IntersimpleConstCollisionAgent(ConstCollisionReward, IntersimpleFlatAgent):
pass
env = IntersimpleConstCollisionAgent(
agent=51,
speed_reward_weight=0.001,
collision_penalty=1000
)
# %%
model = PPO(
"MlpPolicy", env,
learning_rate=3e-6,
verbose=1,
)
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)

View File

@@ -1,35 +0,0 @@
# %%
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from intersim.envs.intersimple import NRasterized
model_name = "ppo_const_image"
env = NRasterized(
agent=51,
)
# %%
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)

View File

@@ -1,33 +0,0 @@
# %%
from stable_baselines3 import PPO
from intersim.envs.intersimple import NRasterizedRandomAgent
import functools
model_name = "ppo_const_image_random"
env = NRasterizedRandomAgent()
# %%
model = PPO(
"CnnPolicy", env,
verbose=1,
)
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)

View File

@@ -1,25 +0,0 @@
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from intersim.envs.intersimple import IntersimpleTargetSpeed
env = IntersimpleTargetSpeed()
model = PPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=25000)
model.save("ppo_intersimple")
print('Done training.')
del model # remove to demonstrate saving and loading
model = PPO.load("ppo_intersimple")
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()

View File

@@ -1,46 +0,0 @@
# %%
from stable_baselines3 import PPO
from intersim.envs.intersimple import IntersimpleReward, speed_reward
import functools
model_name = "ppo_speed"
#def reward(state, action, info):
# speed = state[2].item()
# r = speed if speed < 10 else (10 - 5 * (speed - 10))
# return 0.1 * r
env = IntersimpleReward(
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
# %%
model = PPO(
"MlpPolicy", 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)
# %%

View File

@@ -1,46 +0,0 @@
# %%
from stable_baselines3 import PPO
from intersim.envs.intersimple import NRasterized, speed_reward
import functools
model_name = "ppo_speed_image"
#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=20,
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)
# %%

View File

@@ -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)
# %%

View File

@@ -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)

View File

@@ -1,39 +0,0 @@
# %%
from stable_baselines3 import PPO
from intersim.envs.intersimple import NRasterizedRandomAgent, speed_reward
import functools
model_name = "ppo_speed_image_random"
env = NRasterizedRandomAgent(
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)

View File

@@ -1,50 +0,0 @@
# %%
from stable_baselines3 import PPO
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
model_name = "ppo_speed_lidar"
#def reward(state, action, info):
# speed = state[2].item()
# r = speed if speed < 10 else (10 - 5 * (speed - 10))
# return 0.1 * r
env = IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
# %%
model = PPO(
"MlpPolicy", env,
learning_rate=1e-4,
verbose=1,
tensorboard_log='runs/'
)
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)
# %%

View File

@@ -1,49 +0,0 @@
# %%
from stable_baselines3 import PPO
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
import functools
model_name = "ppo_speed_lidar_random"
#def reward(state, action, info):
# speed = state[2].item()
# r = speed if speed < 10 else (10 - 5 * (speed - 10))
# return 0.1 * r
env = IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
# %%
model = PPO(
"MlpPolicy", env,
learning_rate=1e-4,
verbose=1,
tensorboard_log='runs/'
)
model.learn(total_timesteps=1000000)
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)
# %%

View File

@@ -1,43 +0,0 @@
# %%
from stable_baselines3 import PPO
from intersim.envs.intersimple import IntersimpleFlatRandomAgent, Reward, RewardVisualization, speed_reward
import functools
model_name = "ppo_speed_random"
class IntersimpleRewardRandom(RewardVisualization, Reward, IntersimpleFlatRandomAgent):
"""`IntersimpleFlatAgent` with rewards."""
pass
env = IntersimpleRewardRandom(
reward=functools.partial(
speed_reward,
collision_penalty=0
)
)
# %%
model = PPO(
"MlpPolicy", 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)

View File

@@ -1,39 +0,0 @@
# %%
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from intersim.envs.intersimple import IntersimpleTargetSpeedAgent
model_name = "ppo_tspeed"
env = IntersimpleTargetSpeedAgent(
agent=51,
target_speed=10,
speed_penalty_weight=0.001,
collision_penalty=1000
)
# %%
model = PPO(
"MlpPolicy", env,
learning_rate=3e-6,
verbose=1,
)
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)

View File

@@ -1,31 +0,0 @@
# %%
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from intersim.envs.intersimple import IntersimpleTargetSpeedRandom
model_name = "ppo_tspeed_random"
# %%
env = IntersimpleTargetSpeedRandom(target_speed=10)
# %%
model = PPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=250000)
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()

View File

@@ -1,33 +0,0 @@
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))
if __name__ == '__main__':
import fire
fire.Fire(render_env)

View File

@@ -1,71 +0,0 @@
# %%
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')

View File

@@ -1,52 +0,0 @@
# %%
from sb3_contrib import TRPO
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
model_name = "trpo_speed_lidar"
#def reward(state, action, info):
# speed = state[2].item()
# r = speed if speed < 10 else (10 - 5 * (speed - 10))
# return 0.1 * r
env = IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
# %%
model = TRPO(
"MlpPolicy", env,
learning_rate=1e-4,
verbose=1,
tensorboard_log='runs/',
#use_sde=True,
#sde_sample_freq=4,
)
model.learn(total_timesteps=1000000)
model.save(model_name)
print('Done training.')
del model # remove to demonstrate saving and loading
# %%
model = TRPO.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)
# %%

File diff suppressed because one or more lines are too long

View File

@@ -1,249 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 2,
"source": [
"%cd learners"
],
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"/home/buehrle/dev/InteractionImitation/scratch/etienne/pillbox/learners\n"
]
}
],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": 3,
"source": [
"%load_ext autoreload\n",
"%autoreload 2"
],
"outputs": [],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": 1,
"source": [
"import torch\n",
"import numpy as np"
],
"outputs": [],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": 9,
"source": [
"# save expert demos to ../experts/Intersim/demos.npz\n",
"# make sure to split different experts up\n",
"\n",
"from intersim.envs.simulator import InteractionSimulator\n",
"from intersim.utils import get_map_path, get_svt, SVT_to_stateactions\n",
"import gym\n",
"from tqdm import tqdm\n",
"\n",
"def pillbox_demo(observations, actions, rewards):\n",
" demo = {\n",
" 'env': 'intersim:intersim-v0',\n",
" 'num_trajs': len(observations),\n",
" 'mean_reward': rewards.mean(),\n",
" 'std_reward': rewards.std(),\n",
" }\n",
" demo.update({\n",
" str(i): {\n",
" 'states': o,\n",
" 'actions': a,\n",
" } for i, (o, a) in enumerate(zip(observations, actions))\n",
" })\n",
" return demo\n",
"\n",
"def intersim_expert_demos(loc, track):\n",
" svt, svt_path = get_svt(loc, track)\n",
" osm = get_map_path(loc)\n",
" \n",
" n_actors = svt.simstate.size(1)\n",
" observations = []\n",
" actions = [] ##\n",
" #states, actions = SVT_to_stateactions(svt) ##\n",
" rewards = []\n",
" \n",
" print('Simulating')\n",
" env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm)\n",
" obs, info = env.reset()\n",
" for s in tqdm(svt.simstate[1:]): ##\n",
" #for a in actions: ##\n",
" relative_state = torch.stack((\n",
" obs['relative_state'][..., 0],\n",
" obs['relative_state'][..., 1],\n",
" (obs['relative_state'][..., 2]**2 + obs['relative_state'][..., 3]**2).sqrt(),\n",
" obs['relative_state'][..., 4],\n",
" obs['relative_state'][..., 5],\n",
" ), -1)\n",
" observations.append(torch.cat((\n",
" obs['state'].unsqueeze(1),\n",
" relative_state,\n",
" ), 1))\n",
" obs, r, done, info = env.step(env.target_state(s, mu=.01))\n",
" #obs, r, done, info = env.step(a) ##\n",
" actions.append(info['action_taken'])\n",
" rewards.append(r)\n",
" assert not done, 'Episode terminated during expert demonstration.'\n",
"\n",
" _except_idx = lambda o, i: torch.cat((o[:i], o[i+1:]))\n",
" \n",
" # transpose to per-agent observations and actions\n",
" print('Transposing')\n",
" observations = [torch.stack([_except_idx(o[i], i+1) for o in observations]) for i in range(n_actors)]\n",
" actions = [torch.stack([a[i] for a in actions]) for i in range(n_actors)]\n",
" \n",
" print('Trimming')\n",
" # trim observations and actions to start/end of trajectory\n",
" _alive = lambda o: (~o.isnan().all(2).all(1)).nonzero()\n",
" _start = lambda o: _alive(o).min()\n",
" _end = lambda o: _alive(o).max() + 1\n",
" start_end = [(_start(obs), _end(obs)) for obs in observations]\n",
" observations = [obs[start:end] for obs, (start, end) in zip(observations, start_end)]\n",
" actions = [act[start:end] for act, (start, end) in zip(actions, start_end)]\n",
" \n",
" #print('Cropping')\n",
" ## crop observations to max number of observations\n",
" #max_num_obs = max([(~obs.isnan().all(2)).sum(1).max() for obs in observations])\n",
" #observations = [obs[:, :max_num_obs] for obs in observations]\n",
" \n",
" observations = [o.numpy() for o in observations]\n",
" actions = [a.numpy() for a in actions]\n",
" rewards = np.array(rewards)\n",
" \n",
" return pillbox_demo(observations, actions, rewards)"
],
"outputs": [],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": 10,
"source": [
"demos = intersim_expert_demos(loc=0, track=0)"
],
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Simulating\n",
"Custom Vehicle Trajectory Paths\n",
"Map Path: datasets/maps/DR_USA_Roundabout_FT.osm\n",
"Environment Reset\n"
]
},
{
"output_type": "stream",
"name": "stderr",
"text": [
"100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 3006/3006 [01:17<00:00, 38.87it/s]\n"
]
},
{
"output_type": "stream",
"name": "stdout",
"text": [
"Transposing\n",
"Trimming\n"
]
}
],
"metadata": {
"scrolled": true,
"tags": [
"outputPrepend"
]
}
},
{
"cell_type": "code",
"execution_count": 6,
"source": [
"demos['num_trajs']"
],
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"151"
]
},
"metadata": {},
"execution_count": 6
}
],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": 7,
"source": [
"demos['25']['states'].shape"
],
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"(71, 151, 5)"
]
},
"metadata": {},
"execution_count": 7
}
],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": 8,
"source": [
"np.savez('../experts/intersim:intersim-v0/demos.npz', **demos)"
],
"outputs": [],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": null,
"source": [],
"outputs": [],
"metadata": {}
}
],
"metadata": {
"kernelspec": {
"name": "python3",
"display_name": "Python 3.7.5 64-bit ('.venv': venv)"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.5"
},
"interpreter": {
"hash": "56465d2ea10f338edb3d30adb010c5849fd826fffc543ba31360f3db8b47a703"
}
},
"nbformat": 4,
"nbformat_minor": 4
}

View File

@@ -1,168 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 2,
"source": [
"%cd learners"
],
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"[Errno 2] No such file or directory: 'learners'\n",
"/home/buehrle/dev/InteractionImitation/scratch/etienne/pillbox/learners\n"
]
}
],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": null,
"source": [
"import gym\n",
"from tqdm import tqdm\n",
"\n",
"def rollout(pi, max_steps=1000):\n",
" env = gym.make('intersim:intersim-v0')\n",
" env.reset() # obs = env.reset()\n",
" obs, _, done, _ = env.step(0 * env.action_space.sample())\n",
" \n",
" _except = lambda o, i: torch.cat((o[:i], o[i+1:]))\n",
" \n",
" _relative_state_v = lambda obs: torch.stack((\n",
" obs[..., 0],\n",
" obs[..., 1],\n",
" (obs[..., 2]**2 + obs[..., 3]**2).sqrt(),\n",
" obs[..., 4],\n",
" obs[..., 5],\n",
" ), -1)\n",
" \n",
" for _ in tqdm(range(max_steps)):\n",
" pi_obs = [\n",
" torch.cat((e.unsqueeze(0), _relative_state_v(_except(o, i)))).unsqueeze(0)\n",
" for i, (e, o) in enumerate(zip(obs['state'], obs['relative_state']))\n",
" ]\n",
" \n",
" actions = [pi(o).squeeze() for o in pi_obs]\n",
" actions = torch.stack(actions).unsqueeze(1)\n",
" obs, _, done, _ = env.step(actions)\n",
" env.render(mode='post')\n",
" if done:\n",
" break\n",
" env.close()"
],
"outputs": [],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": 4,
"source": [
"import torch\n",
"import numpy as np\n",
"\n",
"def expert(obs):\n",
" ego = obs[:, 0]\n",
" rel = obs[:, 1:]\n",
" front = torch.stack((torch.cos(ego[:, 3]), torch.sin(ego[:, 3])), -1)\n",
" left = torch.stack((-torch.sin(ego[:, 3]), torch.cos(ego[:, 3])), -1)\n",
" df = (rel[:, :, :2] * front.unsqueeze(1)).sum(-1)\n",
" dl = (rel[:, :, :2] * left.unsqueeze(1)).sum(-1)\n",
"\n",
" df = torch.where(df.isnan(), np.inf * torch.ones_like(df), df)\n",
" dl = torch.where(dl.isnan(), np.inf * torch.ones_like(dl), dl)\n",
" rel = torch.where(rel.isnan(), np.inf * torch.ones_like(rel), rel)\n",
"\n",
" # relative speed in direction of position difference vector\n",
" vrel = rel[:, :, 2] * (rel[:, :, :2] * torch.stack((\n",
" torch.cos(ego[:, 3].unsqueeze(1) + rel[:, :, 3]),\n",
" torch.sin(ego[:, 3].unsqueeze(1) + rel[:, :, 3])),\n",
" -1)).sum(-1)\n",
" vrel = torch.where(vrel.isnan(), np.inf * torch.ones_like(vrel), vrel)\n",
" vrel = torch.maximum(vrel, torch.zeros_like(vrel))\n",
" \n",
" alpha = torch.atan2(dl, df)\n",
" d = (rel[:, :, :2] ** 2).sum(-1)\n",
" attn = torch.exp(-torch.where(alpha > 0, 0.8*alpha, 1*alpha)**2 - 0.01 * d - 0.1*vrel) \n",
" \n",
" act = 10 - ego[:, 2] - 20 * attn.sum(-1)\n",
" \n",
" return act"
],
"outputs": [],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": 5,
"source": [
"rollout(expert, max_steps=500)"
],
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Vehicle Trajectory Paths: /home/buehrle/dev/InteractionImitation/InteractionSimulator/datasets/trackfiles/DR_USA_Roundabout_FT/vehicle_tracks_000.csv\n",
"Map Path: /home/buehrle/dev/InteractionImitation/InteractionSimulator/datasets/maps/DR_USA_Roundabout_FT.osm\n",
"Environment Reset\n"
]
},
{
"output_type": "stream",
"name": "stderr",
"text": [
" 0%| | 0/500 [00:00<?, ?it/s]\n"
]
},
{
"output_type": "error",
"ename": "RuntimeError",
"evalue": "torch.cat(): Sizes of tensors must match except in dimension 0. Got 5 and 6 in dimension 1 (The offending index is 1)",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)",
"\u001b[0;32m/tmp/ipykernel_4266/633354333.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mrollout\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexpert\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m500\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
"\u001b[0;32m/tmp/ipykernel_4266/1839273282.py\u001b[0m in \u001b[0;36mrollout\u001b[0;34m(pi, max_steps)\u001b[0m\n\u001b[1;32m 12\u001b[0m pi_obs = [\n\u001b[1;32m 13\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_except_self\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mo\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mo\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'relative_state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 15\u001b[0m ]\n\u001b[1;32m 16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/tmp/ipykernel_4266/1839273282.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 12\u001b[0m pi_obs = [\n\u001b[1;32m 13\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_except_self\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mo\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mo\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'relative_state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 15\u001b[0m ]\n\u001b[1;32m 16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mRuntimeError\u001b[0m: torch.cat(): Sizes of tensors must match except in dimension 0. Got 5 and 6 in dimension 1 (The offending index is 1)"
]
}
],
"metadata": {}
},
{
"cell_type": "code",
"execution_count": null,
"source": [],
"outputs": [],
"metadata": {}
}
],
"metadata": {
"kernelspec": {
"name": "python3",
"display_name": "Python 3.7.5 64-bit ('.venv': venv)"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.5"
},
"interpreter": {
"hash": "56465d2ea10f338edb3d30adb010c5849fd826fffc543ba31360f3db8b47a703"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

File diff suppressed because one or more lines are too long

View File

@@ -1,140 +0,0 @@
import gym
from gym import spaces
from sklearn.neighbors import KDTree
from scipy.stats import norm
import numpy as np
import warnings
from abc import ABC, abstractmethod
from typing import Dict, Generator, Optional, Union
import torch as th
try:
# Check memory used by replay buffer when possible
import psutil
except ImportError:
psutil = None
from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape
from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBufferSamples
from stable_baselines3.common.vec_env import VecNormalize
from stable_baselines3.common.buffers import ReplayBuffer
class AdRILWrapper(gym.Env):
metadata = {'render.modes': ['human']}
def __init__(self, base_env):
super(AdRILWrapper, self).__init__()
self.base_env = base_env
self.iter = 0
self.observation_space = self.base_env.observation_space
self.action_space = self.base_env.action_space
self.trajs = list()
self.num_trajs = 0
self.curr_state = None
def step(self, action):
next_obs, _, done, info = self.base_env.step(action)
reward = self.iter # Transformed by replay buffer
self.trajs.append((self.curr_state, action, next_obs, done))
if done:
self.num_trajs += 1
self.curr_state = next_obs
return next_obs, reward, done, info
def reset(self):
obs = self.base_env.reset()
self.curr_state = obs
return obs
def render(self, mode='human'):
self.base_env.render(mode=mode)
def close (self):
self.base_env.close()
def get_learner_trajs(self):
return self.trajs
def set_iter(self, k):
self.iter = k
class AdRILReplayBuffer(ReplayBuffer):
def __init__(
self,
buffer_size: int,
observation_space: spaces.Space,
action_space: spaces.Space,
device: Union[th.device, str] = "cpu",
n_envs: int = 1,
optimize_memory_usage: bool = False,
expert_data: dict = dict(),
N_expert: int = 0,
balanced: bool = True,
):
super(AdRILReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs, optimize_memory_usage=optimize_memory_usage)
self.expert_states = expert_data['obs']
self.expert_actions = expert_data['acts']
self.expert_next_states = expert_data['next_obs']
self.expert_dones = expert_data['dones']
n_expert = len(expert_data["obs"])
self.iter = 0
self.N_expert = N_expert
self.N_learner = 0
self.normalizer = 1
self.balanced = balanced
def set_iter(self, k):
self.iter = k
normalizer = 0
for i in range(0, k):
normalizer += 1 ** (-i) # written to support decaying learning rate
self.normalizer = normalizer
def set_n_learner(self, n):
self.N_learner = n
def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
num_samples = len(batch_inds)
if self.balanced:
num_expert_samples = int(num_samples / 2)
batch_inds = batch_inds[:num_expert_samples]
expert_inds = np.random.randint(0, len(self.expert_states), size=num_expert_samples)
# balanced sampling
if self.optimize_memory_usage:
next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
else:
next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
next_obs = np.concatenate((next_obs, self._normalize_obs(self.expert_next_states[expert_inds], env)), axis=0)
obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
obs = np.concatenate((obs, self._normalize_obs(self.expert_states[expert_inds], env)), axis=0)
actions = self.actions[batch_inds, 0, :]
actions = np.concatenate((actions, self.expert_actions[expert_inds].reshape(num_expert_samples, -1)), axis=0)
dones = self.dones[batch_inds]
dones = np.concatenate((dones, self.expert_dones[expert_inds].reshape(num_expert_samples, -1)), axis=0)
# AdRIL Rewards (indicator kernel)
mask1 = (self.rewards[batch_inds] >= 0).astype(np.float32)
mask2 = (self.rewards[batch_inds] < self.iter).astype(np.float32)
r1 = - (1. ** (-self.rewards[batch_inds])) * mask1 * mask2 # Past iter
r2 = np.zeros_like(self.rewards[batch_inds]) * mask1 * (1 - mask2) # current iter
r3 = -self.rewards[batch_inds] * (1 - mask1) # Expert
if self.iter > 0:
rewards = (r1 / self.N_learner) + r2 + r3
else:
rewards = r1 + r2 + r3
rewards = np.concatenate((rewards, np.ones_like(rewards) / self.N_expert), axis=0)
else:
if self.optimize_memory_usage:
next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
else:
next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
actions = self.actions[batch_inds, 0, :]
dones = self.dones[batch_inds]
# AdRIL Rewards (indicator kernel)
mask1 = (self.rewards[batch_inds] >= 0).astype(np.float32)
mask2 = (self.rewards[batch_inds] < self.iter).astype(np.float32)
r1 = - (1. ** (-self.rewards[batch_inds])) * mask1 * mask2 # Past iter
r2 = np.zeros_like(self.rewards[batch_inds]) * mask1 * (1 - mask2) # current iter
r3 = -self.rewards[batch_inds] * (1 - mask1) / self.N_expert # Expert
if self.iter > 0:
rewards = (r1 * 1. / self.N_learner) + r2 + r3
else:
rewards = r1 + r2 + r3
data = (obs, actions, next_obs, dones, rewards)
return ReplayBufferSamples(*tuple(map(self.to_torch, data)))

View File

@@ -1,222 +0,0 @@
import numpy as np
import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from gym.spaces import Discrete
import gym
from stable_baselines3.common.preprocessing import get_action_dim
from tqdm import tqdm
from torch.autograd import Variable
from itertools import repeat
from torch.autograd import grad as torch_grad
from typing import List, Type
import types
# Infinite dataloader
def repeater(data_loader):
for loader in repeat(data_loader):
for data in loader:
yield data
def create_mlp(
input_dim: int, output_dim: int, net_arch: List[int], activation_fn: Type[nn.Module] = nn.ReLU) -> List[nn.Module]:
if len(net_arch) > 0:
modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
else:
modules = []
for idx in range(len(net_arch) - 1):
modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1]))
modules.append(activation_fn())
if output_dim > 0:
last_layer_dim = net_arch[-1] if len(net_arch) > 0 else input_dim
modules.append(nn.Linear(last_layer_dim, output_dim))
return modules
def init_ortho(layer):
if type(layer) == nn.Linear:
nn.init.orthogonal_(layer.weight)
class AdVILPolicy(nn.Module):
def __init__(self, env, mean=None, std=None):
super(AdVILPolicy, self).__init__()
if isinstance(env.action_space, Discrete):
self.net_arch = [64, 64]
self.action_dim = env.action_space.n
self.discrete = True
else:
self.net_arch = [256, 256]
self.action_dim = int(np.prod(env.action_space.shape))
self.low = torch.as_tensor(env.action_space.low)
self.high = torch.as_tensor(env.action_space.high)
self.discrete = False
self.obs_dim = int(np.prod(env.observation_space.shape))
self.observation_space = env.observation_space
net = create_mlp(self.obs_dim, self.action_dim, self.net_arch, nn.ReLU)
if self.discrete:
net.append(nn.Softmax(dim=1))
self.net = nn.Sequential(*net)
self.net.apply(init_ortho)
if mean is not None and std is not None:
self.mean = mean
self.std = std
self.is_normalized = True
else:
self.is_normalized = False
def forward(self, obs):
action = self.net(obs)
return action
def predict(self, obs, state, mask, deterministic):
obs = obs.reshape((-1,) + (self.obs_dim,))
if self.is_normalized:
obs = (obs - self.mean) / self.std
obs = torch.as_tensor(obs)
with torch.no_grad():
actions = self.forward(obs)
if self.discrete:
actions = actions.argmax(dim=1).reshape(-1)
else:
actions = self.low + ((actions + 1.0) / 2.0) * (self.high - self.low)
actions = torch.max(torch.min(actions, self.high), self.low)
actions = actions.cpu().numpy()
return actions, state
class AdVILDiscriminator(nn.Module):
def __init__(self, env):
super(AdVILDiscriminator, self).__init__()
if isinstance(env.action_space, Discrete):
self.net_arch = [64, 64]
self.action_dim = env.action_space.n
else:
self.net_arch = [256, 256]
self.action_dim = int(np.prod(env.action_space.shape))
self.obs_dim = int(np.prod(env.observation_space.shape))
net = create_mlp(self.obs_dim + self.action_dim, 1, self.net_arch, nn.ReLU)
self.net = nn.Sequential(*net)
self.net.apply(init_ortho)
def forward(self, inputs):
output = self.net(inputs)
return output.view(-1)
def pi_update(obs, acts, pi, f, pi_opt, prog):
pi_opt.zero_grad()
obs_v = Variable(obs)
pi_acts = pi(obs_v)
#learner_sa = torch.cat((obs, pi_acts), axis=1)
f_learner = f(obs, acts)
pi_loss = f_learner.mean() + orthogonal_reg(pi) + 2e-1 * (pi_acts - acts).square().mean()
pi_loss.backward()
if prog > 0.1:
torch.nn.utils.clip_grad_norm(pi.parameters(), 40.0)
pi_opt.step()
return pi_loss.item(), (2e-1 * (pi_acts - acts).square().mean()).item()
def orthogonal_reg(pi):
with torch.enable_grad():
reg = 1e-4
orth_loss = torch.zeros(1)
for name, param in pi.named_parameters():
if 'bias' not in name:
x = torch.mm(torch.t(param), param)
x = x * (1. - torch.eye(param.shape[-1]))
orth_loss = orth_loss + reg * (x.square().sum())
return orth_loss
def f_update(obs, acts, pi, f, f_opt, prog):
obs_v = Variable(obs)
pi_acts = pi(obs_v)
#learner_sa = torch.cat((obs, pi_acts), axis=1)
#expert_sa = Variable(torch.cat((obs, acts), axis=1))
f_learner = f(obs, pi_acts)
f_expert = f(obs, acts)
#gp = gradient_penalty((obs, pi_acts), (obs, acts), f)
f_opt.zero_grad()
f_loss = f_expert.mean() - f_learner.mean()# + 10 * gp
f_loss.backward()
if prog > 0.1:
torch.nn.utils.clip_grad_norm(f.parameters(), 40.0)
f_opt.step()
return f_loss.item()
def gradient_penalty(learner_sa, expert_sa, f):
batch_size = expert_sa[0].size()[0]
#alpha = torch.rand(batch_size, 1)
#alpha = alpha.expand_as(expert_sa)
salpha = torch.rand(batch_size, 1, 1)
salpha = salpha.expand_as(expert_sa[0])
aalpha = torch.rand(batch_size, 1)
aalpha = aalpha.expand_as(expert_sa[1])
#interpolated = alpha * expert_sa.data + (1 - alpha) * learner_sa.data
#interpolated = Variable(interpolated, requires_grad=True)
#f_interpolated = f(interpolated.float())
sinterpolated = salpha * expert_sa[0].data + (1 - salpha) * learner_sa[0].data
sinterpolated = Variable(sinterpolated, requires_grad=True)
ainterpolated = aalpha * expert_sa[1].data + (1 - aalpha) * learner_sa[1].data
ainterpolated = Variable(ainterpolated, requires_grad=True)
f_interpolated = f(sinterpolated, ainterpolated)
#gradients = torch_grad(outputs=f_interpolated, inputs=interpolated,
# grad_outputs=torch.ones(f_interpolated.size()),
# create_graph=True, retain_graph=True)[0]
sgradients = torch_grad(outputs=f_interpolated, inputs=sinterpolated,
grad_outputs=torch.ones(f_interpolated.size()),
create_graph=True, retain_graph=True)[0]
agradients = torch_grad(outputs=f_interpolated, inputs=ainterpolated,
grad_outputs=torch.ones(f_interpolated.size()),
create_graph=True, retain_graph=True)[0]
#gradients = gradients.view(batch_size, -1)
sgradients = sgradients.view(batch_size, -1)
agradients = agradients.view(batch_size, -1)
#norm = gradients.norm(2, dim=1).mean().item()
#gradients_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + 1e-12)
gradients_norm = torch.sqrt(torch.sum(sgradients ** 2, dim=1) + torch.sum(agradients ** 2, dim=1) + 1e-12)
# 2 * |f'(x_0)|
return ((gradients_norm - 0.4) ** 2).mean()
def advil_training(data_loader, env, iters=int(1e5), policy_class=AdVILPolicy, discriminator_class=AdVILDiscriminator, lr_pi=8e-6, lr_f=8e-4):
if not isinstance(env.action_space, Discrete):
low = torch.as_tensor(env.action_space.low)
high = torch.as_tensor(env.action_space.high)
if data_loader.dataset.is_normalized:
pi = policy_class(env, data_loader.dataset.mean, data_loader.dataset.std)
else:
pi = policy_class(env)
f = discriminator_class(env)
pi_opt = optim.Adam(pi.parameters(), lr=lr_pi)
last_loss = 0
f_opt = optim.Adam(f.parameters(), lr=lr_f)
data_loader = repeater(data_loader)
for t in tqdm(range(iters)):
data = next(data_loader)
obs = data['obs']
acts = data['acts']
#if isinstance(env.action_space, Discrete):
# acts = nn.functional.one_hot(acts, env.action_space.n)
#else:
# acts = (((acts - low) / (high - low)) * 2.0) - 1.0
pi_loss, mse_reg = pi_update(obs, acts, pi, f, pi_opt, t/iters)
f_loss = f_update(obs, acts, pi, f, f_opt, t/iters)
if t % 100 == 0:
print("pi loss:", pi_loss)
print("mse reg:", mse_reg)
print("f loss:", f_loss)
return pi

View File

@@ -1,155 +0,0 @@
import torch
import torch.nn as nn
def unnormalize(val, mean, std):
val *= std or 1
val += mean or 0
return val
def normalize(val, mean, std):
val -= mean or 0
val /= std or 1
return val
class IntersimPolicy(nn.Module):
def __init__(self, env, mean=None, std=None):
# assert "intersim" in env.unwrapped.spec.id
super().__init__()
self._ego_encoder = nn.Sequential(
# in 5, out 5
nn.Linear(5, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 5),
nn.ReLU(),
)
self._state_encoder = nn.Sequential(
# in 5, out 5
nn.Linear(5, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 5),
nn.ReLU(),
)
self._deepset = lambda e: e.sum(-2)
self._action_decoder = nn.Sequential(
# in 5 + 5, out 1
nn.Linear(5 + 5, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 1),
)
def forward(self, obs):
# obs.shape = (batch=514, 1 + others=150, 5)
# act.shape = (batch=514, 1)
ego = obs[:, 0]#.detach().clone()
rel = obs[:, 1:]#.detach().clone()
nan = rel.isnan().any(-1, keepdim=True)
rel = torch.where(nan, torch.zeros_like(rel), rel) # required because of https://github.com/pytorch/pytorch/issues/15506
d = (rel[:, :, :2] ** 2).sum(-1).sqrt()
front = torch.stack((torch.cos(ego[:, 3]), torch.sin(ego[:, 3])), -1)
left = torch.stack((-torch.sin(ego[:, 3]), torch.cos(ego[:, 3])), -1)
df = (rel[:, :, :2] * front.unsqueeze(1)).sum(-1)
dl = (rel[:, :, :2] * left.unsqueeze(1)).sum(-1)
alpha = torch.atan2(dl, df)
rel[:, :, 0] = d
rel[:, :, 1] = alpha
e = self._ego_encoder(ego)
x = self._state_encoder(rel)
x = torch.where(nan, torch.zeros_like(x), x)
x = self._deepset(x)
a = self._action_decoder(torch.cat((e, x), 1))
return 10 * a
def predict(self, state, mask, deterministic):
#action_distribution = self.forward(obs)
#action = action_distribution.argmax()
#return action
return self.forward(obs)
class IntersimDiscriminator(nn.Module):
def __init__(self, env):
# assert "intersim" in env.unwrapped.spec.id
super().__init__()
self._ego_encoder = nn.Sequential(
# in 5, out 5
nn.Linear(5, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 5),
nn.ReLU(),
)
self._state_encoder = nn.Sequential(
# in 5, out 5
nn.Linear(5, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 5),
nn.ReLU(),
)
self._deepset = lambda e: e.sum(-2)
self._discriminator = nn.Sequential(
# in 5 + 5 + 1, out 1
nn.Linear(5 + 5 + 1, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 1),
)
def forward(self, obs, acts):
# obs.shape = (batch=514, 1 + others=150, 5)
# acts.shape = (batch=514, 1)
# val.shape = (batch=514,)
ego = obs[:, 0]
rel = obs[:, 1:]
nan = rel.isnan().any(-1, keepdim=True)
rel = torch.where(nan, torch.zeros_like(rel), rel) # required because of https://github.com/pytorch/pytorch/issues/15506
d = (rel[:, :, :2] ** 2).sum(-1).sqrt()
front = torch.stack((torch.cos(ego[:, 3]), torch.sin(ego[:, 3])), -1)
left = torch.stack((-torch.sin(ego[:, 3]), torch.cos(ego[:, 3])), -1)
df = (rel[:, :, :2] * front.unsqueeze(1)).sum(-1)
dl = (rel[:, :, :2] * left.unsqueeze(1)).sum(-1)
alpha = torch.atan2(dl, df)
rel[:, :, 0] = d
rel[:, :, 1] = alpha
e = self._ego_encoder(ego)
x = self._state_encoder(rel)
x = torch.where(nan, torch.zeros_like(x), x)
x = self._deepset(x)
v = self._discriminator(torch.cat((e, x, acts), 1))
return v.squeeze(1)

View File

@@ -1,31 +0,0 @@
from typing import Any, Dict, List, Optional, Type
import gym
import torch as th
from torch import nn
from stable_baselines3.common.policies import BasePolicy, register_policy
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, NatureCNN, create_mlp
from stable_baselines3.dqn.policies import DQNPolicy, QNetwork
class SoftQNetwork(QNetwork):
def _predict(self, observation: th.Tensor, deterministic: bool = True) -> th.Tensor:
q_values = self.forward(observation)
probs = nn.functional.softmax(q_values * 10, dim=1)
m = th.distributions.Categorical(probs)
action = m.sample().reshape(-1)
return action
class SQLPolicy(DQNPolicy):
def make_q_net(self) -> SoftQNetwork:
# Make sure we always have separate networks for features extractors etc
net_args = self._update_features_extractor(
self.net_args, features_extractor=None)
return SoftQNetwork(**net_args).to(self.device)
SoftMlpPolicy = SQLPolicy
register_policy("SoftMlpPolicy", SoftMlpPolicy)

View File

@@ -1,61 +0,0 @@
import warnings
from abc import ABC, abstractmethod
from typing import Dict, Generator, Optional, Union
import numpy as np
import torch as th
from gym import spaces
try:
# Check memory used by replay buffer when possible
import psutil
except ImportError:
psutil = None
from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape
from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBufferSamples
from stable_baselines3.common.vec_env import VecNormalize
from stable_baselines3.common.buffers import ReplayBuffer
class SQILReplayBuffer(ReplayBuffer):
def __init__(
self,
buffer_size: int,
observation_space: spaces.Space,
action_space: spaces.Space,
device: Union[th.device, str] = "cpu",
n_envs: int = 1,
optimize_memory_usage: bool = False,
expert_data: dict = dict(),
):
super(SQILReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs, optimize_memory_usage=optimize_memory_usage)
self.expert_states = expert_data['obs']
self.expert_actions = expert_data['acts']
self.expert_next_states = expert_data['next_obs']
self.expert_dones = expert_data['dones']
def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
num_samples = len(batch_inds)
num_expert_samples = int(num_samples / 2)
batch_inds = batch_inds[:num_expert_samples]
expert_inds = np.random.randint(0, len(self.expert_states), size=num_expert_samples)
# Balanced sampling
if self.optimize_memory_usage:
next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
else:
next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
next_obs = np.concatenate((next_obs, self._normalize_obs(self.expert_next_states[expert_inds], env)), axis=0)
obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
obs = np.concatenate((obs, self._normalize_obs(self.expert_states[expert_inds], env)), axis=0)
actions = self.actions[batch_inds, 0, :]
actions = np.concatenate((actions, self.expert_actions[expert_inds].reshape(num_expert_samples, -1)), axis=0)
dones = self.dones[batch_inds]
dones = np.concatenate((dones, self.expert_dones[expert_inds].reshape(num_expert_samples, -1)), axis=0)
# SQIL Rewards
rewards = self.rewards[batch_inds] * 0.
rewards = np.concatenate((rewards, np.ones_like(rewards)), axis=0)
data = (obs, actions, next_obs, dones, rewards)
return ReplayBufferSamples(*tuple(map(self.to_torch, data)))

View File

@@ -1,248 +0,0 @@
from imitation.algorithms import adversarial, bc
from imitation.util import logger, util
from stable_baselines3 import PPO, DQN, SAC
from soft_q import SQLPolicy
from sqil import SQILReplayBuffer
from stable_baselines3.common import policies
from stable_baselines3.common.evaluation import evaluate_policy
from imitation.rewards import discrim_nets
import numpy as np
import argparse
from utils import make_sa_dataloader, make_sads_dataloader, make_sa_dataset, linear_schedule
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
from adril import AdRILWrapper, AdRILReplayBuffer
import os
from gym.spaces import Discrete
import gym
from advil import advil_training
from stable_baselines3.common.running_mean_std import RunningMeanStd
from advil import AdVILPolicy, AdVILDiscriminator
def train_bc(env, n=0):
venv = util.make_vec_env(env, n_envs=8)
if isinstance(venv.action_space, Discrete):
w = 64
else:
w = 256
for i in range(n):
mean_rewards = []
std_rewards = []
for num_trajs in range(0, 26, 5):
if num_trajs == 0:
expert_data = make_sa_dataloader(env, normalize=False)
else:
expert_data = make_sa_dataloader(env, max_trajs=num_trajs, normalize=False)
bc_trainer = bc.BC(venv.observation_space, venv.action_space, expert_data=expert_data,
policy_class=policies.ActorCriticPolicy,
ent_weight=0., l2_weight=0., policy_kwargs=dict(net_arch=[w, w]))
if num_trajs > 0:
bc_trainer.train(n_batches=int(5e5))
def get_policy(*args, **kwargs):
return bc_trainer.policy
model = PPO(get_policy, env, verbose=1)
model.save(os.path.join("learners", env,
"bc_{0}_{1}".format(i, num_trajs)))
mean_reward, std_reward = evaluate_policy(
model, model.get_env(), n_eval_episodes=10)
mean_rewards.append(mean_reward)
std_rewards.append(std_reward)
print("{0} Trajs: {1}".format(num_trajs, mean_reward))
np.savez(os.path.join("learners", env, "bc_rewards_{0}".format(
i)), means=mean_rewards, stds=std_rewards)
def train_gail(env, n=0):
venv = util.make_vec_env(env, n_envs=8)
if isinstance(venv.action_space, Discrete):
w = 64
else:
w = 256
expert_data = make_sads_dataloader(env, max_trajs=5)
logger.configure(os.path.join("learners", "GAIL"))
for i in range(n):
discrim_net = discrim_nets.ActObsMLP(
action_space=venv.action_space,
observation_space=venv.observation_space,
hid_sizes=(w, w),
)
gail_trainer = adversarial.GAIL(venv, expert_data=expert_data, expert_batch_size=32,
gen_algo=PPO("MlpPolicy", venv, verbose=1, n_steps=1024,
policy_kwargs=dict(net_arch=[w, w])),
discrim_kwargs={'discrim_net': discrim_net})
mean_rewards = []
std_rewards = []
for train_steps in range(20):
if train_steps > 0:
if 'Bullet' in env:
gail_trainer.train(total_timesteps=25000)
else:
gail_trainer.train(total_timesteps=16384)
def get_policy(*args, **kwargs):
return gail_trainer.gen_algo.policy
model = PPO(get_policy, env, verbose=1)
mean_reward, std_reward = evaluate_policy(
model, model.env, n_eval_episodes=10)
mean_rewards.append(mean_reward)
std_rewards.append(std_reward)
print("{0} Steps: {1}".format(train_steps, mean_reward))
np.savez(os.path.join("learners", env, "gail_rewards_{0}".format(i)),
means=mean_rewards, stds=std_rewards)
def train_sqil(env, n=0):
venv = gym.make(env)
expert_data = make_sa_dataset(env, max_trajs=5)
for i in range(n):
if isinstance(venv.action_space, Discrete):
model = DQN(SQLPolicy, venv, verbose=1, policy_kwargs=dict(net_arch=[64, 64]), learning_starts=1)
else:
model = SAC('MlpPolicy', venv, verbose=1, policy_kwargs=dict(net_arch=[256, 256]), ent_coef='auto',
learning_rate=linear_schedule(7.3e-4), train_freq=64, gradient_steps=64, gamma=0.98, tau=0.02)
model.replay_buffer = SQILReplayBuffer(model.buffer_size, model.observation_space,
model.action_space, model.device, 1,
model.optimize_memory_usage, expert_data=expert_data)
mean_rewards = []
std_rewards = []
for train_steps in range(20):
if train_steps > 0:
if 'Bullet' in env:
model.learn(total_timesteps=25000, log_interval=1)
else:
model.learn(total_timesteps=16384, log_interval=1)
mean_reward, std_reward = evaluate_policy(
model, model.env, n_eval_episodes=10)
mean_rewards.append(mean_reward)
std_rewards.append(std_reward)
print("{0} Steps: {1}".format(train_steps, mean_reward))
np.savez(os.path.join("learners", env, "sqil_rewards_{0}".format(i)),
means=mean_rewards, stds=std_rewards)
def train_adril(env, n=0, balanced=False):
num_trajs = 20
expert_data = make_sa_dataset(env, max_trajs=num_trajs)
n_expert = len(expert_data["obs"])
expert_sa = np.concatenate((expert_data["obs"], np.reshape(expert_data["acts"], (n_expert, -1))), axis=1)
for i in range(0, n):
venv = AdRILWrapper(gym.make(env))
mean_rewards = []
std_rewards = []
# Create model
if isinstance(venv.action_space, Discrete):
model = DQN(SQLPolicy, venv, verbose=1, policy_kwargs=dict(net_arch=[64, 64]), learning_starts=1)
else:
model = SAC('MlpPolicy', venv, verbose=1, policy_kwargs=dict(net_arch=[256, 256]), ent_coef='auto',
learning_rate=linear_schedule(7.3e-4), train_freq=64, gradient_steps=64, gamma=0.98, tau=0.02)
model.replay_buffer = AdRILReplayBuffer(model.buffer_size, model.observation_space,
model.action_space, model.device, 1,
model.optimize_memory_usage, expert_data=expert_data, N_expert=num_trajs,
balanced=balanced)
if not balanced:
for j in range(len(expert_sa)):
obs = expert_data["obs"][j]
act = expert_data["acts"][j]
next_obs = expert_data["next_obs"][j]
done = expert_data["dones"][j]
model.replay_buffer.add(obs, next_obs, act, -1, done)
for train_steps in range(400):
# Train policy
if train_steps > 0:
if 'Bullet' in env:
model.learn(total_timesteps=1250, log_interval=1000)
else:
model.learn(total_timesteps=25000, log_interval=1000)
if train_steps % 1 == 0: # written to support more complex update schemes
model.replay_buffer.set_iter(train_steps)
model.replay_buffer.set_n_learner(venv.num_trajs)
# Evaluate policy
if train_steps % 20 == 0:
model.set_env(gym.make(env))
mean_reward, std_reward = evaluate_policy(
model, model.env, n_eval_episodes=10)
mean_rewards.append(mean_reward)
std_rewards.append(std_reward)
print("{0} Steps: {1}".format(int(train_steps * 1250), mean_reward))
np.savez(os.path.join("learners", env, "adril_rewards_{0}".format(i)),
means=mean_rewards, stds=std_rewards)
# Update env
if train_steps > 0:
if train_steps % 1 == 0:
venv.set_iter(train_steps + 1)
model.set_env(venv)
def train_advil(env, policy_class=AdVILPolicy, discriminator_class=AdVILDiscriminator,
iters=int(1e5), lr_pi=8e-6, lr_f=8e-4):
venv = gym.make(env)
expert_data = make_sa_dataloader(
env,
normalize=False,
batch_size=1024,
)
pi = advil_training(
expert_data,
venv,
iters=iters,
policy_class=policy_class,
discriminator_class=discriminator_class,
lr_pi=lr_pi,
lr_f=lr_f,
)
return pi
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Train expert policies.')
parser.add_argument(
'-a', '--algo', choices=['bc', 'gail', 'sqil', 'adril', 'advil', 'all'], required=True)
parser.add_argument('-e', '--env', choices=['cartpole', 'lunarlander', 'acrobot', 'pendulum', 'halfcheetah', 'walker', 'hopper', 'ant'],
required=True)
parser.add_argument('-n', '--num_runs', required=False)
args = parser.parse_args()
if args.env == "cartpole":
envname = 'CartPole-v1'
elif args.env == "lunarlander":
envname = 'LunarLander-v2'
elif args.env == "acrobot":
envname = 'Acrobot-v1'
elif args.env == "pendulum":
envname = 'Pendulum-v0'
elif args.env == "halfcheetah":
envname = 'HalfCheetahBulletEnv-v0'
elif args.env == "walker":
envname = 'Walker2DBulletEnv-v0'
elif args.env == "hopper":
envname = 'HopperBulletEnv-v0'
elif args.env == "ant":
envname = 'AntBulletEnv-v0'
else:
print("ERROR: unsupported env.")
if args.num_runs is not None and args.num_runs.isdigit():
num_runs = int(args.num_runs)
else:
num_runs = 1
if args.algo == 'bc':
train_bc(envname, num_runs)
elif args.algo == 'gail':
train_gail(envname, num_runs)
elif args.algo == 'sqil':
train_sqil(envname, num_runs)
elif args.algo == 'adril':
train_adril(envname, num_runs)
elif args.algo == 'advil':
train_advil(envname, num_runs)
elif args.algo == 'all':
train_bc(envname, num_runs)
train_gail(envname, num_runs)
train_sqil(envname, num_runs)
train_adril(envname, num_runs)
train_advil(envname, num_runs)
else:
print("ERROR: unsupported algorithm")

View File

@@ -1,129 +0,0 @@
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from itertools import chain
from typing import Callable, Union, Type, Optional, Dict, Any
# From https://github.com/DLR-RM/rl-baselines3-zoo/blob/8ea4f4a87afa548832ca17e575b351ec5928c1b0/utils/utils.py
def linear_schedule(initial_value: Union[float, str]) -> Callable[[float], float]:
"""
Linear learning rate schedule.
:param initial_value: (float or str)
:return: (function)
"""
if isinstance(initial_value, str):
initial_value = float(initial_value)
def func(progress_remaining: float) -> float:
"""
Progress will decrease from 1 (beginning) to 0
:param progress_remaining: (float)
:return: (float)
"""
return progress_remaining * initial_value
return func
class SADataset(torch.utils.data.Dataset):
def __init__(self, obs, acts, normalize):
if normalize:
obs = np.array(obs)
self.mean = obs.mean(axis=0)
self.std = obs.std(axis=0) + 1e-3
obs = (obs - self.mean) / (self.std)
self.is_normalized = True
else:
self.is_normalized = False
self.obs = torch.tensor(obs)
self.acts = torch.tensor(acts)
def __len__(self):
return len(self.obs)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
obs = self.obs[idx]
acts = self.acts[idx]
sample = {'obs': obs, 'acts': acts}
return sample
def make_sa_dataloader(envname, max_trajs=None, normalize=False, batch_size=32):
demos = np.load(
"../experts/{0}/demos.npz".format(envname), allow_pickle=True)
num_trajs = demos["num_trajs"]
if max_trajs is None:
max_trajs = num_trajs
obs = []
acts = []
for traj in range(min(max_trajs, num_trajs)):
obs.extend(demos[str(traj)].item()['states'])
acts.extend(demos[str(traj)].item()['actions'])
dataset = SADataset(obs, acts, normalize)
dataloader = DataLoader(dataset, batch_size=batch_size,
shuffle=True, num_workers=0)
return dataloader
class SADSDataset(torch.utils.data.Dataset):
def __init__(self, obs, acts, next_obs, traj_lens):
self.obs = torch.tensor(obs)
self.acts = torch.tensor(acts)
self.next_obs = torch.tensor(next_obs)
dones = [[False for _ in range(l - 2)] + [True] for l in traj_lens]
self.dones = torch.tensor(list(chain.from_iterable(dones)))
def __len__(self):
return len(self.obs)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
obs = self.obs[idx]
acts = self.acts[idx]
next_obs = self.next_obs[idx]
dones = self.dones[idx]
sample = {'obs': obs, 'acts': acts,
'next_obs': next_obs, 'dones': dones}
return sample
def make_sads_dataloader(envname, max_trajs=None):
demos = np.load(
"./experts/{0}/demos.npz".format(envname), allow_pickle=True)
num_trajs = demos["num_trajs"]
if max_trajs is None:
max_trajs = num_trajs
obs = []
next_obs = []
acts = []
lens = []
for traj in range(min(max_trajs, num_trajs)):
obs.extend(demos[str(traj)].item()['states'][:-1])
next_obs.extend(demos[str(traj)].item()['states'][1:])
acts.extend(demos[str(traj)].item()['actions'][:-1])
lens.append(len(demos[str(traj)].item()['states']))
dataset = SADSDataset(obs, acts, next_obs, lens)
dataloader = DataLoader(dataset, batch_size=32,
shuffle=False, num_workers=0, drop_last=True)
return dataloader
def make_sa_dataset(envname, max_trajs=None):
demos = np.load("../pillbox/experts/{0}/demos.npz".format(envname), allow_pickle=True)
num_trajs = demos["num_trajs"]
if max_trajs is None:
max_trajs = num_trajs
expert_states = []
expert_actions = []
expert_next_states = []
expert_dones = []
for traj in range(min(max_trajs, num_trajs)):
expert_states.extend(demos[str(traj)].item()['states'][:-1])
expert_next_states.extend(demos[str(traj)].item()['states'][1:])
expert_actions.extend(demos[str(traj)].item()['actions'][:-1])
l = len(demos[str(traj)].item()['states'])
expert_dones.extend([False for _ in range(l - 2)] + [True])
expert_data = dict()
expert_data['obs'] = np.array(expert_states)
expert_data['acts'] = np.array(expert_actions)
expert_data['next_obs'] = np.array(expert_next_states)
expert_data['dones'] = np.array(expert_dones)
return expert_data

View File

@@ -1,9 +0,0 @@
gym
numpy
psutil
scikit_learn
scipy
stable_baselines3
torch
tqdm
imitation

View File

@@ -1,78 +0,0 @@
# %%
import torch
from core.policy import SetPolicy
from tqdm import tqdm
expert_data = torch.load('intersimple-expert-data-setobs2.pt')
states, actions, _, dones = expert_data
policy = SetPolicy(actions.shape[-1])
policy = policy.cuda()
optim = torch.optim.Adam(policy.parameters(), lr=1e-4)
states = states[~dones].cuda()
actions = actions[~dones].cuda()
for _ in tqdm(range(10000)):
optim.zero_grad()
loss = -policy.log_prob(policy(states), actions).mean()
loss.backward()
optim.step()
print('Loss', loss)
torch.save(policy.state_dict(), 'bc-intersimple-setobs2.pt')
# %%
import numpy as np
from core.policy import SetPolicy
from util.wrappers import Setobs, TransformObservation, CollisionPenaltyWrapper
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
import functools
policy = SetPolicy(actions.shape[-1])
policy.load_state_dict(torch.load('bc-intersimple-setobs2.pt'))
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
env = Setobs(
TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
)
obs = env.reset()
env.render(mode='post')
for i in range(300):
#action, _ = policy.predict(torch.tensor(obs))
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action)
env.render(mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()
# %%

View File

@@ -1,74 +0,0 @@
import gym
from core.gail import gail, Buffer
from core.value import Value
from core.policy import Policy
from core.discriminator import Discriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, Minobs
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
expert_data = torch.load('intersimple-expert-data-minobs.pt')
expert_data = Buffer(*expert_data)
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=4000,
rollout_episodes=30,
rollout_steps=200,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
)
torch.save(policy.state_dict(), 'gail-intersimple-minobs.pt')

View File

@@ -1,100 +0,0 @@
# %%
import gym
from core.gail import gail, Buffer
from core.value import Value
from core.policy import Policy
from core.discriminator import Discriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, Minobs
import numpy as np
from gym.wrappers import TransformObservation
from core.reparam_module import ReparamPolicy
from torch.utils.tensorboard import SummaryWriter
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
expert_data = torch.load('intersimple-expert-data-minobs2.pt')
expert_data = Buffer(*expert_data)
# %%
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=500,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=800,
rollout_episodes=50,
rollout_steps=200,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
logger=SummaryWriter(comment='minobs2'),
)
torch.save(policy.state_dict(), 'gail-intersimple-minobs2.pt')
# %%
policy = Policy(env_fn(0).action_space.shape[0])
policy(torch.zeros(env_fn(0).observation_space.shape))
policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load('gail-intersimple-minobs2.pt'))
env = env_fn(0)
env.random_skip = False
obs = env.reset()
env.render(mode='post')
for i in range(300):
#action, _ = policy.predict(torch.tensor(obs))
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action)
env.render(mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()
# %%

View File

@@ -1,74 +0,0 @@
import gym
from core.gail import gail, Buffer
from core.value import Value
from core.policy import Policy
from core.discriminator import Discriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-4)
expert_data = torch.load('intersimple-expert-data-normobs.pt')
expert_data = Buffer(*expert_data)
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=10,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=4000,
rollout_episodes=30,
rollout_steps=100,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
)
torch.save(policy.state_dict(), 'gail-intersimple-normobs.pt')

View File

@@ -1,74 +0,0 @@
import gym
from core.gail import gail, Buffer
from core.value import SetValue
from core.policy import SetPolicy
from core.discriminator import DeepsetDiscriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, Setobs
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = SetPolicy(env_fn(0).action_space.shape[0])
value = SetValue()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
discriminator = DeepsetDiscriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
expert_data = torch.load('intersimple-expert-data-setobs.pt')
expert_data = Buffer(*expert_data)
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=500,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=4000,
rollout_episodes=30,
rollout_steps=200,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
)
torch.save(policy.state_dict(), 'gail-intersimple-setobs.pt')

View File

@@ -1,97 +0,0 @@
# %%
import gym
from core.gail import gail, Buffer
from core.value import SetValue
from core.policy import SetPolicy
from core.discriminator import RecurrentDiscriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, Setobs
import numpy as np
from gym.wrappers import TransformObservation
from core.reparam_module import ReparamPolicy
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
env_fn = lambda i: envs[i]
policy = SetPolicy(env_fn(0).action_space.shape[0])
value = SetValue()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
discriminator = RecurrentDiscriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
expert_data = torch.load('intersimple-expert-data-setobs2.pt')
expert_data = Buffer(*expert_data)
# %%
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=500,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=800,
rollout_episodes=50,
rollout_steps=200,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
)
torch.save(policy.state_dict(), 'gail-intersimple-setobs-recurrent.pt')
# %%
policy = SetPolicy(env_fn(0).action_space.shape[0])
policy(torch.zeros(env_fn(0).observation_space.shape))
policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load('gail-intersimple-setobs-recurrent.pt'))
env = env_fn(0)
obs = env.reset()
env.render(mode='post')
for i in range(300):
#action, _ = policy.predict(torch.tensor(obs))
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action)
env.render(mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()
# %%

View File

@@ -1,101 +0,0 @@
# %%
import gym
from core.gail import gail, Buffer
from core.value import SetValue
from core.policy import SetPolicy
from core.discriminator import DeepsetDiscriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, Setobs
import numpy as np
from gym.wrappers import TransformObservation
from core.reparam_module import ReparamPolicy
from torch.utils.tensorboard import SummaryWriter
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
random_skip=True,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
env_fn = lambda i: envs[i]
policy = SetPolicy(env_fn(0).action_space.shape[0])
value = SetValue()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
discriminator = DeepsetDiscriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
expert_data = torch.load('intersimple-expert-data-setobs2.pt')
expert_data = Buffer(*expert_data)
# %%
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=500,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=800,
rollout_episodes=50,
rollout_steps=200,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
logger=SummaryWriter(comment='setobs2-batchaug'),
)
torch.save(policy.state_dict(), 'gail-intersimple-setobs2.pt')
# %%
policy = SetPolicy(env_fn(0).action_space.shape[0])
policy(torch.zeros(env_fn(0).observation_space.shape))
policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load('gail-intersimple-setobs2.pt'))
env = env_fn(0)
env.random_skip = False
obs = env.reset()
env.render(mode='post')
for i in range(300):
#action, _ = policy.predict(torch.tensor(obs))
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action)
env.render(mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()
# %%

View File

@@ -1,54 +0,0 @@
import gym
from core.gail import gail, Buffer
from core.value import Value
from core.policy import Policy
from core.discriminator import Discriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper
envs = [CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-4)
expert_data = torch.load('intersimple-expert-data.pt')
expert_data = Buffer(*expert_data)
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=10,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=4000,
rollout_episodes=30,
rollout_steps=100,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
)
torch.save(policy.state_dict(), 'gail-intersimple.pt')

View File

@@ -1,97 +0,0 @@
import gym
from options.options import gail
from core.gail import Buffer
from core.value import Value
from core.policy import DiscretePolicy
from core.discriminator import Discriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Minobs
import numpy as np
from options.options import OptionsEnv
from torch.utils.tensorboard import SummaryWriter
from core.reparam_module import ReparamPolicy
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [OptionsEnv(Minobs(
TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
env_fn = lambda i: envs[i]
policy = DiscretePolicy(env_fn(0).action_space.n)
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
expert_data = torch.load('intersimple-expert-data-minobs.pt')
expert_data = Buffer(*expert_data)
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=50,
rollout_episodes=60,
rollout_steps=60,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
logger=SummaryWriter(comment='-options-minobs'),
)
torch.save(policy.state_dict(), 'gail-options-minobs.pt')
# %%
policy = DiscretePolicy(env_fn(0).action_space.n)
policy(torch.zeros(env_fn(0).observation_space.shape))
policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load('gail-options-minobs.pt'))
env = env_fn(0)
obs = env.reset()
env.render(mode='post')
for i in range(300):
#action, _ = policy.predict(torch.tensor(obs))
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action, render_mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()

View File

@@ -1,97 +0,0 @@
import gym
from options.options import gail
from core.gail import Buffer
from core.value import SetValue
from core.policy import SetDiscretePolicy
from core.discriminator import DeepsetDiscriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
import numpy as np
from options.options import OptionsEnv
from torch.utils.tensorboard import SummaryWriter
from core.reparam_module import ReparamPolicy
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [OptionsEnv(Setobs(
TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
env_fn = lambda i: envs[i]
policy = SetDiscretePolicy(env_fn(0).action_space.n)
value = SetValue()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
discriminator = DeepsetDiscriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3)
expert_data = torch.load('intersimple-expert-data-setobs.pt')
expert_data = Buffer(*expert_data)
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=150,
rollout_episodes=60,
rollout_steps=60,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
logger=SummaryWriter(comment='gail-options-setobs'),
)
torch.save(policy.state_dict(), 'gail-options-setobs.pt')
# %%
policy = SetDiscretePolicy(env_fn(0).action_space.n)
policy(torch.zeros(env_fn(0).observation_space.shape))
policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load('gail-options-setobs.pt'))
env = env_fn(0)
obs = env.reset()
env.render(mode='post')
for i in range(300):
#action, _ = policy.predict(torch.tensor(obs))
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action, render_mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()

View File

@@ -1,104 +0,0 @@
# %%
import gym
from options.options import gail
from core.gail import Buffer
from core.value import SetValue
from core.policy import SetDiscretePolicy
from core.discriminator import DeepsetDiscriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
import numpy as np
from options.options import OptionsEnv
from torch.utils.tensorboard import SummaryWriter
from core.reparam_module import ReparamPolicy
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [OptionsEnv(Setobs(
TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5)]) for _ in range(60)]
env_fn = lambda i: envs[i]
policy = SetDiscretePolicy(env_fn(0).action_space.n)
value = SetValue()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
discriminator = DeepsetDiscriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
expert_data = torch.load('intersimple-expert-data-setobs2.pt')
expert_data = Buffer(*expert_data)
# %%
def callback(epoch, value, policy):
if not epoch % 10:
torch.save(policy.state_dict(), f'gail-options-setobs2-{epoch}.pt')
torch.save(value.state_dict(), f'gail-options-setobs2-value-{epoch}.pt')
value, policy = gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=300,
rollout_episodes=60,
rollout_steps=60,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
logger=SummaryWriter(comment='gail-options-setobs2'),
callback=callback,
)
torch.save(policy.state_dict(), 'gail-options-setobs2.pt')
# %%
policy = SetDiscretePolicy(env_fn(0).action_space.n)
policy(torch.zeros(env_fn(0).observation_space.shape))
policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load('gail-options-setobs2.pt'))
env = env_fn(0)
obs = env.reset()
env.render(mode='post')
for i in range(300):
#action, _ = policy.predict(torch.tensor(obs))
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action, render_mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()

View File

@@ -1,39 +0,0 @@
import gym
from core.gail import gail, Buffer
from core.value import Value
from core.policy import Policy
from core.discriminator import Discriminator
import torch.optim
env_fn = lambda _: gym.make('Pendulum-v0')
policy = Policy(env_fn(0).action_space.shape[0])
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
expert_data = torch.load('trpo-pendulum-expert-data.pt')
expert_data = Buffer(*expert_data)
gail(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=10,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=100,
rollout_episodes=20,
rollout_steps=250,
gamma=0.99,
gae_lambda=0.9,
delta=0.01,
backtrack_coeff=0.8,
backtrack_iters=10,
)
torch.save(policy.state_dict(), 'gail-pendulum.pt')

View File

@@ -1,75 +0,0 @@
import gym
from core.gail import gail_ppo, Buffer
from core.value import Value
from core.policy import Policy
from core.discriminator import Discriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, Minobs
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3, weight_decay=1e-3)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
expert_data = torch.load('intersimple-expert-data-minobs.pt')
expert_data = Buffer(*expert_data)
value, policy = gail_ppo(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=4000,
rollout_episodes=30,
rollout_steps=100,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
)
torch.save(policy.state_dict(), 'gail-ppo-intersimple-minobs.pt')

View File

@@ -1,75 +0,0 @@
import gym
from core.gail import gail_ppo, Buffer
from core.value import Value
from core.policy import Policy
from core.discriminator import Discriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
expert_data = torch.load('intersimple-expert-data-normobs.pt')
expert_data = Buffer(*expert_data)
value, policy = gail_ppo(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=4000,
rollout_episodes=30,
rollout_steps=100,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
)
torch.save(policy.state_dict(), 'gail-ppo-intersimple-normobs.pt')

View File

@@ -1,102 +0,0 @@
# %%
import gym
from core.gail import gail_ppo, Buffer
from core.value import SetValue
from core.policy import SetPolicy
from core.discriminator import DeepsetDiscriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, Setobs
import numpy as np
from gym.wrappers import TransformObservation
from core.reparam_module import ReparamPolicy
from torch.utils.tensorboard import SummaryWriter
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
random_skip=True,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
env_fn = lambda i: envs[i]
policy = SetPolicy(env_fn(0).action_space.shape[0])
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
value = SetValue()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
discriminator = DeepsetDiscriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
expert_data = torch.load('intersimple-expert-data-setobs2.pt')
expert_data = Buffer(*expert_data)
# %%
value, policy = gail_ppo(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=500,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=800,
rollout_episodes=50,
rollout_steps=200,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
logger=SummaryWriter(comment='-ppo-setobs2'),
)
torch.save(policy.state_dict(), 'gail-ppo-intersimple-setobs2.pt')
# %%
policy = SetPolicy(env_fn(0).action_space.shape[0])
policy(torch.zeros(env_fn(0).observation_space.shape))
policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load('gail-ppo-intersimple-setobs2.pt'))
env = env_fn(0)
env.random_skip = False
obs = env.reset()
env.render(mode='post')
for i in range(300):
#action, _ = policy.predict(torch.tensor(obs))
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action)
env.render(mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()
# %%

View File

@@ -1,55 +0,0 @@
import gym
from core.gail import gail_ppo, Buffer
from core.value import Value
from core.policy import Policy
from core.discriminator import Discriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper
envs = [CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=3e-4)
expert_data = torch.load('intersimple-expert-data.pt')
expert_data = Buffer(*expert_data)
value, policy = gail_ppo(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=10,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=4000,
rollout_episodes=30,
rollout_steps=100,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
)
torch.save(policy.state_dict(), 'gail-ppo-intersimple.pt')

View File

@@ -1,96 +0,0 @@
import gym
from options.options import gail_ppo, Buffer
from core.value import Value
from core.policy import DiscretePolicy
from core.discriminator import Discriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Minobs
import numpy as np
from options.options import OptionsEnv
from torch.utils.tensorboard import SummaryWriter
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [OptionsEnv(Minobs(
TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
env_fn = lambda i: envs[i]
policy = DiscretePolicy(env_fn(0).action_space.n)
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
value = Value()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
discriminator = Discriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
expert_data = torch.load('intersimple-expert-data-minobs.pt')
expert_data = Buffer(*expert_data)
value, policy = gail_ppo(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=50,
rollout_episodes=60,
rollout_steps=60,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
logger=SummaryWriter(comment='gail-ppo-options-minobs'),
)
torch.save(policy.state_dict(), 'gail-ppo-options-minobs.pt')
# %%
policy = DiscretePolicy(env_fn(0).action_space.n)
policy(torch.zeros(env_fn(0).observation_space.shape))
policy.load_state_dict(torch.load('gail-ppo-options-minobs.pt'))
env = env_fn(0)
obs = env.reset()
env.render(mode='post')
for i in range(300):
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action, render_mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()

View File

@@ -1,96 +0,0 @@
import gym
from options.options import gail_ppo, Buffer
from core.value import SetValue
from core.policy import SetDiscretePolicy
from core.discriminator import DeepsetDiscriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
import numpy as np
from options.options import OptionsEnv
from torch.utils.tensorboard import SummaryWriter
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [OptionsEnv(Setobs(
TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
env_fn = lambda i: envs[i]
policy = SetDiscretePolicy(env_fn(0).action_space.n)
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
value = SetValue()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
discriminator = DeepsetDiscriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3)
expert_data = torch.load('intersimple-expert-data-setobs.pt')
expert_data = Buffer(*expert_data)
value, policy = gail_ppo(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=150,
rollout_episodes=60,
rollout_steps=60,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
logger=SummaryWriter(comment='gail-ppo-options-setobs'),
)
torch.save(policy.state_dict(), 'gail-ppo-options-setobs.pt')
# %%
policy = SetDiscretePolicy(env_fn(0).action_space.n)
policy(torch.zeros(env_fn(0).observation_space.shape))
policy.load_state_dict(torch.load('gail-ppo-options-setobs.pt'))
env = env_fn(0)
obs = env.reset()
env.render(mode='post')
for i in range(300):
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action, render_mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()

View File

@@ -1,103 +0,0 @@
# %%
import gym
from options.options import gail_ppo, Buffer
from core.value import SetValue
from core.policy import SetDiscretePolicy
from core.discriminator import DeepsetDiscriminator
import torch.optim
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
import functools
from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
import numpy as np
from options.options import OptionsEnv
from torch.utils.tensorboard import SummaryWriter
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [OptionsEnv(Setobs(
TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
stop_on_collision=False,
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5)]) for _ in range(60)]
env_fn = lambda i: envs[i]
policy = SetDiscretePolicy(env_fn(0).action_space.n)
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
value = SetValue()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
discriminator = DeepsetDiscriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
expert_data = torch.load('intersimple-expert-data-setobs2.pt')
expert_data = Buffer(*expert_data)
# %%
def callback(epoch, value, policy):
if not epoch % 10:
torch.save(policy.state_dict(), f'gail-ppo-options-setobs2-{epoch}.pt')
torch.save(value.state_dict(), f'gail-ppo-options-setobs2-value-{epoch}.pt')
value, policy = gail_ppo(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
epochs=200,
rollout_episodes=60,
rollout_steps=60,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
logger=SummaryWriter(comment='gail-ppo-options-setobs2'),
callback=callback,
)
torch.save(policy.state_dict(), 'gail-ppo-options-setobs2.pt')
# %%
policy = SetDiscretePolicy(env_fn(0).action_space.n)
policy(torch.zeros(env_fn(0).observation_space.shape))
policy.load_state_dict(torch.load('gail-ppo-options-setobs2.pt'))
env = env_fn(0)
obs = env.reset()
env.render(mode='post')
for i in range(300):
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
obs, reward, done, _ = env.step(action, render_mode='post')
print('step', i, 'reward', reward)
if done:
break
env.close()

File diff suppressed because one or more lines are too long

View File

@@ -1,54 +0,0 @@
import torch
import functools
from core.sampling import rollout_sb3
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
from intersim.expert import NormalizedIntersimpleExpert
from util.wrappers import CollisionPenaltyWrapper, Minobs
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
env = IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
policy = NormalizedIntersimpleExpert(env, mu=0.001)
env = Minobs(TransformObservation(
CollisionPenaltyWrapper(
env,
collision_distance=6, collision_penalty=100
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
))
expert_data = rollout_sb3(env, policy, n_episodes=64, max_steps_per_episode=200)
states, actions, rewards, dones = expert_data
print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
print(f'Observation mean', states[~dones].mean(0))
print(f'Observation std', states[~dones].std(0))
torch.save(expert_data, 'intersimple-expert-data-minobs.pt')

View File

@@ -1,53 +0,0 @@
import torch
import functools
from core.sampling import rollout_sb3
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
from intersim.expert import NormalizedIntersimpleExpert
from util.wrappers import CollisionPenaltyWrapper, Minobs
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
env = IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
policy = NormalizedIntersimpleExpert(env, mu=0.001)
env = Minobs(TransformObservation(
CollisionPenaltyWrapper(
env,
collision_distance=6, collision_penalty=100
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
))
expert_data = rollout_sb3(env, policy, n_episodes=2048, max_steps_per_episode=200)
states, actions, rewards, dones = expert_data
print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
print(f'Observation mean', states[~dones].mean(0))
print(f'Observation std', states[~dones].std(0))
torch.save(expert_data, 'intersimple-expert-data-minobs2.pt')

View File

@@ -1,54 +0,0 @@
import torch
import functools
from core.sampling import rollout_sb3
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
from intersim.expert import NormalizedIntersimpleExpert
from util.wrappers import CollisionPenaltyWrapper
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
env = IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
policy = NormalizedIntersimpleExpert(env, mu=0.001)
env = TransformObservation(
CollisionPenaltyWrapper(
env,
collision_distance=6, collision_penalty=100
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
)
expert_data = rollout_sb3(env, policy, n_episodes=64, max_steps_per_episode=200)
states, actions, rewards, dones = expert_data
print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
print(f'Observation mean', states[~dones].mean(0))
print(f'Observation std', states[~dones].std(0))
torch.save(expert_data, 'intersimple-expert-data-normobs.pt')

View File

@@ -1,54 +0,0 @@
import torch
import functools
from core.sampling import rollout_sb3
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
from intersim.expert import NormalizedIntersimpleExpert
from util.wrappers import CollisionPenaltyWrapper, Setobs
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
env = IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
policy = NormalizedIntersimpleExpert(env, mu=0.001)
env = Setobs(TransformObservation(
CollisionPenaltyWrapper(
env,
collision_distance=6, collision_penalty=100
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
))
expert_data = rollout_sb3(env, policy, n_episodes=64, max_steps_per_episode=200)
states, actions, rewards, dones = expert_data
print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
print(f'Observation mean', states[~dones].mean(0))
print(f'Observation std', states[~dones].std(0))
torch.save(expert_data, 'intersimple-expert-data-setobs.pt')

View File

@@ -1,59 +0,0 @@
import sys
sys.path.append('../../../../')
import torch
import functools
from src.core.sampling import rollout_sb3
from intersim.envs import IntersimpleLidarFlatIncrementingAgent
from intersim.envs.intersimple import speed_reward
from intersim.expert import NormalizedIntersimpleExpert
from src.util.wrappers import CollisionPenaltyWrapper, Setobs
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
env = IntersimpleLidarFlatIncrementingAgent(
loc=0,
track=4,
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
)
policy = NormalizedIntersimpleExpert(env, mu=0.001)
env = Setobs(TransformObservation(
CollisionPenaltyWrapper(
env,
collision_distance=6, collision_penalty=100
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
))
print(env.nv, 'vehicles')
expert_data = rollout_sb3(env, policy, n_episodes=150, max_steps_per_episode=200)
states, actions, rewards, dones = expert_data
print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
print(f'Observation mean', states[~dones].mean(0))
print(f'Observation std', states[~dones].std(0))
torch.save(expert_data, 'intersimple-expert-data-setobs2-loc0-track4.pt')

View File

@@ -1,25 +0,0 @@
import torch
import functools
from core.sampling import rollout_sb3
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
from intersim.expert import NormalizedIntersimpleExpert
from util.wrappers import CollisionPenaltyWrapper
env = CollisionPenaltyWrapper(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
), collision_distance=6, collision_penalty=100)
policy = NormalizedIntersimpleExpert(env.env, mu=0.001)
expert_data = rollout_sb3(env, policy, n_episodes=64, max_steps_per_episode=200)
states, actions, rewards, dones = expert_data
print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
torch.save(expert_data, 'intersimple-expert-data.pt')

View File

@@ -1,63 +0,0 @@
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from core.ppo import ppo
from core.value import Value
from core.policy import Policy
import torch.optim
import numpy as np
from gym.wrappers import TransformObservation
from util.wrappers import Minobs
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [Minobs(TransformObservation(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=0
),
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
value = Value()
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
value, policy = ppo(
env_fn=env_fn,
value=value,
policy=policy,
epochs=50,
rollout_episodes=30,
rollout_steps=200,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
v_opt=v_opt,
v_iters=1000,
)
torch.save(policy.state_dict(), 'ppo-intersimple.pt')

View File

@@ -1,62 +0,0 @@
from intersim.envs import IntersimpleLidarFlatRandom
from intersim.envs.intersimple import speed_reward
import functools
from core.ppo import ppo
from core.value import Value
from core.policy import Policy
import torch.optim
import numpy as np
from gym.wrappers import TransformObservation
from util.wrappers import Minobs
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [Minobs(TransformObservation(IntersimpleLidarFlatRandom(
n_rays=5,
reward=functools.partial(
speed_reward,
collision_penalty=1000
),
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
value = Value()
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3, weight_decay=1e-3)
value, policy = ppo(
env_fn=env_fn,
value=value,
policy=policy,
epochs=50,
rollout_episodes=30,
rollout_steps=100,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
v_opt=v_opt,
v_iters=1000,
)
torch.save(policy.state_dict(), 'ppo-intersimple.pt')

View File

@@ -1,61 +0,0 @@
from intersim.envs import IntersimpleLidarFlat
from intersim.envs.intersimple import speed_reward
import functools
from core.ppo import ppo
from core.value import Value
from core.policy import Policy
import torch.optim
import numpy as np
from gym.wrappers import TransformObservation
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
[0, -np.pi, -20, -20, -np.pi, -1e-1],
]).reshape(-1)
obs_max = np.array([
[1000, 1000, 20, np.pi, 1e-1, 0.],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
[50, np.pi, 20, 20, np.pi, 1e-1],
]).reshape(-1)
envs = [TransformObservation(IntersimpleLidarFlat(
n_rays=5,
agent=51,
reward=functools.partial(
speed_reward,
collision_penalty=10
),
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)) for _ in range(30)]
env_fn = lambda i: envs[i]
policy = Policy(env_fn(0).action_space.shape[0])
value = Value()
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
value, policy = ppo(
env_fn=env_fn,
value=value,
policy=policy,
epochs=50,
rollout_episodes=30,
rollout_steps=100,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
pi_opt=pi_opt,
pi_iters=100,
v_opt=v_opt,
v_iters=1000,
)
torch.save(policy.state_dict(), 'ppo-intersimple.pt')

Some files were not shown because too many files have changed in this diff Show More