adding rollout and trajectory flattening utilities. including old implementation in this commit

generate_demos.sh

@@ -6,4 +6,5 @@
 # expert_class:str='NRasterizedRouteIncrementingAgent',  
 # expert_args:dict={mu:0.001}):
 
-python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]'
+# python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]'
+python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]' --filename='expert_new.pkl'

src/data/expert.py

@@ -5,13 +5,16 @@ import gym
 import intersim.envs.intersimple
 import pickle
 from tqdm import tqdm
-import imitation.data.rollout as rollout
 from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
-from imitation.data.wrappers import RolloutInfoWrapper
 import copy
 import os
 import numpy as np
 
+import imitation.data.rollout as rollout
+from imitation.data.wrappers import RolloutInfoWrapper
+
+from src.util.rollout import rollout_and_save, flatten_trajectories, make_sample_until
+
 class IntersimExpert(BasePolicy):
     
     def __init__(self, intersim_env, mu=0, *args, **kwargs):
@@ -149,9 +152,10 @@ def single_agent_expert(expert='NormalizedIntersimpleExpert',
     Expert = globals()[expert]
     env = Env(**env_args)
     policy = Expert(env, **policy_args)
+    # single_agent_demonstrations_old(env, policy, **kwargs)
     single_agent_demonstrations(env, policy, **kwargs)
 
-def single_agent_demonstrations(env, policy, 
+def single_agent_demonstrations_old(env, policy, 
                                     path=None, min_timesteps=None, 
                                     min_episodes=None, video=False,
                                     env_args={}, policy_args={}):
@@ -183,8 +187,7 @@ def single_agent_demonstrations(env, policy,
     path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
     suntil = rollout.make_sample_until(
         min_timesteps=min_timesteps, 
-            min_episodes=min_episodes,
+        min_episodes=min_episodes)
-        )
     
     rollout.rollout_and_save(
         path=path,
@@ -194,6 +197,41 @@ def single_agent_demonstrations(env, policy,
         rng=NoShuffleRNG()
     )
 
+def single_agent_demonstrations(env, policy, 
+                                path=None, min_timesteps=None, 
+                                min_episodes=None, video=False,
+                                env_args={}, policy_args={}):
+    """Rollout and save expert demos.
+    
+    Usage:
+        python -m intersimple.expert <flags>
+    Args: 
+        env (class): intersimple environment
+        policy (BasePolicy): intersimple policy
+        path (str): path to store output
+        min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save
+        min_episodes (int): min number of episodes for call to rollout.rollout_and_save
+        video (bool): whether to save a video of the expert until a single environment instantiation stops
+        env_args (dict): dictionary of kwargs when instantiating environment class
+        policy_args (dict): dictionary of kwargs when instantiating Expert policy
+    """
+    
+    if min_timesteps is None and min_episodes is None:
+        min_episodes = env.nv # one episode per vehicle being controlled in environment (hopefully an incrementing agent environment)
+
+    if video:
+        save_video(env, policy)
+
+    path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
+    suntil = make_sample_until(min_timesteps=min_timesteps,min_episodes=min_episodes)
+    
+    rollout_and_save(
+        path=path,
+        policy=policy,
+        env=env,
+        sample_until=suntil
+    )
+
 def multi_agent_demonstrations(expert='IntersimExpert',path=None, env_args={}, policy_args={}):
     """
     Run and save the `intersim' multiagent environment demonstration

src/util/rollout.py

@@ -0,0 +1,206 @@
+# Borrowed heavily from https://github.com/HumanCompatibleAI/imitation/tree/master/src/imitation/data
+
+import pickle
+import numpy as np
+import logging
+import os
+import pathlib
+from typing import Optional, List, Dict
+
+def generate_trajectories(
+    policy,
+    env,
+    sample_until,
+    rng: Optional[np.random.RandomState] = None, # np.random to shuffle
+) -> List[dict]:
+    """Generate trajectory dictionaries from a policy and an environment.
+
+    Args:
+      policy: a stable_baselines3 policy or algorithm trained on the gym environment
+      env: The environment to interact with.
+      sample_until: A function determining the termination condition.
+          It takes a sequence of trajectories, and returns a bool.
+          Most users will want to use one of `min_episodes` or `min_timesteps`.
+      rng: used for shuffling trajectories.
+
+    Returns:
+      Sequence of trajectories, satisfying `sample_until`. 
+    """
+    trajectories = []
+    while not sample_until(trajectories):
+        
+        # sample a trajectory
+        ob, done = env.reset(), False
+        ob_list, rew_list, act_list, info_list = [], [], [], []
+        while not done:
+            act, _ = policy.predict(ob)
+            next_ob, rew, done, info = env.step(act) # ignore infos
+            ob_list.append(ob)
+            act_list.append(act)
+            rew_list.append(rew)
+            info_list.append(info)
+            ob = next_ob
+        ob_list.append(ob)
+
+        traj = {
+            'obs':np.stack(ob_list),
+            'acts':np.stack(act_list),
+            'rews':np.stack(rew_list),
+            'infos': info_list,
+            'terminal': True
+        }
+        trajectories.append(traj)
+    
+    # Shuffle trajectories
+    if rng:
+        rng.shuffle(trajectories)
+
+    # Sanity checks.
+    for trajectory in trajectories:
+        n_steps = len(trajectory['acts'])
+        # extra 1 for the end
+        exp_obs = (n_steps + 1,) + env.observation_space.shape
+        real_obs = trajectory['obs'].shape
+        assert real_obs == exp_obs, f"expected shape {exp_obs}, got {real_obs}"
+        exp_act = (n_steps,) + env.action_space.shape
+        real_act = trajectory['acts'].shape
+        assert real_act == exp_act, f"expected shape {exp_act}, got {real_act}"
+        exp_rew = (n_steps,)
+        real_rew = trajectory['rews'].shape
+        assert real_rew == exp_rew, f"expected shape {exp_rew}, got {real_rew}"
+
+    return trajectories
+
+
+def flatten_trajectories(trajectory_list: list) -> list:
+    """
+    Turn a list of trajectories into a (longer) list of transitions with appropriate fields
+
+    Args:
+        trajectory_list (list): list of trajectory dicts with keys:
+            obs (np.ndarray): (T, *O) tensor of all observations in T-step trajectory
+            acts (np.ndarray): (T-1, *A) tensor of all actions in T-step trajectory
+            infos (list[dict]): (T-1)-length list of all information dictionaries
+            terminal (bool): True if the trajectory ends at the last step
+            rews (np.ndarray): (T-1, 1) tensor of rewards along trajectory
+    Returns:
+        transition_list (list): list of all transition dicts with keys:
+            obs (np.ndarray): (*O) tensor of single-step observation
+            acts (np.ndarray): (*A) tensor of single-step action
+            infos (dict): single-step information dictionary
+            next_obs (np.ndarray): (*O) tensor of next observation
+            rews (np.ndarray): (1) tensor of single-step reward
+            dones (bool): whether state is terminal
+    """
+    transition_list = []
+    for traj in trajectory_list:
+        T = traj['obs'].shape[0]
+        if traj['infos']:
+            infos = traj['infos'] 
+        else: 
+            infos = [{}] * T
+        for i in range(T-1):
+            transition_list.append({
+               'obs': traj['obs'][i],
+               'acts': traj['acts'][i],
+               'next_obs': traj['obs'][i+1],
+               'rews': traj['rews'][i],
+               'dones': (i==T-2),
+               'infos': infos[i],
+            })
+    return transition_list
+
+
+def make_sample_until(min_timesteps: Optional[int]=None, min_episodes: Optional[int]=None):
+    """Returns a termination condition sampling for a number of timesteps and episodes.
+    
+    Args:
+        min_timesteps: Sampling will not stop until there are at least this many
+            timesteps.
+        min_episodes: Sampling will not stop until there are at least this many
+            episodes.
+
+    Returns:
+        A termination condition which given a list of trajectories returns true if the condition is met.
+
+    Raises:
+        ValueError if neither of n_timesteps and n_episodes are set, or if either are
+            non-positive.
+    """
+    if min_timesteps is None and min_episodes is None:
+        raise ValueError(
+            "At least one of min_timesteps and min_episodes needs to be non-None"
+        )
+
+    conditions = []
+    if min_timesteps is not None:
+        if min_timesteps <= 0:
+            raise ValueError(
+                f"min_timesteps={min_timesteps} if provided must be positive"
+            )
+
+        def timestep_cond(trajectories):
+            if len(trajectories) == 0:
+                return False
+            timesteps = sum(len(t['obs']) - 1 for t in trajectories)
+            return timesteps >= min_timesteps   
+        conditions.append(timestep_cond)
+
+    if min_episodes is not None:
+        if min_episodes <= 0:
+            raise ValueError(
+                f"min_episodes={min_episodes} if provided must be positive"
+            )
+        conditions.append(lambda trajectories: len(trajectories) >= min_episodes)
+
+    def sample_until(trajs: List[dict]) -> bool:
+        for cond in conditions:
+            if not cond(trajs):
+                return False
+        return True
+
+    return sample_until
+
+def rollout_and_save(
+    path: str,
+    policy,
+    env,
+    sample_until,
+    *,
+    exclude_infos: bool = True,
+    **kwargs,
+) -> None:
+    """Generate policy rollouts and save them to a pickled list of trajectories.
+
+    The `.infos` field of each Trajectory is set to `None` to save space.
+
+    Args:
+      path: Rollouts are saved to this path.
+      policy: a stable_baselines3 policy or algorithm trained on the gym environment
+      env: The environment to interact with.
+      sample_until: End condition for rollout sampling.
+      exclude_infos: If True, then exclude `infos` from pickle by setting
+        this field to None. Excluding `infos` can save a lot of space during
+        pickles.
+      **kwargs: Passed through to `generate_trajectories`.
+    """
+    trajs = generate_trajectories(policy, env, sample_until, **kwargs)
+    if exclude_infos:
+        [traj.update(infos=None) for traj in trajs]
+    save(path, trajs)
+
+def save(path: str, trajectories: List[dict]) -> None:
+    """Save a sequence of Trajectories to disk.
+
+    Args:
+        path: Trajectories are saved to this path.
+        trajectories: The trajectories to save.
+    """
+    p = pathlib.Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = f"{path}.tmp"
+    with open(tmp_path, "wb") as f:
+        pickle.dump(trajectories, f)
+    # Ensure atomic write
+    os.replace(tmp_path, path)
+    logging.info(f"Dumped demonstrations to {path}.")