Cython in RLlib to speed up gym.env processing?

How severe does this issue affect your experience of using Ray?

  • Medium: It contributes to significant difficulty to complete my task, but I can work around it.

Hello Ray users,

I would like to raise the hand for experience exchange: Has anyone ever tried to use Cython together with RLllib?

As I am struggling currently a bit with the memory footprint and computational runtime of the logic inside my custom gymnasium environment, I am looking for every option to improve that. While for sure potential is lying in data structures, alogrithm logic in step() method and similar, I believe that Cython could reduce the computational burden of the environment logic during PPO training.

Would be good to think out-of-the-box and check for pre-requisities for a successful application.

@PhilippWillms, great question! So, in general the logic inside of the environment is independent of RLlib, i.e. you can use whatever language fits your needs. Many simulators are written in C/C++.
When it comes then to connect RLlib with the environment, it needs at best a gymnasium.Env that wraps somehow your code (like this environment vectorization can be achieved with ease, see num_envs_per_worker).

If such a gymnasium wrapping is not possible for you, we still have the ExternalEnv and the PolicyClient methods that can deal with environments that run in their very own process and possess some kind of API.

In case of Cython the first option should be possible to implement and I would opt for it from my experience. Only switch to option two, if there is no other way to achieve what you want as this option will come with its own hurdles (e.g. when it comes to scaling, fault tolerance, etc.)

Before I go over and try to rewrite my code, I would however profile my code to see where exactly the bottleneck lays and if I can overcome it with a faster implementation in Cython.