Move files to src
This commit is contained in:
ebuehrle
@@ -1,154 +0,0 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
from intersim.collisions import state_to_polygon
|
||||
|
||||
def feasible(env, plan, ch, method='exact'):
|
||||
"""Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
|
||||
if ch == 0:
|
||||
return True
|
||||
|
||||
# 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)
|
||||
|
||||
# check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
|
||||
if method=='circle':
|
||||
valid = check_future_collisions_fast(env, [full_plan])
|
||||
elif method=='ncircles':
|
||||
valid = check_future_collisions_ncircles(env, [full_plan])
|
||||
elif method=='exact':
|
||||
valid = check_future_collisions_exact(env, [full_plan])
|
||||
else:
|
||||
raise NotImplementedError('Invalid collision-checking method')
|
||||
return valid.item()
|
||||
|
||||
def check_future_collisions_ncircles(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 check_future_collisions_circle(env, actions):
|
||||
"""Compute collision information for circular vehicle approximations
|
||||
|
||||
Args:
|
||||
env (gym.Env): current environment state
|
||||
actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
|
||||
Returns:
|
||||
states (torch.Tensor): tensor of shape (B, T, nv, 5) of future states based on the action profiles
|
||||
collision_tensor (torch.Tensor): tensor of shape (B, T, nv) of bools indicating which plan collides with which vehicles in which time frame
|
||||
false: colliding, true: not colliding
|
||||
"""
|
||||
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)
|
||||
|
||||
collision_tensor = distance > min_distance
|
||||
assert collision_tensor.shape == (B, T, nv)
|
||||
return states, collision_tensor
|
||||
|
||||
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
|
||||
"""
|
||||
_, collision_tensor = check_future_collisions_circle(env, actions)
|
||||
return collision_tensor.all(-1).all(-1)
|
||||
|
||||
def check_future_collisions_exact(env, actions):
|
||||
"""
|
||||
Checks whether `env._agent` would collide with other agents assuming `actions` as input.
|
||||
|
||||
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
|
||||
"""
|
||||
# First check with simple circle collision check
|
||||
states, collision_tensor = check_future_collisions_circle(env, actions)
|
||||
(B, T, nv, _) = states.shape
|
||||
# For those that have colliding circles, check exactly
|
||||
colliding_mask = ~collision_tensor
|
||||
|
||||
ego_states = states[:, :, env._agent:env._agent+1, :].expand(states.shape)
|
||||
assert ego_states.shape == states.shape
|
||||
|
||||
# get dimensions
|
||||
lengths = env._env._lengths.expand(states.shape[:3])
|
||||
widths = env._env._widths.expand(states.shape[:3])
|
||||
ego_lengths = lengths[:, :, env._agent:env._agent+1].expand(lengths.shape)
|
||||
ego_widths = widths[:, :, env._agent:env._agent+1].expand(widths.shape)
|
||||
assert lengths.shape == widths.shape == ego_lengths.shape == ego_widths.shape == (B, T, nv)
|
||||
|
||||
# For every collision instance between ego and other vehicle, check whether rectangles intersect
|
||||
exact_collisions = torch.zeros_like(collision_tensor[colliding_mask])
|
||||
for i, (ego_state, ego_length, ego_width, other_state, other_length, other_width) in enumerate(zip(
|
||||
ego_states[colliding_mask], ego_lengths[colliding_mask], ego_widths[colliding_mask],
|
||||
states[colliding_mask], lengths[colliding_mask], widths[colliding_mask]
|
||||
)):
|
||||
assert ego_state.shape == other_state.shape == (5,)
|
||||
assert ego_length.shape == ego_width.shape == other_length.shape == other_width.shape == ()
|
||||
p_ego = state_to_polygon(ego_state, ego_length, ego_width)
|
||||
p_other = state_to_polygon(other_state, other_length, other_width)
|
||||
exact_collisions[i] = p_ego.intersects(p_other)
|
||||
|
||||
collision_tensor[colliding_mask] = ~exact_collisions
|
||||
return collision_tensor.all(-1).all(-1)
|
||||
@@ -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)
|
||||
@@ -1,148 +0,0 @@
|
||||
import gym
|
||||
import torch
|
||||
from .collisions import feasible
|
||||
import numpy as np
|
||||
|
||||
class OptionsEnv(gym.Wrapper):
|
||||
|
||||
def __init__(self, env, options=[(v,t) for v in [0,2,4,6,8] for t in [5]], *args, **kwargs):
|
||||
"""option 0 is treated as safe fallback"""
|
||||
|
||||
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,)),
|
||||
})
|
||||
|
||||
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):
|
||||
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.options)
|
||||
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.options)))
|
||||
|
||||
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):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.observation_space = self.observation_space['obs']
|
||||
|
||||
def _after_choice(self):
|
||||
self._transition_buffer = []
|
||||
|
||||
def _after_step(self):
|
||||
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 self._transition_buffer
|
||||
|
||||
def sample_ll(self, policy):
|
||||
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):
|
||||
self.obs = {'obs': np.copy(self.s), 'mask': np.copy(self.m)}
|
||||
self.r = 0
|
||||
self.steps = 0
|
||||
|
||||
def _after_step(self):
|
||||
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 {
|
||||
'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):
|
||||
self.discriminator = discriminator
|
||||
return self.sample(policy)
|
||||
|
||||
class RenderOptions(LLOptions):
|
||||
|
||||
def _after_step(self):
|
||||
super()._after_step()
|
||||
self.env.render()
|
||||
|
||||
def close(self, *args, **kwargs):
|
||||
self.env.close(*args, **kwargs)
|
||||
|
||||
def available_actions(env, options):
|
||||
"""Return mask of available actions given current `env` state."""
|
||||
valid = np.array([feasible(env, generate_plan(env, i, options), i) for i in range(len(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, 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
|
||||
@@ -1,59 +0,0 @@
|
||||
import stable_baselines3
|
||||
from torch.distributions import Categorical
|
||||
|
||||
class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
|
||||
"""
|
||||
Class for high-level options policy (generator)
|
||||
"""
|
||||
def __init__(self, observation_space, *args, **kwargs):
|
||||
super().__init__(observation_space['obs'], *args, **kwargs)
|
||||
|
||||
def _prior_distribution(self, s):
|
||||
"""
|
||||
Return prior distribution over high-level options (before masking)
|
||||
Args:
|
||||
s (torch.tensor): observation
|
||||
Returns:
|
||||
values (torch.tensor): values from critic
|
||||
dist (torch.distributions): prior distribution over actions
|
||||
"""
|
||||
latent_pi, latent_vf, latent_sde = self._get_latent(s)
|
||||
distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
|
||||
values = self.value_net(latent_vf)
|
||||
return values, distribution.distribution
|
||||
|
||||
def predict(self, obs, eps=1e-6):
|
||||
"""
|
||||
Will mask invalid states before making action selections
|
||||
Args:
|
||||
obs: dict with keys:
|
||||
obs (torch.tensor): (*,o) true observations
|
||||
mask (torch.tensor): (*,m) mask over valid actions
|
||||
Returns:
|
||||
ch (torch.tensor): (*,a) sampled actions
|
||||
values (torch.tensor): (*,) predicted value at observation
|
||||
log_probs (torch.tensor): (*,) log probabilities of selected actions
|
||||
"""
|
||||
s, m = obs['obs'], obs['mask']
|
||||
values, prior = self._prior_distribution(s)
|
||||
posterior = Categorical((prior.probs + eps) * m)
|
||||
ch = posterior.sample()
|
||||
return ch, values, posterior.log_prob(ch)
|
||||
|
||||
def evaluate_actions(self, obs, ch, eps=1e-6):
|
||||
"""
|
||||
Evaluate particular actions
|
||||
Args:
|
||||
obs: dict with keys:
|
||||
obs (torch.tensor): (*,o) true observations
|
||||
mask (torch.tensor): (*,m) masks over valid actions
|
||||
ch (torch.tensor): (*,a) selected actions
|
||||
Returns:
|
||||
values (torch.tensor): (*,) predicted value at observation
|
||||
log_probs (torch.tensor): (*,) log probabilities of selected actions
|
||||
ent (torch.tensor): (*,) entropy of each distribution over actions
|
||||
"""
|
||||
s, m = obs['obs'], obs['mask']
|
||||
values, prior = self._prior_distribution(s)
|
||||
posterior = Categorical((prior.probs + eps) * m)
|
||||
return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train
|
||||
@@ -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()
|
||||
@@ -1,59 +0,0 @@
|
||||
import pickle
|
||||
import imitation.data.rollout as rollout
|
||||
from options import LLOptions, OptionsEnv
|
||||
from intersim.envs import NRasterized
|
||||
import itertools
|
||||
import stable_baselines3
|
||||
from policy import OptionsCnnPolicy
|
||||
from train import flatten_transitions
|
||||
import numpy as np
|
||||
|
||||
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
|
||||
@@ -1,36 +0,0 @@
|
||||
import numpy as np
|
||||
import itertools
|
||||
|
||||
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()
|
||||
@@ -1,5 +1,8 @@
|
||||
# %%
|
||||
from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
|
||||
import sys
|
||||
sys.path.append('../../../')
|
||||
|
||||
from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
|
||||
from imitation.algorithms import adversarial
|
||||
import stable_baselines3
|
||||
import torch.utils.data
|
||||
@@ -16,9 +19,9 @@ import pathlib
|
||||
from imitation.util import logger
|
||||
from stable_baselines3.common.env_util import make_vec_env
|
||||
from tqdm import tqdm
|
||||
from gail.policy import OptionsCnnPolicy
|
||||
from gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
|
||||
from gail.train import train_discriminator, train_generator
|
||||
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}
|
||||
@@ -78,7 +81,7 @@ if __name__ == '__main__':
|
||||
# %%
|
||||
model = stable_baselines3.PPO.load(model_name)
|
||||
|
||||
env = RenderOptions(NRasterized(**env_settings))
|
||||
env = RenderOptions(NRasterized(**env_settings), options=ALL_OPTIONS)
|
||||
|
||||
for s in env.sample_ll(model):
|
||||
if s['dones']:
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
# %%
|
||||
from gail.discriminator import CnnDiscriminatorFlatAction
|
||||
import sys
|
||||
sys.path.append('../../../')
|
||||
|
||||
from src.discriminator import CnnDiscriminatorFlatAction
|
||||
from imitation.algorithms import adversarial
|
||||
import stable_baselines3
|
||||
import torch.utils.data
|
||||
@@ -16,16 +19,16 @@ import pathlib
|
||||
from imitation.util import logger
|
||||
from stable_baselines3.common.env_util import make_vec_env
|
||||
from tqdm import tqdm
|
||||
from gail.policy import OptionsCnnPolicy
|
||||
from gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
|
||||
from gail.train import train_discriminator, train_generator
|
||||
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=20, expert_batch_size=32, generator_steps=1024, discount=0.99):
|
||||
def train(expert_data, epochs=100, expert_batch_size=64, generator_steps=1024, discount=0.99):
|
||||
env = NRasterizedRouteRandomAgent(**env_settings)
|
||||
env.discount = discount
|
||||
|
||||
@@ -61,27 +64,28 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di
|
||||
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
|
||||
|
||||
# %%
|
||||
if __name__ == '__main__':
|
||||
# %%
|
||||
|
||||
with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteRandomAgentw70h70mppx1.pkl", "rb") as f:
|
||||
trajectories = pickle.load(f)
|
||||
transitions = rollout.flatten_trajectories(trajectories)
|
||||
generator = train(transitions, epochs=100)
|
||||
|
||||
generator.save(model_name)
|
||||
|
||||
# %%
|
||||
def video(model_name, env):
|
||||
model = stable_baselines3.PPO.load(model_name)
|
||||
|
||||
env = RenderOptions(NRasterizedRouteRandomAgent(**env_settings))
|
||||
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user