add video saving and uploading support to `train_*` scripts

[yawen-d]

Description

Closes #523.

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

Testing

• [x] Add video_saving tests in tests/scripts/test_scripts.py
• [x] Add video uploading test in tests/util/test_wb_logger.py

[codecov[bot]]

Codecov Report

Merging #524 (9b9ea3d) into master (de36306) will decrease coverage by 0.02%. The diff coverage is 100.00%.

:exclamation: Current head 9b9ea3d differs from pull request most recent head 5ca705f. Consider uploading reports for the commit 5ca705f to get more accurate results

@@            Coverage Diff             @@
##           master     #524      +/-   ##
==========================================
- Coverage   96.95%   96.93%   -0.03%     
==========================================
  Files          84       84              
  Lines        7460     7369      -91     
==========================================
- Hits         7233     7143      -90     
+ Misses        227      226       -1     

Impacted Files	Coverage Δ
src/imitation/algorithms/preference_comparisons.py	98.98% <100.00%> (-0.19%)	:arrow_down:
src/imitation/scripts/common/common.py	97.22% <100.00%> (ø)
src/imitation/scripts/common/train.py	100.00% <100.00%> (ø)
...ion/scripts/config/train_preference_comparisons.py	84.72% <100.00%> (-0.62%)	:arrow_down:
src/imitation/scripts/train_adversarial.py	96.29% <100.00%> (+1.62%)	:arrow_up:
src/imitation/scripts/train_imitation.py	94.11% <100.00%> (+0.17%)	:arrow_up:
.../imitation/scripts/train_preference_comparisons.py	98.38% <100.00%> (+0.02%)	:arrow_up:
src/imitation/scripts/train_rl.py	100.00% <100.00%> (ø)
src/imitation/util/logger.py	100.00% <100.00%> (ø)
src/imitation/util/video_wrapper.py	100.00% <100.00%> (ø)
... and 10 more

:mega: We’re building smart automated test selection to slash your CI/CD build times. Learn more

[AdamGleave]

I'm still a bit backlogged, @Rocamonde could you review this please?

[AdamGleave]

Closing in favor of https://github.com/HumanCompatibleAI/imitation/pull/597