diff --git a/.gitignore b/.gitignore index a919c684d4..17b247decf 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ eggs/ lib/ lib64/ parts/ +!robosuite/controllers/config/default/parts !robosuite/controllers/parts sdist/ var/ diff --git a/robosuite/controllers/composite/composite_controller_factory.py b/robosuite/controllers/composite/composite_controller_factory.py index 3a67b227c5..a531c1e5d2 100644 --- a/robosuite/controllers/composite/composite_controller_factory.py +++ b/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 diff --git a/robosuite/controllers/config/default/parts/ik_pose.json b/robosuite/controllers/config/default/parts/ik_pose.json new file mode 100644 index 0000000000..db6191e617 --- /dev/null +++ b/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 + } \ No newline at end of file diff --git a/robosuite/controllers/config/default/parts/joint_position.json b/robosuite/controllers/config/default/parts/joint_position.json new file mode 100644 index 0000000000..bc25f5cb91 --- /dev/null +++ b/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 + } \ No newline at end of file diff --git a/robosuite/controllers/config/default/parts/joint_torque.json b/robosuite/controllers/config/default/parts/joint_torque.json new file mode 100644 index 0000000000..8d648428ea --- /dev/null +++ b/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 + } \ No newline at end of file diff --git a/robosuite/controllers/config/default/parts/joint_velocity.json b/robosuite/controllers/config/default/parts/joint_velocity.json new file mode 100644 index 0000000000..4a9c7c86aa --- /dev/null +++ b/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 + } \ No newline at end of file diff --git a/robosuite/controllers/config/default/parts/osc_pose.json b/robosuite/controllers/config/default/parts/osc_pose.json new file mode 100644 index 0000000000..30b20677e7 --- /dev/null +++ b/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 + } \ No newline at end of file diff --git a/robosuite/controllers/config/default/parts/osc_position.json b/robosuite/controllers/config/default/parts/osc_position.json new file mode 100644 index 0000000000..c75e51238c --- /dev/null +++ b/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 + } \ No newline at end of file diff --git a/robosuite/controllers/parts/controller_factory.py b/robosuite/controllers/parts/controller_factory.py index b5c6331852..8c780f8d48 100644 --- a/robosuite/controllers/parts/controller_factory.py +++ b/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, diff --git a/robosuite/demos/demo_collect_and_playback_data.py b/robosuite/demos/demo_collect_and_playback_data.py index ef44c4b11d..0dd3a48061 100644 --- a/robosuite/demos/demo_collect_and_playback_data.py +++ b/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) diff --git a/robosuite/demos/demo_control.py b/robosuite/demos/demo_control.py index 8d7fab6cdb..1622b255af 100644 --- a/robosuite/demos/demo_control.py +++ b/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 = { diff --git a/robosuite/demos/demo_device_control.py b/robosuite/demos/demo_device_control.py index 224ed6882e..673b6551bd 100644 --- a/robosuite/demos/demo_device_control.py +++ b/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 diff --git a/robosuite/demos/demo_domain_randomization.py b/robosuite/demos/demo_domain_randomization.py index 5f5503cf96..e1113ec5ad 100644 --- a/robosuite/demos/demo_domain_randomization.py +++ b/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 = {} diff --git a/robosuite/utils/input_utils.py b/robosuite/utils/input_utils.py index 1931b7539b..a7f6860323 100644 --- a/robosuite/utils/input_utils.py +++ b/robosuite/utils/input_utils.py @@ -39,7 +39,7 @@ def choose_environment(): return envs[k] -def choose_controller(): +def choose_controller(part_controllers=False): """ Prints out controller options, and returns the requested controller name @@ -47,7 +47,7 @@ def choose_controller(): str: Chosen controller name """ # get the list of all controllers - controllers = list(suite.ALL_COMPOSITE_CONTROLLERS) + controllers = list(suite.ALL_PART_CONTROLLERS) if part_controllers else list(suite.ALL_COMPOSITE_CONTROLLERS) # Select controller to use print("Here is a list of controllers in the suite:\n")