[Question] Suport "native" environments

[martinitus]

❓ Question

Hey SB3 team!

I was experimenting the last days with rewriting my working custom python environment in rust and exposing it to python via pyo3 and maturin.

Doing so I came about a couple of issues, most of them I could solve. However, I needed to disable one or two isintance(env,Env) checks within my local stable_baselines3 package, as inheritance from non built-ins seems not to be possible in pyo3 at the moment. I also could not find a solution to patch sb3 or "inject" the base type into my custom environment class imported from rust.

I therefore have two questions:

1. Is there interest in this topic? Currently all i have is mixed into my own codebase, but i could produce a minimal standalone repository demonstrating how this would be feasible.
2. Could you imagine a contribution that includes a mechanism to disables certain "safety checks" (such as the isintance(env,Env) checks into which advanced users could opt in - allowing them to use environments that come from native code like C++ or rust?

Curious about your thoughts!

Remarks:

• I am aware, that I could just write a thin wrapper in python that capsules my rust environment, but it would just be nice if i wouldn't need that :-)
• For complicated environments, the observation and step function may also cost some computation/performance, so implementing them in native code seems to make sense.
• For applications where there already is a lot of rust code, defining the environment directly in rust would kinda feel natural to me.

Finally a random screenshot to hopefully ignite some curiosity ;-)

Checklist

• I have checked that there is no similar issue in the repo
• I have read the documentation
• If code there is, it is minimal and working
• If code there is, it is formatted using the markdown code blocks for both code and stack traces.

[martinitus]

Oh - one more issue i just stumbled across while testing with SubprocVecEnv ist pickling/unpickling of types defined via pyo3 seems also not possible at the moment. 🤷🏼

Update:
Nevermind this is also solvable, as you don't send the actual environments to the subprocesses, but just their factory function. One can avoid the pickling issues with some proper python import handling.

[araffin]

Hello,

yes, this easiest is probably to write a wrapper:

import gymnasium as gym
from stable_baselines3 import PPO, SAC


class CustomEnv(gym.Wrapper):
    def __init__(self, env: gym.Env) -> None:
        self.env = env
        self.observation_space = self.env.observation_space
        self.action_space = self.env.action_space
        self.metadata = {"render_modes": []}

    def reset(self, *, seed=None, options=None):
        return self.env.reset()

    def step(self, action):
        return self.env.step(action)


env = gym.make("Pendulum-v1")

PPO("MlpPolicy", CustomEnv(env)).learn(1)
SAC("MlpPolicy", CustomEnv(env)).learn(200)

and if you do things in rust, I would recommend doing the parallelization in rust (to avoid GIL limitatins), and so write a VecEnvWrapper.

You have multiple examples of those in the doc: https://stable-baselines3.readthedocs.io/en/master/guide/examples.html#sb3-with-isaac-lab-brax-procgen-envpool

[martinitus]

Ahh, thanks for pointing me into the direction of VecEnv. I just started digging into how this is actually done within SB3.

I will leave a couple of remarks here, where I found things slightly confusing and see a chance for small improvements. Lots of them would be breaking changes I think, but maybe my thoughts help some other interested reader who follows :-)

First Point:

My first understanding was, that SB3 is completely built around the gymnasium / gym env class. But if I understand that correctly now - that's actually not the case! Instead everything is built around the VecEnv (and maybe VecEnvWrapper) class. I think my first idea came mostly from the name of the type alias GymEnv = Union[gym.Env, "VecEnv"] that is used in most of the user facing API (model constructor, etc).

So one small but significant change would be to simply rename GymEnv to something that does not include Gym :-)

Another possibility that would go a bit further would be to remove the type alias entirely and always use use VecEnv in the user facing and internal API. That would unclutter a lot of the internals where the current code has to deal with the case that the env might not be a VecEnv. Instead it would put the responsibility of appropriately wrapping a gymnasium environment onto the user.
Maybe it would also make the documentation for how to write custom environments clearer:

• either write a gymnasium environment and use one of the wrappers,
• or if no compatibility with gymasium is required a user could also directly implement a VecEnv (as you proposed above)

[martinitus]

Another Point:

Yes, I could implement the parallelization in rust, in fact I'll probably go with that! But I could also free the GIL while the rust implementation runs the step code. That would allow using a multi-threaded VecEnv implementation in python that could truly run multiple steps in parallel in different python threads.

[araffin]

Instead everything is built around the VecEnv

yes, from the developer guide: "All the environments handled internally are assumed to be VecEnv (gym.Env are automatically wrapped)." (https://stable-baselines3.readthedocs.io/en/master/guide/developer.html)

small but significant change would be to simply rename GymEnv

In the past, the two APIs used to be very similar, but that's not the case anymore: https://stable-baselines3.readthedocs.io/en/master/guide/vec_envs.html#vecenv-api-vs-gym-api

nstead it would put the responsibility of appropriately wrapping a gymnasium environment onto the user.

That used to be the case, but we chose to do repetitive things that the user would have to do instead (like changing from channel last to channel first, or converting gym env to VecEnv). As a user, it makes life easier.

Maybe it would also make the documentation for how to write custom environments clearer:

If something is not clear, we welcome contribution to improve our doc =)

[martinitus]

Hah, thank you to pointing me to some of those links, and some of the insights into development history :-)

I will see if i can find spots where I could (subjectively) improve the documentation.

I can also report, that my rust experiments went very well. I got a speedup from ~100 fps to ~50000 fps (depending on the policy, as some actions require more processing than others). To be fair, my first python based implementation was horrible - on the rust side i actuall put some thought into optimizing the implementation, avoiding allocations and so forth.

The parallelization works beautifully with rayon. Also my rust based vec-env pre-allocates the observation and reward buffer (implemented via rust ndarray) for its nested environments. The poplulated rust ndarrays are then be exposed with zero copy as numpy arrays back to python.

Overall, i now have a solid 100% load on 16 cores where I had like 20 % on some cores with the SupProcVecEnv implementation.

Two more remarks:

If instead of using "classical inheritance" via abc.ABC SB3 would use structural subtyping (via Protocols), I could save myself some wrappers and boilerplate empty methods. Obviously that would require some changes to all the isintance(v, VecEnv) calls, but i think that could be achieved in a backwards compatible way and would give more options to users :)

I still have to dig into the rollout buffers, but in principle it should be possible to have the rust side write directly into preallocated torch buffers. that would eliminate another round of copy operations.

[martinitus]

Also I'll close this as complete. Thanks for your help :)