Issues with Ray RLlib and Unity Environment Integration: Episode Management and Unity_3d_local_env

Hello everyone,

I’ve been working on integrating Unity ML-Agents with Ray RLlib and have come across a few issues with the implementation of the ray unity3d environment wrapper. Specifically, I’m finding the way it handles episode management for individual agents quite strange on each step.

From my understanding, the wrapper doesn’t seem to reset each agent’s episode based on when the agent reaches its individual end state. Instead, it waits until the global step horizon is reached before resetting all agents. This seems quite different from how the ML-Agents framework itself handles episode termination, where each agent’s episode ends independently.

Is anyone else having experience with this? I was really hoping to use Ray RLlib for its scalability, but this implementation feels quite disconnected from how ML-Agents operates. Has anyone found a way to make the RLlib environment behave more closely to ML-Agents’ episode management? Any advice or insights would be appreciated!