Gymnasium Single Frame Render with TorchRL

reinforcement-learning
1

I can do the following with Stable-Baselines3, but unsure how to do it with TorchRL. I simply want a single frame image to be saved off, not a full rollout video. I could not find a solution in the TorchRL docs. Please let me know if I am missing something. Thank you!

# initial conditions image
img = env.render()
images = wandb.Image(img, caption=f"Initial Condition State for Seed {env_seed}")
wandb.log({"Initial Conditions": images})
2

If you have a GymEnv you can do GymEnv.whatever and if whatever is not registered in the GymEnv class you will get it from the base env (ie, your gym env)

from torchrl.envs import GymEnv
env = GymEnv("Pendulum-v1")
env.reset()
print(env.render())

Hope that helps!
(if you want the rendered frames just create the env with from_pixels=True, pixels_only=False and if you want just one frame, well call env.step(data) - but that will make a step, not return just one rendered image :wink: )