diff --git a/scratch/etienne/intersimple/data/expert.py b/scratch/etienne/intersimple/data/expert.py index fe4be59..fcecf42 100644 --- a/scratch/etienne/intersimple/data/expert.py +++ b/scratch/etienne/intersimple/data/expert.py @@ -1,4 +1,4 @@ -from intersim.envs.intersimple import Intersimple +from intersim.envs.intersimple import Intersimple, InfoFilter from stable_baselines3.common.policies import BasePolicy import gym import intersim.envs.intersimple @@ -119,6 +119,7 @@ def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomA 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, @@ -127,7 +128,8 @@ def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomA sample_until=rollout.make_sample_until( min_timesteps=min_timesteps, min_episodes=min_episodes, - ) + ), + exclude_infos=not include_infos, ) if __name__ == '__main__': diff --git a/scratch/etienne/intersimple/gail_options_image_alltracks.py b/scratch/etienne/intersimple/gail_options_image_alltracks.py index 71d4741..84b418b 100644 --- a/scratch/etienne/intersimple/gail_options_image_alltracks.py +++ b/scratch/etienne/intersimple/gail_options_image_alltracks.py @@ -331,7 +331,7 @@ if __name__ == '__main__': #env_class = NRasterized #env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2} - files = ['../../../expert_data/DR_USA_Roundabout_FT0/track%04i/expert.pkl'%(i) for i in range(5)] + files = ['../../../expert_data/DR_USA_Roundabout_FT/track%04i/expert.pkl'%(i) for i in range(5)] transitions=load_experts(files) generator = train( diff --git a/scratch/johannes/evaluation.py b/scratch/johannes/evaluation.py new file mode 100644 index 0000000..67fec87 --- /dev/null +++ b/scratch/johannes/evaluation.py @@ -0,0 +1,90 @@ + +def evaluate_policy_simple( + model, + env: gym.Env, + n_eval_episodes: int = 10, + deterministic: bool = True, + render: bool = False, + callback = None, + reward_threshold = None, + return_episode_rewards: bool = False, + warn: bool = True, +): + """ + Runs policy for ``n_eval_episodes`` episodes and returns average reward. + If a vector env is passed in, this divides the episodes to evaluate onto the + different elements of the vector env. This static division of work is done to + remove bias. See https://github.com/DLR-RM/stable-baselines3/issues/402 for more + details and discussion. + + .. note:: + If environment has not been wrapped with ``Monitor`` wrapper, reward and + episode lengths are counted as it appears with ``env.step`` calls. If + the environment contains wrappers that modify rewards or episode lengths + (e.g. reward scaling, early episode reset), these will affect the evaluation + results as well. You can avoid this by wrapping environment with ``Monitor`` + wrapper before anything else. + + :param model: The RL agent you want to evaluate. + :param env: The gym environment or ``VecEnv`` environment. + :param n_eval_episodes: Number of episode to evaluate the agent + :param deterministic: Whether to use deterministic or stochastic actions + :param render: Whether to render the environment or not + :param callback: callback function to do additional checks, + called after each step. Gets locals() and globals() passed as parameters. + :param reward_threshold: Minimum expected reward per episode, + this will raise an error if the performance is not met + :param return_episode_rewards: If True, a list of rewards and episode lengths + per episode will be returned instead of the mean. + :param warn: If True (default), warns user about lack of a Monitor wrapper in the + evaluation environment. + :return: Mean reward per episode, std of reward per episode. + Returns ([float], [int]) when ``return_episode_rewards`` is True, first + list containing per-episode rewards and second containing per-episode lengths + (in number of steps). + """ + episode_rewards = [] + episode_lengths = [] + + episode_counts = 0 + + current_rewards = 0 + current_lengths = 0 + observations = env.reset() + states = None + while (episode_counts < n_eval_episodes): + actions, states = model.predict(observations, state=states, deterministic=deterministic) + observations, rewards, dones, infos = env.step(actions) + print(env._env.t) + current_rewards += rewards + current_lengths += 1 + + # unpack values so that the callback can access the local variables + reward = rewards + done = dones + info = infos + if info['collision']: + print("COLLISION") + + if callback is not None: + callback(locals(), globals()) + + if dones: + episode_rewards.append(current_rewards) + episode_lengths.append(current_lengths) + episode_counts += 1 + current_rewards = 0 + current_lengths = 0 + if states is not None: + states *= 0 + + if render: + env.render() + + mean_reward = np.mean(episode_rewards) + std_reward = np.std(episode_rewards) + if reward_threshold is not None: + assert mean_reward > reward_threshold, "Mean reward below threshold: " f"{mean_reward:.2f} < {reward_threshold:.2f}" + if return_episode_rewards: + return episode_rewards, episode_lengths + return mean_reward, std_reward diff --git a/scratch/johannes/gail_options_image.py b/scratch/johannes/gail_options_image.py new file mode 100644 index 0000000..1ea179d --- /dev/null +++ b/scratch/johannes/gail_options_image.py @@ -0,0 +1,96 @@ +# %% +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, speed_reward +import itertools +import functools +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 +from src.evaluation.evaluation import Evaluation + +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 epoch 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) + + eval_env = env + ev = Evaluation(eval_env, n_eval_episodes=100) + ev.evaluate(epoch, generator, discriminator, expert_data) + + return generator + +# %% +if __name__ == '__main__': + # %% + + with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedInfoAgent51w36h36mppx2.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) diff --git a/scratch/johannes/raytune_simple.py b/scratch/johannes/raytune_simple.py new file mode 100644 index 0000000..5979c11 --- /dev/null +++ b/scratch/johannes/raytune_simple.py @@ -0,0 +1,69 @@ +"""This example demonstrates basic Ray Tune random search and grid search.""" +import time + +import ray +from ray import tune + + +def evaluation_fn(step, width, height): + time.sleep(0.1) + return (0.1 + width * step / 100)**(-1) + height * 0.1 + +def easy_objective(config): + # Hyperparameters + width, height = config["width"], config["height"] + + mydata = ray.get(ray_data) + print(mydata) + + for step in range(config["steps"]): + # Iterative training function - can be any arbitrary training procedure + intermediate_score = evaluation_fn(step, width, height) + # Feed the score back back to Tune. + tune.report(iterations=step, mean_loss=intermediate_score) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing") + parser.add_argument( + "--server-address", + type=str, + default=None, + required=False, + help="The address of server to connect to if using " + "Ray Client.") + args, _ = parser.parse_known_args() + if args.server_address is not None: + ray.init(f"ray://{args.server_address}") + else: + ray.init(configure_logging=False) + + # This will do a grid search over the `activation` parameter. This means + # that each of the two values (`relu` and `tanh`) will be sampled once + # for each sample (`num_samples`). We end up with 2 * 50 = 100 samples. + # The `width` and `height` parameters are sampled randomly. + # `steps` is a constant parameter. + + import numpy as np + N = 3 + data = np.random.rand(N,N,N) + ray_data = ray.put(data) + + + analysis = tune.run( + easy_objective, + metric="mean_loss", + mode="min", + num_samples=5 if args.smoke_test else 50, + config={ + "steps": 5 if args.smoke_test else 100, + "width": tune.uniform(0, 20), + "height": tune.uniform(-100, 100), + "activation": tune.grid_search(["relu", "tanh"]) + }) + + print("Best hyperparameters found were: ", analysis.best_config) \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py index d44975d..aebb694 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,3 +1,3 @@ from src.data.expert_data import generate_expert_data, load_expert_data from src.data.data_utils import InteractionDatasetSingleAgent -from src.metrics import metrics \ No newline at end of file +from src.evaluation.metrics import metrics \ No newline at end of file diff --git a/src/evaluation/__init__.py b/src/evaluation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/evaluation/evaluation.py b/src/evaluation/evaluation.py new file mode 100644 index 0000000..cbbd0da --- /dev/null +++ b/src/evaluation/evaluation.py @@ -0,0 +1,87 @@ +import torch +import numpy as np +from stable_baselines3.common.vec_env import VecEnv +from stable_baselines3.common.evaluation import evaluate_policy + +from intersim.envs.intersimple import Intersimple +from src.evaluation.metrics import nanmean, divergence, visualize_distribution + +class Evaluation: + def __init__(self, eval_env, n_eval_episodes=10): + # if env is a VecEnv, the code needs to be adapted, since the callback will be called after each step, + # so transitions of different envs will be mixed and the total number of episodes could be larger than n_eval_episodes! + assert not isinstance(eval_env, VecEnv) + self.env = eval_env + self.n_eval_episodes = n_eval_episodes + self.reset() + + def reset(self): + self._n_collisions = 0 + self._trajectories = [] + self._episode_done = True + self._accelerations = [] + + def evaluate(self, epoch, generator, discriminator, expert_data): + self.reset() + metrics = {} + + episode_rewards, episode_lengths = evaluate_policy( + generator, + self.env, + n_eval_episodes=self.n_eval_episodes, + callback=self.evaluate_policy_callback, + return_episode_rewards=True + ) + + collision_rate = self._n_collisions / self.n_eval_episodes + metrics['collision_rate'] = collision_rate + + assert len(self._trajectories) >= self.n_eval_episodes + + # velocities produced by generator + policy_velocities = torch.cat([torch.stack(t)[:,2] for t in self._trajectories]) + # if episodes terminate without collisions, then the state is fully nan + policy_velocities = policy_velocities[~torch.isnan(policy_velocities)] + + # expert velocities + extract_state = lambda info: info['projected_state'][info['agent']] + expert_velocities = torch.stack([extract_state(info) for info in expert_data.infos])[:,2] + expert_velocities = expert_velocities[~torch.isnan(expert_velocities)] + + metrics['avg_velocity_loss'] = (expert_velocities.mean() - policy_velocities.mean()).item() + metrics['velocity_divergence'] = divergence(policy_velocities, expert_velocities, type='js') + + + # accelerations produced by generator + policy_accelerations = torch.tensor(self._accelerations) + # expert accelerations + extract_accel = lambda info: info['action_taken'][info['agent']] + expert_accelerations = torch.cat([extract_accel(info) for info in expert_data.infos]) + + metrics['acceleration_divergence'] = divergence(policy_accelerations, expert_accelerations, type='js') + visualize_distribution(expert_accelerations, policy_accelerations, 'output/_action_viz{:02}'.format(epoch)) + + print(metrics) + return metrics + + def evaluate_policy_callback(self, local_vars, global_vars): + venv_i = local_vars['i'] + info = local_vars['info'] + done = local_vars['done'] + _agent = info['agent'] + env = local_vars['env'].envs[venv_i] + assert isinstance(env, Intersimple) + + # Increase collision counter if episode terminated with a collision + if info['collision']: + assert done + self._n_collisions += 1 + + # if last episode is done, start new trajectory + # this is currently not necessary, only if velocity is to be averaged over individual trajectories first + # and then averaging over all trajectories + if self._episode_done: + self._trajectories.append([]) + self._trajectories[-1].append(info['projected_state'][_agent]) + self._accelerations.append(info['action_taken'][_agent]) + self._episode_done = done diff --git a/src/metrics.py b/src/evaluation/metrics.py similarity index 100% rename from src/metrics.py rename to src/evaluation/metrics.py diff --git a/src/gail/options.py b/src/gail/options.py index 4c09935..b39fe1f 100644 --- a/src/gail/options.py +++ b/src/gail/options.py @@ -40,7 +40,7 @@ class OptionsEnv(gym.Wrapper): # action 0 is considered safe fallback self.m[0] = True - self.ch, self.value, self.log_prob = generator.policy.predict({ + self.ch, self.value, self.log_prob = generator.policy.forward({ 'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device), 'mask': self.m.unsqueeze(0).to(generator.policy.device), }) diff --git a/src/gail/train.py b/src/gail/train.py index 31edd46..ee2715c 100644 --- a/src/gail/train.py +++ b/src/gail/train.py @@ -9,10 +9,11 @@ def flatten_transitions(transitions): 'dones': np.stack(list(t['dones'] for t in transitions), axis=0), } -def train_discriminator(env, generator, discriminator, num_samples): +def train_discriminator(env, generator, discriminator, num_samples, n_updates=1): transitions = list(itertools.islice(env.sample_ll(generator), num_samples)) generator_samples = flatten_transitions(transitions) - discriminator.train_disc(gen_samples=generator_samples) + for _ in range(n_updates): + 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)) diff --git a/src/policies/options.py b/src/policies/options.py index 9bf71c9..851edd4 100644 --- a/src/policies/options.py +++ b/src/policies/options.py @@ -22,7 +22,7 @@ class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy): values = self.value_net(latent_vf) return values, distribution.distribution - def predict(self, obs, eps=1e-6): + def forward(self, obs, eps=1e-6): """ Will mask invalid states before making action selections Args: