Recommended way for parallel simulation compatibility with Gymnasium

[johnsutor]

I plan on creating a vector environment using Stable Baselines3, in conjunction with Genesis. Is there a recommended way for handling parallel training with Genesis parallel simulation?

[jgillick]

The Go2 example (github) is a pretty good demo and has an option for multiple parallel environments, but it's using rsl-rl, not Stable Baselines3.

Stable Baselines3 might not be ideal for this case, since it uses NumPy (vs pytorch or JAX) and all the vectorized work would remain on the CPU, versus pushing things to the GPU.

How parallelism works

I'm not exactly sure what part of parallel simulation and RL training you're trying to understand, but here's a high-level.

TL;DR: Everything works with arrays of arrays.

The observations returned by your environment's reset & step functions tell the training framework how many parallel environments are running. For example, if you have 100 parallel environments, your observation will be an array of 100 entries, each containing a list of observations for an environment. Then the training framework will call your step function with actions, which is an array of 100 entries, each with an array of actions that match the action space defined by your environment.

Assuming your action space is only for actuators, you can pass these actions directly to genesis to move the DOF actuators:

def step(self, actions):
    robot.control_dofs_position(actions)
    scene.step()

Genesis Forge Framework

I've been working on a framework that feels like Gymnasium meets Isaac Lab and uses the SKRL reinforcement library by default, which is GPU optimized and runs on a mac. I haven't published it yet, but here's my working example: https://github.com/jgillick/spider-bot/tree/main/simulate/genesis The framework I'm building is in the genesis_forge directory.

I'm not sure how much of this is portable to Stable Baselines3, but here's how it works with SKRL:

The environment will define the action and observation spaces.

class SpiderRobotEnv(GenesisEnv):
    # ...

    @property
    def observation_space(self):
        return spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(84,),
            dtype=np.float32,
        )

    @property
    def action_space(self):
        return spaces.Box(
            low=-1.0,
            high=1.0,
            shape=(24,),
            dtype=np.float32,
        )

    # ...

Then, in the training script, build the environment and wrap it for SKRL (you'll probably need to create a custom wrapper for Stable Baselines3).

NUM_ENVS = 100
env = SpiderRobotEnv(num_envs=NUM_ENVS)
env.build()

skrl_env = SkrlEnvWapper(env)

or

NUM_ENVS = 100
env = SpiderRobotEnv(num_envs=NUM_ENVS)
env.build()

skrl_env = create_skrl_env(env)

Then train!

runner = Runner(skrl_env, cfg)
runner.run("train")

[johnsutor]

That looks awesome! I have a custom class that creates parallel classes under the hood, but was wondering if there is a proper Genesis way to achieve this. For example, ManiSkill explicitly has a Vector Env for Stable Baselines 3 support.

[jgillick]

From what I have seen, there's no standard Genesis way to do this. The framework is so new that it largely provides the simulations, sensors, and utilities, but lacks many environmental best practices, training classes, or wrappers. That said, I could be wrong; I've only been working with it for a few weeks.

But that's also the inspiration around the framework I'm creating. I'd like it to be easier to get up and running with less boilerplate work.