Add support for saving videos of policies on a environment for evaluation during and after training

[yawen-d]

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)})

[AdamGleave]

Assigning to @samuelarnesen to get PR https://github.com/HumanCompatibleAI/imitation/pull/524 over finish line once he starts

[Rocamonde]

Sounds good, feel free to request a review @samuelarnesen once you make some progress. I had already reviewed #524 before.