babyai
babyai copied to clipboard
Use new gym API, gym's vectorized environments and fix model backprop
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.