Merge pull request #548 from Abhiram824/demo_fixes

Demo fixes

.gitignore

@@ -18,6 +18,7 @@ eggs/
 lib/
 lib64/
 parts/
+!robosuite/controllers/config/default/parts
 !robosuite/controllers/parts
 sdist/
 var/

robosuite/controllers/composite/composite_controller_factory.py

@@ -1,3 +1,4 @@
+import copy
 import json
 import os
 import pathlib
@@ -18,6 +19,57 @@ def validate_composite_controller_config(config: dict):
             raise ValueError
 
 
+def is_part_controller_config(config: Dict):
+    """
+    Checks if a controller config is a part config as a opposed to a composite
+    controller config.
+
+    Args:
+        config (dict): Controller configuration
+
+    Returns:
+        bool: True if the config is in the for the arm-only, False otherwise
+    """
+
+    PART_CONTROLLER_TYPES = ["JOINT_VELOCITY", "JOINT_TORQUE", "JOINT_POSITION", "OSC_POSITION", "OSC_POSE", "IK_POSE"]
+    if "body_parts_controller_configs" not in config and "type" in config:
+        return config["type"] in PART_CONTROLLER_TYPES
+    return False
+
+
+def refactor_composite_controller_config(controller_config, robot_type, arms):
+    """
+    Checks if a controller config is in the format from robosuite versions <= 1.4.1.
+    If this is the case, converts the controller config to the new composite controller
+    config format in robosuite versions >= 1.5. If the robot has a default
+    controller config use that and override the arms with the old controller config.
+    If not just use the old controller config for arms.
+
+    Args:
+        old_controller_config (dict): Old controller config
+
+    Returns:
+        dict: New controller config
+    """
+    if not is_part_controller_config(controller_config):
+        return controller_config
+
+    config_dir = pathlib.Path(robosuite.__file__).parent / "controllers/config/robots/"
+    name = robot_type.lower()
+    configs = os.listdir(config_dir)
+    if f"default_{name}.json" in configs:
+        new_controller_config = load_composite_controller_config(robot=name)
+    else:
+        new_controller_config = {}
+        new_controller_config["type"] = "BASIC"
+        new_controller_config["body_parts"] = {}
+
+    for arm in arms:
+        new_controller_config["body_parts"][arm] = copy.deepcopy(controller_config)
+        new_controller_config["body_parts"][arm]["gripper"] = {"type": "GRIP"}
+    return new_controller_config
+
+
 def load_composite_controller_config(controller: Optional[str] = None, robot: Optional[str] = None) -> Optional[Dict]:
     """
     Utility function that loads the desired composite controller and returns the loaded configuration as a dict

robosuite/controllers/config/default/parts/ik_pose.json

@@ -0,0 +1,7 @@
+{
+    "type" : "IK_POSE",
+    "ik_pos_limit": 0.02,
+    "ik_ori_limit": 0.05,
+    "interpolation": null,
+    "ramp_ratio": 0.2
+  }

robosuite/controllers/config/default/parts/joint_position.json

@@ -0,0 +1,15 @@
+{
+    "type": "JOINT_POSITION",
+    "input_max": 1,
+    "input_min": -1,
+    "output_max": 0.05,
+    "output_min": -0.05,
+    "kp": 50,
+    "damping_ratio": 1,
+    "impedance_mode": "fixed",
+    "kp_limits": [0, 300],
+    "damping_ratio_limits": [0, 10],
+    "qpos_limits": null,
+    "interpolation": null,
+    "ramp_ratio": 0.2
+  }

robosuite/controllers/config/default/parts/joint_torque.json

@@ -0,0 +1,10 @@
+{
+    "type": "JOINT_TORQUE",
+    "input_max": 1,
+    "input_min": -1,
+    "output_max": 0.1,
+    "output_min": -0.1,
+    "torque_limits": null,
+    "interpolation": null,
+    "ramp_ratio": 0.2
+  }

robosuite/controllers/config/default/parts/joint_velocity.json

@@ -0,0 +1,11 @@
+{
+    "type" : "JOINT_VELOCITY",
+    "input_max": 1,
+    "input_min": -1,
+    "output_max": 0.5,
+    "output_min": -0.5,
+    "kp": 3.0,
+    "velocity_limits": [-1,1],
+    "interpolation": null,
+    "ramp_ratio": 0.2
+  }

robosuite/controllers/config/default/parts/osc_pose.json

@@ -0,0 +1,18 @@
+{
+    "type": "OSC_POSE",
+    "input_max": 1,
+    "input_min": -1,
+    "output_max": [0.05, 0.05, 0.05, 0.5, 0.5, 0.5],
+    "output_min": [-0.05, -0.05, -0.05, -0.5, -0.5, -0.5],
+    "kp": 150,
+    "damping_ratio": 1,
+    "impedance_mode": "fixed",
+    "kp_limits": [0, 300],
+    "damping_ratio_limits": [0, 10],
+    "position_limits": null,
+    "orientation_limits": null,
+    "uncouple_pos_ori": true,
+    "control_delta": true,
+    "interpolation": null,
+    "ramp_ratio": 0.2
+  }

robosuite/controllers/config/default/parts/osc_position.json

@@ -0,0 +1,16 @@
+{
+    "type": "OSC_POSITION",
+    "input_max": 1,
+    "input_min": -1,
+    "output_max": [0.05, 0.05, 0.05],
+    "output_min": [-0.05, -0.05, -0.05],
+    "kp": 150,
+    "damping_ratio": 1,
+    "impedance_mode": "fixed",
+    "kp_limits": [0, 300],
+    "damping_ratio_limits": [0, 10],
+    "position_limits": null,
+    "control_delta": true,
+    "interpolation": null,
+    "ramp_ratio": 0.2
+  }

robosuite/controllers/parts/controller_factory.py

@@ -52,7 +52,7 @@ def load_part_controller_config(custom_fpath=None, default_controller=None):
 
         # Store the default controller config fpath associated with the requested controller
         custom_fpath = os.path.join(
-            os.path.dirname(__file__), "..", "controllers/config/parts/{}.json".format(default_controller.lower())
+            os.path.dirname(__file__), "..", "config/default/parts/{}.json".format(default_controller.lower())
         )
 
     # Assert that the fpath to load the controller is not empty
@@ -122,7 +122,7 @@ def arm_controller_factory(name, params):
             ori_interpolator = deepcopy(interpolator)
             ori_interpolator.set_states(dim=4, ori="quat")
 
-        from robosuite.controllers.arm.ik import InverseKinematicsController
+        from robosuite.controllers.parts.arm.ik import InverseKinematicsController
 
         return InverseKinematicsController(
             interpolator_pos=interpolator,

robosuite/demos/demo_collect_and_playback_data.py

@@ -47,7 +47,7 @@ def collect_random_trajectory(env, timesteps=1000, max_fr=None):
                 time.sleep(diff)
 
 
-def playback_trajectory(env, ep_dir):
+def playback_trajectory(env, ep_dir, max_fr=None):
     """Playback data from an episode.
 
     Args:
@@ -69,13 +69,22 @@ def playback_trajectory(env, ep_dir):
         dic = np.load(state_file)
         states = dic["states"]
         for state in states:
+            start = time.time()
             env.sim.set_state_from_flattened(state)
             env.sim.forward()
+            env.viewer.update()
             env.render()
             t += 1
             if t % 100 == 0:
                 print(t)
 
+            if max_fr is not None:
+                elapsed = time.time() - start
+                diff = 1 / max_fr - elapsed
+                if diff > 0:
+                    time.sleep(diff)
+    env.close()
+
 
 if __name__ == "__main__":
 
@@ -120,4 +129,4 @@ def playback_trajectory(env, ep_dir):
     _ = input("Press any key to begin the playback...")
     print("Playing back the data...")
     data_directory = env.ep_directory
-    playback_trajectory(env, data_directory)
+    playback_trajectory(env, data_directory, args.max_fr)

robosuite/demos/demo_control.py

@@ -48,6 +48,7 @@
 from typing import Dict
 
 import robosuite as suite
+from robosuite.controllers.composite.composite_controller_factory import refactor_composite_controller_config
 from robosuite.utils.input_utils import *
 
 MAX_FR = 25  # max frame rate for running simluation
@@ -92,10 +93,14 @@
     joint_dim = 6 if options["robots"] == "UR5e" else (16 if options["robots"] == "GR1" else 7)
 
     # Choose controller
-    controller_name = choose_controller()
+    controller_name = choose_controller(part_controllers=True)
 
     # Load the desired controller
-    options["controller_configs"] = suite.load_part_controller_config(default_controller=controller_name)
+    arm_controller_config = suite.load_part_controller_config(default_controller=controller_name)
+    robot = options["robots"][0] if isinstance(options["robots"], list) else options["robots"]
+    options["controller_configs"] = refactor_composite_controller_config(
+        arm_controller_config, robot, ["right", "left"]
+    )
 
     # Define the pre-defined controller actions to use (action_dim, num_test_steps, test_value)
     controller_settings = {

robosuite/demos/demo_device_control.py

@@ -95,6 +95,7 @@
 
 import robosuite as suite
 from robosuite import load_composite_controller_config
+from robosuite.controllers.composite.composite_controller import WholeBody
 from robosuite.wrappers import VisualizationWrapper
 
 if __name__ == "__main__":
@@ -108,7 +109,12 @@
     parser.add_argument("--arm", type=str, default="right", help="Which arm to control (eg bimanual) 'right' or 'left'")
     parser.add_argument("--switch-on-grasp", action="store_true", help="Switch gripper control on gripper action")
     parser.add_argument("--toggle-camera-on-grasp", action="store_true", help="Switch camera angle on gripper action")
-    parser.add_argument("--controller", type=str, default="BASIC", help="Choice of controller. Can be 'ik' or 'osc'")
+    parser.add_argument(
+        "--controller",
+        type=str,
+        default=None,
+        help="Choice of controller. Can be generic (eg. 'BASIC' or 'WHOLE_BODY_MINK_IK') or json file (see robosuite/controllers/config for examples) or None to get the robot's default controller if it exists",
+    )
     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")
@@ -168,6 +174,10 @@
         from robosuite.devices import SpaceMouse
 
         device = SpaceMouse(env=env, pos_sensitivity=args.pos_sensitivity, rot_sensitivity=args.rot_sensitivity)
+    elif args.device == "mjgui":
+        from robosuite.devices.mjgui import MJGUI
+
+        device = MJGUI(env=env)
     else:
         raise Exception("Invalid device choice: choose either 'keyboard' or 'spacemouse'.")
 
@@ -185,65 +195,54 @@
 
         # Initialize device control
         device.start_control()
-        action_dict = device.input2action()
-
-        # get key with "delta" in it
-        action_key = [key for key in action_dict.keys() if "delta" in key]
-        assert len(action_key) == 1, "Error: Key with 'delta' in it not found!"
-        action_key = action_key[0]
-        gripper_key = [key for key in action_dict.keys() if "gripper" in key]
-        assert len(gripper_key) == 1, "Error: Key with 'gripper' in it not found!"
-        gripper_key = gripper_key[0]
-
+        all_prev_gripper_actions = [
+            {
+                f"{robot_arm}_gripper": np.repeat([0], robot.gripper[robot_arm].dof)
+                for robot_arm in robot.arms
+                if robot.gripper[robot_arm].dof > 0
+            }
+            for robot in env.robots
+        ]
+
+        # 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[0] if args.config == "bimanual" else env.robots[args.arm == "left"]
+            active_robot = env.robots[device.active_robot]
 
             # Get the newest action
-            action_dict = device.input2action()
+            input_ac_dict = device.input2action()
 
-            # If action_dict is none, then this a reset so we should break
-            if action_dict is None:
+            # If action is none, then this a reset so we should break
+            if input_ac_dict is None:
                 break
 
-            action = action_dict[action_key]
-            grasp = action_dict[gripper_key]
-
-            # If the current grasp is active (1) and last grasp is not (-1) (i.e.: grasping input just pressed),
-            # toggle arm control and / or camera viewing angle if requested
-            if last_grasp < 0 < grasp:
-                if args.switch_on_grasp:
-                    args.arm = "left" if args.arm == "right" else "right"
-                if args.toggle_camera_on_grasp:
-                    cam_id = (cam_id + 1) % num_cam
-                    env.viewer.set_camera(camera_id=cam_id)
-            # Update last grasp
-            last_grasp = grasp
-
-            # Fill out the rest of the action space if necessary
-            rem_action_dim = env.action_dim - action.size
-            if rem_action_dim > 0:
-                # Initialize remaining action space
-                rem_action = np.zeros(rem_action_dim)
-                # This is a multi-arm setting, choose which arm to control and fill the rest with zeros
-                if args.arm == "right":
-                    action = np.concatenate([action, rem_action])
-                elif args.arm == "left":
-                    action = np.concatenate([rem_action, action])
+            from copy import deepcopy
+
+            action_dict = deepcopy(input_ac_dict)  # {}
+            # set arm actions
+            for arm in active_robot.arms:
+                if isinstance(active_robot.composite_controller, WholeBody):  # input type passed to joint_action_policy
+                    controller_input_type = active_robot.composite_controller.joint_action_policy.input_type
+                else:
+                    controller_input_type = active_robot.part_controllers[arm].input_type
+
+                if controller_input_type == "delta":
+                    action_dict[arm] = input_ac_dict[f"{arm}_delta"]
+                elif controller_input_type == "absolute":
+                    action_dict[arm] = input_ac_dict[f"{arm}_abs"]
                 else:
-                    # Only right and left arms supported
-                    print(
-                        "Error: Unsupported arm specified -- "
-                        "must be either 'right' or 'left'! Got: {}".format(args.arm)
-                    )
-            elif rem_action_dim < 0:
-                # We're in an environment with no gripper action space, so trim the action space to be the action dim
-                action = action[: env.action_dim]
-
-            # Step through the simulation and render
-            obs, reward, done, info = env.step(action)
+                    raise ValueError
+
+            # Maintain gripper state for each robot but only update the active robot with action
+            env_action = [robot.create_action_vector(all_prev_gripper_actions[i]) for i, robot in enumerate(env.robots)]
+            env_action[device.active_robot] = active_robot.create_action_vector(action_dict)
+            env_action = np.concatenate(env_action)
+            for gripper_ac in all_prev_gripper_actions[device.active_robot]:
+                all_prev_gripper_actions[device.active_robot][gripper_ac] = action_dict[gripper_ac]
+
+            env.step(env_action)
             env.render()
 
             # limit frame rate if necessary

robosuite/demos/demo_domain_randomization.py

@@ -2,6 +2,8 @@
 Script to showcase domain randomization functionality.
 """
 
+import mujoco
+
 import robosuite.macros as macros
 from robosuite.utils.input_utils import *
 from robosuite.wrappers import DomainRandomizationWrapper
@@ -10,7 +12,7 @@
 macros.USING_INSTANCE_RANDOMIZATION = True
 
 if __name__ == "__main__":
-
+    assert mujoco.__version__ == "3.1.1", "Script requires mujoco-py version 3.1.1 to run"
     # Create dict to hold options that will be passed to env creation call
     options = {}