def recode_and_save_trajectories(traj_or_policy_file, save_loc, traj_len, seed,
args):
if "skills" in traj_or_policy_file:
trajs = convert_trajs(traj_or_policy_file, traj_len, *args)
else:
trajs = rollout_policy(traj_or_policy_file, traj_len, seed, *args)
# assert len(trajs) == 1
for traj in trajs:
assert traj.obs.shape[0] == traj_len + 1
assert traj.acts.shape[0] == traj_len
trajs = [dataclasses.replace(traj, infos=None) for traj in trajs]
types.save(save_loc, trajs)
def rollout_and_save(
path: str,
policy: Union[BaseAlgorithm, BasePolicy],
venv: VecEnv,
sample_until: GenTrajTerminationFn,
*,
unwrap: bool = True,
exclude_infos: bool = True,
verbose: 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.
venv: The vectorized environments.
sample_until: End condition for rollout sampling.
unwrap: If True, then save original observations and rewards (instead of
potentially wrapped observations and rewards) by calling
`unwrap_traj()`.
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.
verbose: If True, then print out rollout stats before saving.
deterministic_policy: Argument from `generate_trajectories`.
"""
trajs = generate_trajectories(policy, venv, sample_until, **kwargs)
if unwrap:
trajs = [unwrap_traj(traj) for traj in trajs]
if exclude_infos:
trajs = [dataclasses.replace(traj, infos=None) for traj in trajs]
if verbose:
stats = rollout_stats(trajs)
logging.info(f"Rollout stats: {stats}")
types.save(path, trajs)
def test_save_trajectories(self, trajectory_rew, use_chdir, tmpdir):
"""Check that trajectories are properly saved."""
if use_chdir:
# Test no relative path without directory edge-case.
chdir_context = pushd(tmpdir)
save_dir = ""
else:
chdir_context = contextlib.nullcontext()
save_dir = tmpdir
trajs = [trajectory_rew]
save_path = pathlib.Path(save_dir, "trajs.pkl")
with chdir_context:
types.save(str(save_path), trajs)
with open(save_path, "rb") as f:
loaded_trajs = pickle.load(f)
assert len(trajs) == len(loaded_trajs)
for t1, t2 in zip(trajs, loaded_trajs):
d1, d2 = dataclasses.asdict(t1), dataclasses.asdict(t2)
assert d1.keys() == d2.keys()
for k, v in d1.items():
assert np.array_equal(v, d2[k])