Fix mypy type erorrs in skipped files

[AdamGleave]

Bug description

Some files are failing mypy type checks (but passing on pytype) and being skipped by EXCLUDE_MYPY in .circleci/config.yml introduced in https://github.com/HumanCompatibleAI/imitation/pull/534 We should fix these errors, by ignoring spurious errors and for true-positive errors refactoring the code or adding assert checks where appropriate.

Steps to reproduce

Run mypy on the excluded files

Environment