import os import numpy as np import gymnasium as gym import matplotlib.pyplot as plt from gymnasium.vector import SyncVectorEnv import robosuite as suite from robosuite.wrappers import GymWrapper def make_env(seed: int, camera_names: list[str]): def _thunk(): # Create a robosuite env (example: Lift) rs_env = suite.make( env_name="Lift", robots="Panda", has_renderer=False, has_offscreen_renderer=True, use_camera_obs=True, camera_names=camera_names, horizon=400, control_freq=20, use_object_obs=False, camera_heights=128, camera_widths=128, ) # Wrap to Gym-style API env = GymWrapper(rs_env, flatten_obs=False) # returns obs, reward, done, info (classic gym style) env.metadata = {} # hack since robosuite sets metadata to None env.reset(seed=seed) return env return _thunk # Build a 2-env vector n_envs = 2 camera_names = ["agentview", "robot0_eye_in_hand"] venv = SyncVectorEnv([make_env(i, camera_names) for i in range(n_envs)], autoreset_mode=gym.vector.AutoresetMode.SAME_STEP) obs, info = venv.reset() for i in range(10): actions = venv.action_space.sample() # samples a batch of 2 actions automatically for vector envs obs, rewards, terminated, truncated, infos = venv.step(actions) image = np.concatenate([obs[camera_names[0] + "_image"][j] for j in range(n_envs)], axis=1) # save each image to a file in /tmp/ image = image.astype(np.uint8) plt.imshow(image) filename = f"/tmp/image_{i}.png" os.makedirs(os.path.dirname(filename), exist_ok=True) plt.savefig(filename) print(f"Saved image to {filename}") plt.close()