python -m render_options --model_name='gail_options_image_mid_wcollision' --env='NRasterizedRoute' --options=True --width=36 --height=36 --m_per_px=2 --agent=50 --stop_on_collision=False

import torch, os
from src.data import load_experts
folder = 'expert_data/DR_USA_Roundabout_FT/track0000'
single_agent = os.path.join(folder, 'expert.pkl')
multi_agent = os.path.join(folder,'joint_expert_states.pt')
multi_agent_actions = os.path.join(folder,'joint_expert_actions.pt')
demonstrations = load_experts([single_agent], flatten=False)
demonstrations[0].__dict__.keys()
len(demonstrations[0].obs)
single_agent_lengths = [len(demonstration.obs) for demonstration in demonstrations]
states = torch.load(multi_agent)
actions = torch.load(multi_agent_actions)
multi_agent_lengths = [sum(~torch.isnan(states[:,i,0])).item() for i in range(states.shape[1])]

single_agent_actions = [demonstration.acts for demonstration in demonstrations]
multi_agent_actions = [actions[~torch.isnan(actions[:,i,0])] for i in range(actions.shape[1])]

import pickle
with open(single_agent, "rb") as f:
    new_trajectories = pickle.load(f)