Would like to be able to run `eval_policy` for atari environments

[dfilan]

Problem

I would like to be able to run imitation.scripts.eval_policy on an atari environment, to roll out a policy, so that I can then do clever things with the rollouts.

If I just set environment.gym_id to "AsteroidsNoFrameskip-v4", the environment is not wrapped in an appropriate form: it first needs to be put into an AtariWrapper from SB3, and then either a FrameStack or a VecFrameStack, so that each obs has 4 frames, in order to be fed into the policy that will be downloaded from HuggingFace.

The problem is that I can't write a config to use these wrappers, since the script doesn't allow them to be manually set.

Solution

Let me set the wrappers that I put around the environment.

Possible alternative solutions

Automatically detect if the environment is an Atari game, and apply wrappers appropriately.

[dfilan]

To be honest, this is probably low priority, since the only Atari policy we have saved in the HuggingFace repo is for Asteroids (which doesn't get a particularly good score anyway), and home-trained policies could manually apply wrappers themselves - but it is a bit sad that this is so hard, given that this repo and seals nominally both support atari environments and CNNs.