Best way to add a GPU parallelized environment/framework?

[StoneT2000]

Been hearing about TorchRL a lot recently (hope to see some of you at ICLR!). Recently my lab and I have been working on / released in beta a GPU parallelized robotics simulation called ManiSkill (think of it as Isaac but rendering is also parallelized). I'm looking to add support for TorchRL for ManiSkill but had a bit of trouble trying to figure out how to do it without making the RL run too slow and to fully leverage the GPU parallelization on both the state data as well as rendering (for image observations and evaluation videos).

PPO Training and Env Creation Setup

For example of something I didn't quite follow (I assume I implemented something wrong), I am trying to implement PPO baselines for my environment following this: https://github.com/pytorch/rl/tree/7348af330c466bcd6ab092b9ef1cc7f3e9631475/sota-implementations/ppo

To do this I have so far made a ManiSkillWrapper class that inherits GymWrapper and it seems like it works (ManiSkill is mostly gymnasium compliant with the exception that with num_envs=1 or more than 1 the outputs are always batched torch tensors)

To test I just checked this worked

import gymnasium
import mani_skill
from wrapper import ManiSkillWrapper
env = ManiSkillWrapper(gymnasium.make("PushCube-v1", obs_mode="rgbd", num_envs=4))
obs = env.reset(seed=0)
td = env.rand_step()

and td has all the data that I expected for observation, terminations etc. and specs

However looking at the PPO implementation for mujuco it seems to be single env only? Wasn't sure how to best adapt that code.

Looking at the Atari one it uses the ParallelEnv class which seemed promising but I am not sure if ParallelEnv class is the one I should try and inherit as I found that IsaacGym has a wrapper implementation but inherits GymWrapper.

Evaluation with parallel rendered videos

Our simulator supports rendering e.g. 16 videos quickly on a single process with a GPU, I was wondering if there was anything in TorchRL that might not let this work easily with the existing VideoRecorder wrapper. If not our repo provides a video recording wrapper but I would like to try and stick to as many pure torch rl modules as possible.

Any advice would be greatly appreciated!

[vmoens]

Hey!
I hope you heard good things about the lib!
If you want to check that the env works ok you can run check_env_specs which should do the job.

Other than that, you can indicate how many envs are in your env by passing a batch_size arg to __init__. You'll find examples of this in jumanji and Brax envs.
From there, each spec will need to have a leading dim that matches the batch size.

Let us know how things go. If you want torchrl to support your lib from within we can think about that, but you can also build things independently which will offer you more flexibility on the long run.

Feel free to join the discord and ask questions there if you want more interactions with the community (GitHub discussions are also fine!)

[StoneT2000]

I'll follow up on Discord later thanks!