adding functions to save joint expert states and actions for repeated use in metrics, adding option to flatten loaded trajectories, adding class to not shuffle trajectories when saving experts to make sure it lines up with the joint states. checked that it does

This commit is contained in:
Arec
2021-11-29 14:09:57 -08:00
parent 284d4af637
commit d34fa5774d
3 changed files with 87 additions and 11 deletions

View File

@@ -1 +1,22 @@
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
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)