Problem
- I personally found logging videos during training is really useful as another dimension of explaining experiment results.
- Concretely, this issue advocates for adding support for saving videos of policies on a environment for evaluation during and after training, including
scripts.train_rl, scripts.train_preference_comparisons, scripts.train_adversarial and scripts.train_bc.
- Also, it would be nice to add support for uploading saved videos to Weights & Biases during and after training.
Solution
- Write a
record_and_save_video() function in imitation.util.video_wrapper that takes in a policy, eval_venv, and a logger to save the video of a policy evaluated on an environment to a designated path.
def record_and_save_video(
output_dir: str,
policy: policies.BasePolicy,
eval_venv: vec_env.VecEnv,
video_kwargs: Mapping[str, Any],
logger: Optional[sb_logger.Logger] = None,
) -> None:
...
- Upload the video to weights & biases within
WandbOutputFormat.write() by adding the following:
if key != "video":
self.wandb_module.log({key: value}, step=step)
else:
self.wandb_module.log({"video": self.wandb_module.Video(value)})
