ragged-buffer VS enn-zoo

We are also releasing enn-trainer, a PPO implementation that takes full advantage of the Entity Gym interface. Variable-length observations are efficiently processed using ragged sample buffers and a general ragged batch transformer implementation that can be applied to any Entity Gym environment. With many performance optimizations still missing, enn-trainer can already reach a throughput of 10s of thousands of samples per second per GPU when it is not bottlenecked by stepping the environment. More typically, environments implemented in Python reach thousands of samples per second, but can share a single GPU between multiple concurrent training runs.

One of Rust's many strengths is that it can be seamlessly integrated with Python and speed up critical code sections. I recently wrote a small library with an efficient ragged array datatype, and I figured it would make for a good example of how to set up a Rust Python package with PyO3 and maturin that interoperates with numpy. There are a lot of little details that took me quite a while to figure out:

While we have not had the time to run careful experiments that meet our standard of rigor, preliminary evaluations on a number of standard RL environments have looked quite promising compared to baselines with vision-based policies. Entity Gym’s flexible API makes it comparatively effortless to interface with many kinds of environments that would be quite cumbersome to integrate with existing RL frameworks and I’m quite excited to see what happens when Entity Gym is applied to other interesting tasks. If you want to give this a shot, our tutorials for implementing Entity Gym environments and training policies with enn-trainer will have you up and running in no time.