diff --git a/scratch/arec/intersimple/gail_options_image.py b/scratch/arec/intersimple/gail_options_image.py index f6de1d1..26ef3af 100644 --- a/scratch/arec/intersimple/gail_options_image.py +++ b/scratch/arec/intersimple/gail_options_image.py @@ -18,7 +18,7 @@ from stable_baselines3.common.env_util import make_vec_env from tqdm import tqdm import logging -#logging.basicConfig(level=logging.DEBUG) +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 @@ -216,13 +216,60 @@ def check_future_collisions_fast(env, actions): 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_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): @@ -314,7 +361,7 @@ if __name__ == '__main__': transitions, env_class=env_class, env_settings=env_settings, - epochs=10, + epochs=2, discrim_batch_size=32, generator_steps=2048, discount=0.99 diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py index 72a33e2..fb2329f 100644 --- a/scratch/etienne/intersimple/gail_options_image.py +++ b/scratch/etienne/intersimple/gail_options_image.py @@ -16,6 +16,7 @@ import tempfile import pathlib from imitation.util import logger from stable_baselines3.common.env_util import make_vec_env +from tqdm import tqdm model_name = 'gail_options_image' env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2} @@ -33,17 +34,17 @@ class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy): values = self.value_net(latent_vf) return values, distribution.distribution - def predict(self, obs): + def predict(self, obs, eps=1e-6): s, m = obs['obs'], obs['mask'] values, prior = self._prior_distribution(s) - posterior = Categorical(prior.probs * m) + posterior = Categorical((prior.probs + eps) * m) ch = posterior.sample() return ch, values, posterior.log_prob(ch) - def evaluate_actions(self, obs, ch): + def evaluate_actions(self, obs, ch, eps=1e-6): s, m = obs['obs'], obs['mask'] values, prior = self._prior_distribution(s) - posterior = Categorical(prior.probs * m) + posterior = Categorical((prior.probs + eps) * m) return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train class OptionsEnv(gym.Wrapper): @@ -72,10 +73,10 @@ class OptionsEnv(gym.Wrapper): self.episode_start = False if self.done: self.s = self.env.reset() - self.m = available_actions(self.env) 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), @@ -86,18 +87,16 @@ class OptionsEnv(gym.Wrapper): assert not self.done assert self.plan - assert feasible(self.env, self.plan, self.ch) + #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.nextm = available_actions(self.env) self._after_step() self.s = self.nexts - self.m = self.nextm yield from self._transitions() @@ -132,6 +131,7 @@ class HLOptions(OptionsEnv): 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 @@ -146,7 +146,7 @@ class HLOptions(OptionsEnv): def _transitions(self): yield { - 'obs': {'obs': self.s, 'mask': self.m}, + 'obs': self.obs, 'action': self.ch, 'reward': self.r.detach(), 'episode_start': self.episode_start, @@ -250,7 +250,7 @@ def check_future_collisions_exact(env, actions): 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_coqllisions_circle(env, actions) + 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 @@ -282,12 +282,15 @@ def check_future_collisions_exact(env, actions): def feasible(env, plan, ch): """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) - valid = check_future_collisions_fast(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor - return ch == 0 or valid.item() + # 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_exact(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor + return valid.item() def flatten_transitions(transitions): return { @@ -323,7 +326,7 @@ def train_generator(env, generator, discriminator, num_samples): generator.train() -def train(expert_data, epochs=10, expert_batch_size=32, generator_steps=2048, discount=0.99): +def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99): env = NRasterized(**env_settings) env.discount = discount @@ -356,7 +359,7 @@ def train(expert_data, epochs=10, expert_batch_size=32, generator_steps=2048, di generator.tensorboard_log, ) - for _ in range(epochs): + for _ in tqdm(range(epochs)): train_discriminator(LLOptions(env), generator, discriminator, num_samples=expert_batch_size) train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps) @@ -368,8 +371,6 @@ if __name__ == '__main__': with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f: trajectories = pickle.load(f) - import pdb - pdb.set_trace() transitions = rollout.flatten_trajectories(trajectories) generator = train(transitions) @@ -380,7 +381,7 @@ if __name__ == '__main__': env = RenderOptions(NRasterized(**env_settings)) - for s in env.sample_ll(generator): + for s in env.sample_ll(model): if s['dones']: break diff --git a/scratch/etienne/intersimple/render/gail_image_ani.mp4 b/scratch/etienne/intersimple/render/gail_image_ani.mp4 deleted file mode 100644 index 1816bbb..0000000 Binary files a/scratch/etienne/intersimple/render/gail_image_ani.mp4 and /dev/null differ diff --git a/scratch/etienne/intersimple/render/gail_image_observation.mp4 b/scratch/etienne/intersimple/render/gail_image_observation.mp4 deleted file mode 100644 index 52719da..0000000 Binary files a/scratch/etienne/intersimple/render/gail_image_observation.mp4 and /dev/null differ diff --git a/scratch/etienne/intersimple/render/gail_options_image_ani.mp4 b/scratch/etienne/intersimple/render/gail_options_image_ani.mp4 new file mode 100644 index 0000000..830d7a6 Binary files /dev/null and b/scratch/etienne/intersimple/render/gail_options_image_ani.mp4 differ diff --git a/scratch/etienne/intersimple/render/gail_options_image_observation.mp4 b/scratch/etienne/intersimple/render/gail_options_image_observation.mp4 new file mode 100644 index 0000000..7937fa8 Binary files /dev/null and b/scratch/etienne/intersimple/render/gail_options_image_observation.mp4 differ