Remove old code

src/eval_main.py

@@ -10,7 +10,7 @@ import src.gail.options as options_envs
 from src.evaluation.metrics import divergence, visualize_distribution
 from src.core.policy import SetPolicy, SetDiscretePolicy
 from src.core.reparam_module import ReparamPolicy
-from src.gail2 import envs as options_envs2
+from src.options import envs as options_envs2
 
 from typing import Optional, List, Dict, Tuple
 import torch

src/evaluation/evaluation.py

@@ -6,7 +6,7 @@ from typing import Callable, Dict, Optional
 import os
 import pickle
 from tqdm import tqdm
-from src.gail2.envs import OptionsEnv
+from src.options.envs import OptionsEnv
 
 class IntersimpleEvaluation:
     """

src/gail2/envs.py

@@ -1,111 +0,0 @@
-import gym
-import numpy as np
-from src.gail2.wrappers import Wrapper, Setobs, TransformObservation
-from intersim.envs import IntersimpleLidarFlatIncrementingAgent
-
-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 NormalizedOptionsEvalEnv(**kwargs):
-    return OptionsEnv(Setobs(
-        TransformObservation(IntersimpleLidarFlatIncrementingAgent(
-            n_rays=5,
-            **kwargs,
-        ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
-    ), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5)])
-
-def NormalizedContinuousEvalEnv(**kwargs):
-    return Setobs(
-        TransformObservation(IntersimpleLidarFlatIncrementingAgent(
-            n_rays=5,
-            **kwargs,
-        ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
-    )
-
-class OptionsEnv(Wrapper):
-
-    def __init__(self, env, options):
-        super().__init__(env)
-        self.ll_action_space = env.action_space
-        self.options = options
-        self.action_space = gym.spaces.Discrete(len(options))
-        self.max_plan_length = max(t for _, t in options)
-    
-    def plan(self, option):
-        target_v, t = option
-        current_v = self.env._env.state[self.env._agent, 1].item()
-        dt = self.env._env._dt
-        a = (target_v - current_v) / (t * dt)
-        a = self.env._normalize(a)
-        a = a * np.ones((t,))
-        a += 0.01 * np.random.randn(*a.shape)
-        a = np.clip(a, self.ll_action_space.low, self.ll_action_space.high)
-        return a
-
-    def execute_plan(self, obs, option, render_mode=None):
-        observations = np.zeros((self.max_plan_length + 1, *self.env.observation_space.shape))
-        actions = np.zeros((self.max_plan_length + 1, *self.ll_action_space.shape))
-        rewards = np.zeros((self.max_plan_length + 1,))
-        env_done = np.ones((self.max_plan_length + 1,), dtype=bool)
-        plan_done = np.ones((self.max_plan_length + 1,), dtype=bool)
-        infos = []
-
-        observations[0] = obs
-        env_done[0] = False
-        for k, u in enumerate(self.plan(option)):
-            plan_done[k] = False
-            o, r, d, i = super().step(u)
-            actions[k] = u
-            rewards[k] = r
-            env_done[k] = d
-            infos.append(i)
-            observations[k+1] = o
-
-            if render_mode is not None:
-                self.env.render(render_mode)
-
-            if d:
-                break
-        
-        n_steps = k + 1
-        return observations, actions, rewards, env_done, plan_done, infos, n_steps
-
-    def step(self, action, render_mode=None):
-        a = int(action)
-        assert a == action
-        ll_obs, ll_actions, ll_rewards, ll_env_done, ll_plan_done, ll_infos, ll_steps = self.execute_plan(self.last_obs, self.options[a], render_mode)
-        hl_obs = ll_obs[ll_steps]
-        hl_reward = (ll_rewards * ~ll_plan_done).sum().item()
-        hl_done = ll_env_done[ll_steps-1].item()
-        hl_infos = {
-            'll': {
-                'observations': ll_obs,
-                'actions': ll_actions,
-                'rewards': ll_rewards,
-                'env_done': ll_env_done,
-                'plan_done': ll_plan_done,
-                'infos': ll_infos,
-                'steps': ll_steps,
-            }
-        }
-        self.last_obs = hl_obs
-        return hl_obs, hl_reward, hl_done, hl_infos
-    
-    def reset(self, *args, **kwargs):
-        self.last_obs = super().reset(*args, **kwargs)
-        return self.last_obs

src/gail2/options.py

@@ -1,128 +0,0 @@
-import gym
-import numpy as np
-import torch
-from stable_baselines3.common.vec_env import DummyVecEnv as VecEnv
-
-from core.reparam_module import ReparamPolicy
-from tqdm import tqdm
-from core.gail import Buffer, train_discriminator, roll_buffer, TerminalLogger
-from dataclasses import dataclass
-from core.trpo import trpo_step
-from core.ppo import ppo_step
-import torch.nn.functional as F
-
-@dataclass
-class OptionsRollout:
-    hl: Buffer
-    ll: Buffer
-
-def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value,
-         v_opt, v_iters, epochs, rollout_episodes, rollout_steps, gamma,
-         gae_lambda, delta, backtrack_coeff, backtrack_iters, cg_iters=10, cg_damping=0.1, wasserstein=False, wasserstein_c=None, logger=TerminalLogger()):
-
-    policy(torch.zeros(env_fn(0).observation_space.shape))
-    policy = ReparamPolicy(policy)
-
-    logger.add_scalar('expert/mean_episode_length', (~expert_data.dones).sum() / expert_data.states.shape[0])
-    logger.add_scalar('expert/mean_reward_per_episode', expert_data.rewards[~expert_data.dones].sum() / expert_data.states.shape[0])
-
-    for epoch in tqdm(range(epochs)):
-        hl_data, ll_data = rollout(env_fn, policy, rollout_episodes, rollout_steps)
-        generator_data = OptionsRollout(Buffer(*hl_data), Buffer(*ll_data))
-
-        generator_data.ll.actions += 0.1 * torch.randn_like(generator_data.ll.actions)
-
-        logger.add_scalar('gen/mean_episode_length', (~generator_data.ll.dones).sum() / generator_data.ll.states.shape[0], epoch)
-        logger.add_scalar('gen/mean_reward_per_episode', generator_data.hl.rewards[~generator_data.hl.dones].sum() / generator_data.hl.states.shape[0], epoch)
-
-        discriminator, loss = train_discriminator(expert_data, generator_data.ll, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
-        if wasserstein:
-            generator_data.ll.rewards = discriminator(generator_data.ll.states, generator_data.ll.actions)
-        else:
-            generator_data.ll.rewards = -F.logsigmoid(discriminator(generator_data.ll.states, generator_data.ll.actions))
-        logger.add_scalar('disc/final_loss', loss, epoch)
-        logger.add_scalar('disc/mean_reward_per_episode', generator_data.ll.rewards[~generator_data.ll.dones].sum() / generator_data.ll.states.shape[0], epoch)
-
-        #assert generator_data.ll.rewards.shape == generator_data.ll.dones.shape
-        generator_data.hl.rewards = torch.where(~generator_data.ll.dones, generator_data.ll.rewards, torch.tensor(0.)).sum(-1)
-
-        value, policy = trpo_step(value, policy, generator_data.hl.states, generator_data.hl.actions, generator_data.hl.rewards, generator_data.hl.dones, gamma, gae_lambda, delta, backtrack_coeff, backtrack_iters, v_opt, v_iters, cg_iters, cg_damping)
-        expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
-    
-    return value, policy
-
-def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value,
-         v_opt, v_iters, epochs, rollout_episodes, rollout_steps, gamma,
-         gae_lambda, clip_ratio, pi_opt, pi_iters, target_kl=None, max_grad_norm=None, wasserstein=False, wasserstein_c=None, logger=TerminalLogger()):
-
-    logger.add_scalar('expert/mean_episode_length', (~expert_data.dones).sum() / expert_data.states.shape[0])
-    logger.add_scalar('expert/mean_reward_per_episode', expert_data.rewards[~expert_data.dones].sum() / expert_data.states.shape[0])
-
-    for epoch in range(epochs):
-        hl_data, ll_data = rollout(env_fn, policy, rollout_episodes, rollout_steps)
-        generator_data = OptionsRollout(Buffer(*hl_data), Buffer(*ll_data))
-
-        generator_data.ll.actions += 0.1 * torch.randn_like(generator_data.ll.actions)
-
-        logger.add_scalar('gen/mean_episode_length', (~generator_data.ll.dones).sum() / generator_data.ll.states.shape[0], epoch)
-        logger.add_scalar('gen/mean_reward_per_episode', generator_data.hl.rewards[~generator_data.hl.dones].sum() / generator_data.hl.states.shape[0], epoch)
-
-        discriminator, loss = train_discriminator(expert_data, generator_data.ll, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
-        if wasserstein:
-            generator_data.ll.rewards = discriminator(generator_data.ll.states, generator_data.ll.actions)
-        else:
-            generator_data.ll.rewards = -F.logsigmoid(discriminator(generator_data.ll.states, generator_data.ll.actions))
-        logger.add_scalar('disc/final_loss', loss, epoch)
-        logger.add_scalar('disc/mean_reward_per_episode', generator_data.ll.rewards[~generator_data.ll.dones].sum() / generator_data.ll.states.shape[0], epoch)
-        
-        #assert generator_data.ll.rewards.shape == generator_data.ll.dones.shape
-        generator_data.hl.rewards = torch.where(~generator_data.ll.dones, generator_data.ll.rewards, torch.tensor(0.)).sum(-1)
-
-        value, policy = ppo_step(value, policy, generator_data.hl.states, generator_data.hl.actions, generator_data.hl.rewards, generator_data.hl.dones, clip_ratio, gamma, gae_lambda, pi_opt, pi_iters, v_opt, v_iters, target_kl, max_grad_norm)
-        expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
-    
-    return value, policy
-
-def rollout(env_fn, policy, n_episodes, max_steps_per_episode):
-    env = env_fn(0)
-
-    states = torch.zeros(n_episodes, max_steps_per_episode + 1, *env.observation_space.shape)
-    actions = torch.zeros(n_episodes, max_steps_per_episode + 1, *env.action_space.shape)
-    rewards = torch.zeros(n_episodes, max_steps_per_episode + 1)
-    dones = torch.ones(n_episodes, max_steps_per_episode + 1, dtype=bool)
-
-    ll_states = torch.zeros(n_episodes, max_steps_per_episode, env.max_plan_length + 1, *env.observation_space.shape)
-    ll_actions = torch.zeros(n_episodes, max_steps_per_episode, env.max_plan_length + 1, *env.ll_action_space.shape)
-    ll_rewards = torch.zeros(n_episodes, max_steps_per_episode, env.max_plan_length + 1)
-    ll_dones = torch.ones(n_episodes, max_steps_per_episode, env.max_plan_length + 1, dtype=bool)
-
-    env = VecEnv(list(map(lambda i: (lambda: env_fn(i)), range(n_episodes))))
-
-    states[:, 0] = torch.tensor(env.reset()).clone().detach()
-    dones[:, 0] = False
-
-    for s in tqdm(range(max_steps_per_episode), 'Rollout'):
-        actions[:, s] = policy.sample(policy(states[:, s])).clone().detach()
-
-        clipped_actions = actions[:, s]
-        if isinstance(env.action_space, gym.spaces.Box):
-            clipped_actions = torch.clamp(clipped_actions, torch.from_numpy(env.action_space.low), torch.from_numpy(env.action_space.high))
-
-        o, r, d, info = env.step(clipped_actions)
-        states[:, s + 1] = torch.tensor(o).clone().detach()
-        rewards[:, s] = torch.tensor(r).clone().detach()
-        dones[:, s + 1] = torch.tensor(d).clone().detach()
-
-        ll_states[:, s] = torch.from_numpy(np.stack([i['ll']['observations'] for i in info])).clone().detach()
-        ll_actions[:, s] = torch.from_numpy(np.stack([i['ll']['actions'] for i in info])).clone().detach()
-        ll_rewards[:, s] = torch.from_numpy(np.stack([i['ll']['rewards'] for i in info])).clone().detach()
-        ll_dones[:, s] = torch.from_numpy(np.stack([i['ll']['plan_done'] for i in info])).clone().detach()
-
-    dones = dones.cumsum(1) > 0
-
-    states = states[:, :max_steps_per_episode]
-    actions = actions[:, :max_steps_per_episode]
-    rewards = rewards[:, :max_steps_per_episode]
-    dones = dones[:, :max_steps_per_episode]
-
-    return (states, actions, rewards, dones), (ll_states, ll_actions, ll_rewards, ll_dones)

src/gail2/test_options.py

@@ -1,54 +0,0 @@
-from intersim.envs import IntersimpleLidarFlat
-from options import OptionsEnv
-import gym
-import numpy as np
-
-def test_obs_shape():
-    options = [(0, 5), (5, 5), (10, 5)]
-    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
-    assert env.reset().shape == (36,)
-
-def test_act_space():
-    options = [(0, 5), (5, 5), (10, 5)]
-    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
-    assert env.action_space == gym.spaces.Discrete(3)
-
-def test_plan():
-    options = [(0, 5), (5, 5), (10, 5)]
-    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
-    env.reset()
-    plan = env.plan(options[0])
-    assert np.allclose(plan, -13.998268127441406 * np.ones((5,)))
-
-def test_plan2():
-    options = [(0, 5), (5, 5), (10, 5)]
-    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
-    obs = env.reset()
-    states, actions, rewards, dones, plan_done, infos, n_steps = env.execute_plan(obs, options[0])
-    assert states.shape == (6, 36)
-    assert rewards.shape == (6,)
-    assert dones.shape == (6,)
-    assert len(infos) == 5
-
-def test_step():
-    options = [(0, 5), (5, 5), (10, 5)]
-    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
-    env.reset()
-    obs, reward, done, _ = env.step(0)
-    assert obs.shape == (36,)
-    assert reward == 5.0
-    assert done == False
-
-def test_ll_step():
-    options = [(0, 5), (5, 5), (10, 5)]
-    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
-    env.reset()
-    _, _, _, info = env.step(0)
-    assert info['ll']['observations'].shape == (6, 36)
-    assert info['ll']['actions'].shape == (6, 1)
-    assert info['ll']['rewards'].shape == (6,)
-    assert info['ll']['env_done'].shape == (6,)
-    assert info['ll']['plan_done'].shape == (6,)
-    assert info['ll']['plan_done'][5] == True
-    assert info['ll']['steps'] == 5
-    assert len(info['ll']['infos']) == 5

src/gail2/wrappers.py

@@ -1,74 +0,0 @@
-import numpy as np
-import gym
-
-class Wrapper(gym.Wrapper):
-    def __getattr__(self, name):
-        return getattr(self.env, name)
-
-class TransformObservation(gym.wrappers.TransformObservation):
-    def __getattr__(self, name):
-        return getattr(self.env, name)
-
-class CollisionPenaltyWrapper(Wrapper):
-
-    def __init__(self, env, collision_distance, collision_penalty, *args, **kwargs):
-        super().__init__(env, *args, **kwargs)
-        self.penalty = collision_penalty
-        self.distance = collision_distance
-    
-    def step(self, action):
-        obs, reward, done, info = super().step(action)
-        reward = -self.penalty if (obs.reshape(-1, 6)[1:, 0] < self.distance).any() else reward
-
-        self.env._rewards.pop()
-        self.env._rewards.append(reward)
-
-        return obs, reward, done, info
-
-class Minobs(Wrapper):
-    """ Meant to be used as wrapper around LidarObservation """
-
-    def __init__(self, env, *args, **kwargs):
-        super().__init__(env, *args, **kwargs)
-        n_rays = int(self.observation_space.shape[0] / 6) - 1
-        self.observation_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=((1 + n_rays) * 2,))
-    
-    def minobs(self, obs):
-        """ ego v, psidot ; (for each ray,) rel. distance, rel. velocity in ego forward direction """
-        obs = obs.reshape(-1, 6)
-        obs = np.concatenate((obs[:1, [2, 4]], obs[1:, [0, 2]]), axis=0)
-        return obs.reshape(-1)
-
-    def reset(self):
-        return self.minobs(super().reset())
-    
-    def step(self, action):
-        obs, reward, done, info = super().step(action)
-        return self.minobs(obs), reward, done, info
-
-class Setobs(Wrapper):
-    """ Meant to be used as wrapper around LidarObservation """
-
-    def __init__(self, env, *args, **kwargs):
-        super().__init__(env, *args, **kwargs)
-        self.n_rays = int(self.observation_space.shape[0] / 6) - 1
-        self.observation_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.n_rays, 6))
-    
-    def obs(self, obs):
-        obs = obs.reshape(-1, 6)
-
-        ego = obs[:1, [2, 4]] # v, psidot
-        ego = np.tile(ego, (self.n_rays, 1))
-        
-        other = obs[1:, [0, 1, 2]] # distance, angle, velocity component in ego forward direction
-        other = np.stack((other[:, 0], np.cos(other[:, 1]), np.sin(other[:, 1]), other[:, 2]), axis=-1)
-        
-        obs = np.concatenate((ego, other), axis=-1) 
-        return obs
-
-    def reset(self):
-        return self.obs(super().reset())
-    
-    def step(self, action):
-        obs, reward, done, info = super().step(action)
-        return self.obs(obs), reward, done, info

src/options/envs.py

@@ -1,6 +1,6 @@
 import gym
 import numpy as np
-from src.gail2.wrappers import Wrapper, Setobs, TransformObservation
+from src.util.wrappers import Wrapper, Setobs, TransformObservation
 from intersim.envs import IntersimpleLidarFlatIncrementingAgent
 
 obs_min = np.array([