Merge pull request #537 from ARISE-Initiative/maxfr

add max frame rate feature to relevant demo scripts

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...")

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
+    )

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

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)

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,12 +38,19 @@
 
         # 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)
             if done:
                 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()

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)

robosuite/demos/demo_renderers.py

@@ -1,12 +1,15 @@
 import argparse
 import json
+import time
 
 import numpy as np
 
 import robosuite as suite
 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.")

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)