diff --git a/robosuite/demos/demo_collect_and_playback_data.py b/robosuite/demos/demo_collect_and_playback_data.py index 1bb71c316a..ef44c4b11d 100644 --- a/robosuite/demos/demo_collect_and_playback_data.py +++ b/robosuite/demos/demo_collect_and_playback_data.py @@ -7,6 +7,7 @@ import argparse import os +import time from glob import glob import numpy as np @@ -15,7 +16,7 @@ from robosuite.wrappers import DataCollectionWrapper -def collect_random_trajectory(env, timesteps=1000): +def collect_random_trajectory(env, timesteps=1000, max_fr=None): """Run a random policy to collect trajectories. The rollout trajectory is saved to files in npz format. @@ -24,18 +25,27 @@ def collect_random_trajectory(env, timesteps=1000): Args: env (MujocoEnv): environment instance to collect trajectories from timesteps(int): how many environment timesteps to run for a given trajectory + max_fr (int): if specified, pause the simulation whenever simulation runs faster than max_fr """ env.reset() dof = env.action_dim for t in range(timesteps): + start = time.time() action = np.random.randn(dof) env.step(action) env.render() if t % 100 == 0: print(t) + # limit frame rate if necessary + if max_fr is not None: + elapsed = time.time() - start + diff = 1 / max_fr - elapsed + if diff > 0: + time.sleep(diff) + def playback_trajectory(env, ep_dir): """Playback data from an episode. @@ -74,6 +84,12 @@ def playback_trajectory(env, ep_dir): parser.add_argument("--robots", nargs="+", type=str, default="Panda", help="Which robot(s) to use in the env") parser.add_argument("--directory", type=str, default="/tmp/") parser.add_argument("--timesteps", type=int, default=1000) + parser.add_argument( + "--max_fr", + default=20, + type=int, + help="Sleep when simluation runs faster than specified frame rate; 20 fps is real time.", + ) args = parser.parse_args() # create original environment @@ -98,7 +114,7 @@ def playback_trajectory(env, ep_dir): # collect some data print("Collecting some random data...") - collect_random_trajectory(env, timesteps=args.timesteps) + collect_random_trajectory(env, timesteps=args.timesteps, max_fr=args.max_fr) # playback some data _ = input("Press any key to begin the playback...") diff --git a/robosuite/demos/demo_composite_robot.py b/robosuite/demos/demo_composite_robot.py index a178032276..9fb14c8139 100644 --- a/robosuite/demos/demo_composite_robot.py +++ b/robosuite/demos/demo_composite_robot.py @@ -1,4 +1,5 @@ import argparse +import time from typing import Dict, List, Union import numpy as np @@ -14,6 +15,7 @@ def create_and_test_env( robots: Union[str, List[str]], controller_config: Dict, headless: bool = False, + max_fr: int = None, ): config = { @@ -38,9 +40,18 @@ def create_and_test_env( # Runs a few steps of the simulation as a sanity check for i in range(100): + start = time.time() + action = np.random.uniform(low, high) obs, reward, done, _ = env.step(action) + # limit frame rate if necessary + if max_fr is not None: + elapsed = time.time() - start + diff = 1 / max_fr - elapsed + if diff > 0: + time.sleep(diff) + env.close() @@ -53,6 +64,12 @@ def create_and_test_env( parser.add_argument("--grippers", nargs="+", type=str, default=["PandaGripper"]) parser.add_argument("--env", type=str, default="Lift") parser.add_argument("--headless", action="store_true") + parser.add_argument( + "--max_fr", + default=20, + type=int, + help="Sleep when simluation runs faster than specified frame rate; 20 fps is real time.", + ) args = parser.parse_args() @@ -63,4 +80,6 @@ def create_and_test_env( create_composite_robot(name, base=args.base, robot=args.robot, grippers=args.grippers) controller_config = load_composite_controller_config(controller="BASIC", robot=name) - create_and_test_env(env="Lift", robots=name, controller_config=controller_config, headless=args.headless) + create_and_test_env( + env="Lift", robots=name, controller_config=controller_config, headless=args.headless, max_fr=args.max_fr + ) diff --git a/robosuite/demos/demo_control.py b/robosuite/demos/demo_control.py index 32d0dafd01..8d7fab6cdb 100644 --- a/robosuite/demos/demo_control.py +++ b/robosuite/demos/demo_control.py @@ -44,11 +44,14 @@ """ +import time from typing import Dict import robosuite as suite from robosuite.utils.input_utils import * +MAX_FR = 25 # max frame rate for running simluation + if __name__ == "__main__": # Create dict to hold options that will be passed to env creation call @@ -143,6 +146,7 @@ while count < num_test_steps: action = neutral.copy() for i in range(steps_per_action): + start = time.time() if controller_name in {"IK_POSE", "OSC_POSE"} and count > 2: # Set this value to be the scaled axis angle vector vec = np.zeros(3) @@ -153,10 +157,23 @@ total_action = np.tile(action, n) env.step(total_action) env.render() + + # limit frame rate if necessary + elapsed = time.time() - start + diff = 1 / MAX_FR - elapsed + if diff > 0: + time.sleep(diff) for i in range(steps_per_rest): + start = time.time() total_action = np.tile(neutral, n) env.step(total_action) env.render() + + # limit frame rate if necessary + elapsed = time.time() - start + diff = 1 / MAX_FR - elapsed + if diff > 0: + time.sleep(diff) count += 1 # Shut down this env before starting the next test diff --git a/robosuite/demos/demo_device_control.py b/robosuite/demos/demo_device_control.py index 1788cc2c14..224ed6882e 100644 --- a/robosuite/demos/demo_device_control.py +++ b/robosuite/demos/demo_device_control.py @@ -89,6 +89,7 @@ """ import argparse +import time import numpy as np @@ -111,6 +112,12 @@ parser.add_argument("--device", type=str, default="keyboard") parser.add_argument("--pos-sensitivity", type=float, default=1.0, help="How much to scale position user inputs") parser.add_argument("--rot-sensitivity", type=float, default=1.0, help="How much to scale rotation user inputs") + parser.add_argument( + "--max_fr", + default=20, + type=int, + help="Sleep when simluation runs faster than specified frame rate; 20 fps is real time.", + ) args = parser.parse_args() # Get controller config @@ -189,6 +196,8 @@ gripper_key = gripper_key[0] while True: + start = time.time() + # Set active robot active_robot = env.robots[0] if args.config == "bimanual" else env.robots[args.arm == "left"] @@ -236,3 +245,10 @@ # Step through the simulation and render obs, reward, done, info = env.step(action) env.render() + + # limit frame rate if necessary + if args.max_fr is not None: + elapsed = time.time() - start + diff = 1 / args.max_fr - elapsed + if diff > 0: + time.sleep(diff) diff --git a/robosuite/demos/demo_gripper_selection.py b/robosuite/demos/demo_gripper_selection.py index 1f87a800b8..1da85c9ddd 100644 --- a/robosuite/demos/demo_gripper_selection.py +++ b/robosuite/demos/demo_gripper_selection.py @@ -2,11 +2,15 @@ This script shows you how to select gripper for an environment. This is controlled by gripper_type keyword argument. """ +import time + import numpy as np import robosuite as suite from robosuite import ALL_GRIPPERS +MAX_FR = 25 # max frame rate for running simluation + if __name__ == "__main__": for gripper in ALL_GRIPPERS: @@ -34,6 +38,7 @@ # Run random policy for t in range(100): + start = time.time() env.render() action = np.random.uniform(low, high) observation, reward, done, info = env.step(action) @@ -41,5 +46,11 @@ print("Episode finished after {} timesteps".format(t + 1)) break + # limit frame rate if necessary + elapsed = time.time() - start + diff = 1 / MAX_FR - elapsed + if diff > 0: + time.sleep(diff) + # close window env.close() diff --git a/robosuite/demos/demo_random_action.py b/robosuite/demos/demo_random_action.py index 91993bab3e..6ce751c1c9 100644 --- a/robosuite/demos/demo_random_action.py +++ b/robosuite/demos/demo_random_action.py @@ -1,5 +1,9 @@ +import time + from robosuite.utils.input_utils import * +MAX_FR = 25 # max frame rate for running simluation + if __name__ == "__main__": # Create dict to hold options that will be passed to env creation call @@ -53,6 +57,14 @@ # do visualization for i in range(10000): + start = time.time() + action = np.random.uniform(low, high) obs, reward, done, _ = env.step(action) env.render() + + # limit frame rate if necessary + elapsed = time.time() - start + diff = 1 / MAX_FR - elapsed + if diff > 0: + time.sleep(diff) diff --git a/robosuite/demos/demo_renderers.py b/robosuite/demos/demo_renderers.py index dc31a16e94..d4d0d351d4 100644 --- a/robosuite/demos/demo_renderers.py +++ b/robosuite/demos/demo_renderers.py @@ -1,5 +1,6 @@ import argparse import json +import time import numpy as np @@ -7,6 +8,8 @@ import robosuite.utils.transform_utils as T from robosuite.utils.input_utils import * +MAX_FR = 25 # max frame rate for running simluation + def str2bool(v): if v.lower() in ("yes", "true", "t", "y", "1"): @@ -102,9 +105,17 @@ def print_command(char, info): # do visualization for i in range(10000): + start = time.time() + action = np.random.uniform(low, high) obs, reward, done, _ = env.step(action) env.render() + # limit frame rate if necessary + elapsed = time.time() - start + diff = 1 / MAX_FR - elapsed + if diff > 0: + time.sleep(diff) + env.close_renderer() print("Done.") diff --git a/robosuite/scripts/collect_human_demonstrations.py b/robosuite/scripts/collect_human_demonstrations.py index a78ad2d639..2cd28732c3 100644 --- a/robosuite/scripts/collect_human_demonstrations.py +++ b/robosuite/scripts/collect_human_demonstrations.py @@ -20,7 +20,7 @@ from robosuite.wrappers import DataCollectionWrapper, VisualizationWrapper -def collect_human_trajectory(env, device, arm): +def collect_human_trajectory(env, device, arm, max_fr): """ Use the device (keyboard or SpaceNav 3D mouse) to collect a demonstration. The rollout trajectory is saved to files in npz format. @@ -30,6 +30,7 @@ def collect_human_trajectory(env, device, arm): env (MujocoEnv): environment to control device (Device): to receive controls from the device arms (str): which arm to control (eg bimanual) 'right' or 'left' + max_fr (int): if specified, pause the simulation whenever simulation runs faster than max_fr """ env.reset() @@ -53,6 +54,8 @@ def collect_human_trajectory(env, device, arm): # Loop until we get a reset from the input or the task completes while True: + start = time.time() + # Set active robot active_robot = env.robots[device.active_robot] @@ -103,6 +106,13 @@ def collect_human_trajectory(env, device, arm): else: task_completion_hold_count = -1 # null the counter if there's no success + # limit frame rate if necessary + if max_fr is not None: + elapsed = time.time() - start + diff = 1 / max_fr - elapsed + if diff > 0: + time.sleep(diff) + # cleanup for end of data collection episodes env.close() @@ -228,6 +238,12 @@ def gather_demonstrations_as_hdf5(directory, out_dir, env_info): default="mjviewer", help="Use Mujoco's builtin interactive viewer (mjviewer) or OpenCV viewer (mujoco)", ) + parser.add_argument( + "--max_fr", + default=20, + type=int, + help="Sleep when simluation runs faster than specified frame rate; 20 fps is real time.", + ) args = parser.parse_args() # Get controller config @@ -298,5 +314,5 @@ def gather_demonstrations_as_hdf5(directory, out_dir, env_info): # collect demonstrations while True: - collect_human_trajectory(env, device, args.arm) + collect_human_trajectory(env, device, args.arm, args.max_fr) gather_demonstrations_as_hdf5(tmp_directory, new_dir, env_info)