Refactor options GAIL training script

scratch/etienne/intersimple/gail/collisions.py

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

scratch/etienne/intersimple/gail/options.py

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

scratch/etienne/intersimple/gail/policy.py

@@ -0,0 +1,59 @@
+import stable_baselines3
+from torch.distributions import Categorical
+
+class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
+    """
+    Class for high-level options policy (generator)
+    """
+    def __init__(self, observation_space, *args, **kwargs):
+        super().__init__(observation_space['obs'], *args, **kwargs)
+
+    def _prior_distribution(self, s):
+        """
+        Return prior distribution over high-level options (before masking)
+        Args:
+            s (torch.tensor): observation
+        Returns:
+            values (torch.tensor): values from critic
+            dist (torch.distributions): prior distribution over actions
+        """
+        latent_pi, latent_vf, latent_sde = self._get_latent(s)
+        distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
+        values = self.value_net(latent_vf)
+        return values, distribution.distribution
+
+    def predict(self, obs, eps=1e-6):
+        """
+        Will mask invalid states before making action selections
+        Args:
+            obs: dict with keys:
+                obs (torch.tensor): (*,o)  true observations
+                mask (torch.tensor): (*,m) mask over valid actions
+        Returns:
+            ch (torch.tensor): (*,a) sampled actions
+            values (torch.tensor): (*,) predicted value at observation
+            log_probs (torch.tensor): (*,) log probabilities of selected actions
+        """
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical((prior.probs + eps) * m)
+        ch = posterior.sample()
+        return ch, values, posterior.log_prob(ch)
+
+    def evaluate_actions(self, obs, ch, eps=1e-6):
+        """
+        Evaluate particular actions
+        Args:
+            obs: dict with keys:
+                obs (torch.tensor): (*,o) true observations
+                mask (torch.tensor): (*,m) masks over valid actions
+            ch (torch.tensor): (*,a) selected actions
+        Returns:
+            values (torch.tensor): (*,) predicted value at observation
+            log_probs (torch.tensor): (*,) log probabilities of selected actions
+            ent (torch.tensor): (*,) entropy of each distribution over actions
+        """
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical((prior.probs + eps) * m)
+        return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train

scratch/etienne/intersimple/gail/test_options.py

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

scratch/etienne/intersimple/gail/train.py

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