dppo/script/test_robomimic_render.py
2024-09-03 21:03:27 -04:00

37 lines
976 B
Python

"""
Test Robomimic rendering, no GUI
"""
import os
import time
from gym import spaces
import robosuite as suite
os.environ["MUJOCO_GL"] = "egl"
if __name__ == "__main__":
env = suite.make(
env_name="TwoArmTransport",
robots=["Panda", "Panda"],
has_renderer=False,
has_offscreen_renderer=True,
use_camera_obs=True,
camera_heights=96,
camera_widths=96,
camera_names="shouldercamera0",
render_gpu_device_id=0,
horizon=20,
)
obs, done = env.reset(), False
print("Finished resetting!")
low, high = env.action_spec
action_space = spaces.Box(low=low, high=high)
steps, time_stamp = 0, time.time()
while True:
while not done:
obs, reward, done, info = env.step(action_space.sample())
steps += 1
obs, done = env.reset(), False
print(f"FPS: {steps / (time.time() - time_stamp)}")
steps, time_stamp = 0, time.time()