Pufferlib already had a pretty good model before: https://puffer.ai/ocean.html?env=tetris