diff --git a/robosuite/demos/demo_composite_robot.py b/robosuite/demos/demo_composite_robot.py new file mode 100644 index 0000000000..2b24d6c0bf --- /dev/null +++ b/robosuite/demos/demo_composite_robot.py @@ -0,0 +1,25 @@ +import argparse + +import numpy as np + +import robosuite as suite +import robosuite.utils.test_utils as tu +from robosuite.controllers import load_composite_controller_config +from robosuite.utils.robot_composition_utils import create_composite_robot + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + + parser.add_argument("--robot", type=str, required=True) + parser.add_argument("--base", type=str, default=None) + parser.add_argument("--grippers", nargs="+", type=str, default=["PandaGripper"]) + parser.add_argument("--env", type=str, default="Lift") + + args = parser.parse_args() + + name = f"Custom{args.robot}" + create_composite_robot(name, base=args.base, robot=args.robot, grippers=args.grippers) + controller_config = load_composite_controller_config(controller="BASIC", robot=name) + + tu.create_and_test_env(env="Lift", robots=name, controller_config=controller_config) diff --git a/robosuite/models/assets/robots/baxter/robot.xml b/robosuite/models/assets/robots/baxter/robot.xml index b1142f0e8c..349b140872 100644 --- a/robosuite/models/assets/robots/baxter/robot.xml +++ b/robosuite/models/assets/robots/baxter/robot.xml @@ -1,6 +1,6 @@ + - @@ -9,23 +9,19 @@ - - - - @@ -39,13 +35,11 @@ - - @@ -55,13 +49,11 @@ - - @@ -69,133 +61,125 @@ - - + - + - + - + - - - - - - - - + + + + + + + + - + - - - - - - - - - - - + + + + + + + + + + + + - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -204,66 +188,60 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -277,7 +255,6 @@ - @@ -285,7 +262,6 @@ - @@ -293,7 +269,5 @@ - - - - + + \ No newline at end of file diff --git a/robosuite/models/assets/robots/iiwa/robot.xml b/robosuite/models/assets/robots/iiwa/robot.xml index 849d7bc2d4..71f58b5c79 100644 --- a/robosuite/models/assets/robots/iiwa/robot.xml +++ b/robosuite/models/assets/robots/iiwa/robot.xml @@ -2,13 +2,13 @@ - - - - - - - + + + + + + + @@ -31,51 +31,97 @@ - - + + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -86,4 +132,4 @@ - + \ No newline at end of file diff --git a/robosuite/models/robots/__init__.py b/robosuite/models/robots/__init__.py index 36b7a36c9e..e9188e4a3c 100644 --- a/robosuite/models/robots/__init__.py +++ b/robosuite/models/robots/__init__.py @@ -1,3 +1,15 @@ from .robot_model import RobotModel, create_robot from .manipulators import * from .compositional import * + + +def is_robosuite_robot(robot: str) -> bool: + """ + robot is robosuite repo robot if can import robot class from robosuite.models.robots + """ + try: + module = __import__("robosuite.models.robots", fromlist=[robot]) + getattr(module, robot) + return True + except (ImportError, AttributeError): + return False diff --git a/robosuite/utils/robot_composition_utils.py b/robosuite/utils/robot_composition_utils.py new file mode 100644 index 0000000000..cb757f2eff --- /dev/null +++ b/robosuite/utils/robot_composition_utils.py @@ -0,0 +1,67 @@ +from typing import List, Optional, Tuple, Union + +from robosuite.models.robots.robot_model import REGISTERED_ROBOTS, RobotModel +from robosuite.robots import register_robot_class +from robosuite.utils.log_utils import ROBOSUITE_DEFAULT_LOGGER +from robosuite.utils.robot_utils import check_bimanual + +BASE_TARGET_MAPPING = { + "RethinkMount": "FixedBaseRobot", + "RethinkMinimalMount": "FixedBaseRobot", + "NullMount": "FixedBaseRobot", + "OmronMobileBase": "WheeledRobot", + "NullMobileBase": "WheeledRobot", + "NoActuationBase": "LeggedRobot", + "Spot": "LeggedRobot", + "SpotFloating": "LeggedRobot", +} + + +def get_target_type(base) -> str: + """ + Returns the target type of the robot + """ + return BASE_TARGET_MAPPING[base] + + +def create_composite_robot( + name: str, robot: str, base: Optional[str] = None, grippers: Optional[Union[str, List[str], Tuple[str]]] = None +) -> RobotModel: + """ + Factory function to create a composite robot + """ + + bimanual_robot = check_bimanual(robot) + grippers = grippers if type(grippers) == list or type(grippers) == tuple else [grippers] + + # perform checks and issues warning if necessary + if bimanual_robot and len(grippers) == 1: + grippers = grippers + grippers + if not bimanual_robot and len(grippers) == 2: + ROBOSUITE_DEFAULT_LOGGER.warning( + f"Grippers {grippers} supplied for single gripper robot.\ + Using gripper {grippers[0]}." + ) + if not base: + base = REGISTERED_ROBOTS[robot]().default_base + if robot in ["Tiago", "GR1"] and base: + ROBOSUITE_DEFAULT_LOGGER.warning(f"Defined custom base when using {robot} robot. Ignoring base.") + if robot in ["Tiago"]: + base = "NullMobileBase" + elif robot == "GR1": + base = "NoActuationBase" + + target_type = get_target_type(base) + + class_dict = { + "default_base": property(lambda self: base), + "default_arms": property(lambda self: {"right": robot}), + "default_gripper": property( + lambda self: {"right": grippers[0], "left": grippers[1]} if bimanual_robot else {"right": grippers[0]} + ), + } + + CustomCompositeRobotClass = type(name, (REGISTERED_ROBOTS[robot],), class_dict) + register_robot_class(target_type)(CustomCompositeRobotClass) + + return CustomCompositeRobotClass diff --git a/tests/test_robots/test_composite_robots.py b/tests/test_robots/test_composite_robots.py new file mode 100644 index 0000000000..c72892598b --- /dev/null +++ b/tests/test_robots/test_composite_robots.py @@ -0,0 +1,100 @@ +""" +Script to test composite robots: + +$ pytest -s tests/test_robots/test_composite_robots.py +""" +import logging +from typing import Dict, List, Union + +import numpy as np +import pytest + +import robosuite as suite +import robosuite.utils.robot_composition_utils as cu +from robosuite.controllers import load_composite_controller_config +from robosuite.models.grippers import GRIPPER_MAPPING +from robosuite.models.robots import is_robosuite_robot +from robosuite.utils.log_utils import ROBOSUITE_DEFAULT_LOGGER + +ROBOSUITE_DEFAULT_LOGGER.setLevel(logging.ERROR) + +TEST_ROBOTS = ["Baxter", "IIWA", "Jaco", "Kinova3", "Panda", "Sawyer", "UR5e", "Tiago", "SpotArm", "GR1"] +TEST_BASES = [ + "RethinkMount", + "RethinkMinimalMount", + "NullMount", + "OmronMobileBase", + "NullMobileBase", + "NoActuationBase", + "Spot", + "SpotFloating", +] + +# If you would like to visualize the scene during testing, +# set render to True and increase env_steps to a larger value. +def create_and_test_env( + env: str, + robots: Union[str, List[str]], + controller_config: Dict, + render: bool = True, + env_steps: int = 20, +): + + config = { + "env_name": env, + "robots": robots, + "controller_configs": controller_config, + } + + env = suite.make( + **config, + has_renderer=render, + has_offscreen_renderer=False, + ignore_done=True, + use_camera_obs=False, + reward_shaping=True, + control_freq=20, + ) + env.reset() + low, high = env.action_spec + low = np.clip(low, -1, 1) + high = np.clip(high, -1, 1) + + # Runs a few steps of the simulation as a sanity check + for i in range(env_steps): + action = np.random.uniform(low, high) + obs, reward, done, _ = env.step(action) + if render: + env.render() + + env.close() + + +@pytest.mark.parametrize("robot", TEST_ROBOTS) +@pytest.mark.parametrize("base", TEST_BASES) +def test_composite_robot_base_combinations(robot, base): + if is_robosuite_robot(robot): + if robot in ["Tiago", "GR1", "SpotArm"]: + pytest.skip(f"Skipping {robot} for now since it we typically do not attach it to another base.") + elif base in ["NullMobileBase", "NoActuationBase", "Spot", "SpotFloating"]: + pytest.skip(f"Skipping {base} for now since comopsite robots do not use {base}.") + else: + cu.create_composite_robot(name="CompositeRobot", robot=robot, base=base, grippers="RethinkGripper") + controller_config = load_composite_controller_config(controller="BASIC", robot="CompositeRobot") + create_and_test_env(env="Lift", robots="CompositeRobot", controller_config=controller_config, render=False) + + +@pytest.mark.parametrize("robot", TEST_ROBOTS) +@pytest.mark.parametrize("gripper", GRIPPER_MAPPING.keys()) +def test_composite_robot_gripper_combinations(robot, gripper): + if is_robosuite_robot(robot): + if robot in ["Tiago"]: + base = "NullMobileBase" + elif robot == "GR1": + base = "NoActuationBase" + else: + base = "RethinkMount" + + cu.create_composite_robot(name="CompositeRobot", robot=robot, base=base, grippers=gripper) + controller_config = load_composite_controller_config(controller="BASIC", robot="CompositeRobot") + create_and_test_env(env="Lift", robots="CompositeRobot", controller_config=controller_config, render=False)