Pytorch implementation of World Models