babyai icon indicating copy to clipboard operation
babyai copied to clipboard

Use new gym API, gym's vectorized environments and fix model backprop

Open saleml opened this issue 2 years ago • 0 comments

This PR follows this one in Minigrid and should at the very leaast await gym 0.25.0. In the meantime, environment and setup files use versions of gym and gym-minigrid that are not official releases.

The goal of the PR is to follow gym's new API, and use gym's vectorized environments instead of babyai's ParallelEnv for RL (both training and evaluation).

Also, to make some parts of the code used for IL and RL more readable, I replaced some convoluted torch.transpose + torch.reshape/ torch.view + torch.unsqueeze with einops.rearrange.

Lastly, a recent change in torch.autograd made a part of the code in model.py break, because of in-place operations. This should hopefully be fixed now.

P.S. I had a weird bug with train_rl on MacOS when using multiple processes (which didn't occur with Linux, or with just one process). So although the changes in train_rl seem huge, they are actually minimal as I just indented most of the file under a if __name__ = '__main__' as suggested here, in order to fix the bug.

saleml avatar Jul 10 '22 23:07 saleml