diff --git a/examples/robosuite/dual_table_workspace.scenic b/examples/robosuite/dual_table_workspace.scenic new file mode 100644 index 000000000..c4824c746 --- /dev/null +++ b/examples/robosuite/dual_table_workspace.scenic @@ -0,0 +1,18 @@ +# examples/robosuite/dual_table_workspace.scenic +model scenic.simulators.robosuite.model + +# Arena Setup +back_table = new Table at (-0.6, 0, 0.8) +front_table = new Table at (0.6, 0, 0.8) + +# OBJECTS +red_cube = new Box at (Range(-0.7, -0.5), Range(-0.1, 0.1), 0.83), + with color (1, 0, 0, 1), + + +green_cube = new Box at (Range(0.5, 0.7), Range(-0.1, 0.1), 0.83), + with color (0, 1, 0, 1), + with width 0.05, with length 0.05, with height 0.05 + +# ROBOT +ego = new Panda at (0, 0, 0) \ No newline at end of file diff --git a/examples/robosuite/four_table_workspace.scenic b/examples/robosuite/four_table_workspace.scenic new file mode 100644 index 000000000..456f42cc0 --- /dev/null +++ b/examples/robosuite/four_table_workspace.scenic @@ -0,0 +1,49 @@ +# examples/robosuite/four_table_workspace.scenic +model scenic.simulators.robosuite.model + +# CAMERA CONFIGURATION +# param camera_view = "sideview" + +TABLE_DISTANCE = 1.0 + +# WORKSPACE: Four tables in cross formation +back_table = new Table at (-TABLE_DISTANCE, 0, 0.8) +front_table = new Table at (TABLE_DISTANCE, 0, 0.8) +right_table = new Table at (0, TABLE_DISTANCE, 0.8) +left_table = new Table at (0, -TABLE_DISTANCE, 0.8) + +# TABLE 1 (Back): Primitive objects with random positions +ball = new Ball at (Range(-1.1, -0.9), Range(-0.1, 0.1), 0.83), + with color (1, 0.5, 0, 1) + +box = new Box at (Range(-1.1, -0.9), Range(-0.1, 0.1), 0.83), + with color (0, 0, 1, 1) + +capsule = new Capsule at (Range(-1.1, -0.9), Range(-0.1, 0.1), 0.83), + with color (0.5, 0.5, 0.5, 1) + +cylinder = new Cylinder at (Range(-1.1, -0.9), Range(-0.1, 0.1), 0.83), + with color (0, 1, 1, 1) + +# TABLE 2 (Front): Nuts and food items +square_nut = new SquareNut at (Range(0.9, 1.1), Range(-0.1, 0.1), 0.85) + +round_nut = new RoundNut at (Range(0.9, 1.1), Range(-0.1, 0.1), 0.85) + +milk = new Milk at (Range(0.9, 1.1), Range(-0.1, 0.1), 0.85) + +cereal = new Cereal at (Range(0.9, 1.1), Range(-0.1, 0.1), 0.85) + +# TABLE 3 (Right): Various objects +can = new Can at (Range(-0.1, 0.1), Range(0.9, 1.1), 0.85) + +bread = new Bread at (Range(-0.1, 0.1), Range(0.9, 1.1), 0.85) + +bottle = new Bottle at (Range(-0.1, 0.1), Range(0.9, 1.1), 0.85) + +hammer = new Hammer at (Range(-0.1, 0.1), Range(0.9, 1.1), 0.85) + +# TABLE 4 (Left): Intentionally left empty + +# ROBOT +ego = new Panda at (0, 0, 0) \ No newline at end of file diff --git a/examples/robosuite/stack_lift.scenic b/examples/robosuite/stack_lift.scenic new file mode 100644 index 000000000..488003e02 --- /dev/null +++ b/examples/robosuite/stack_lift.scenic @@ -0,0 +1,28 @@ +model scenic.simulators.robosuite.model + +work_table = new Table at (0.6, 0, 0.8), + with width 0.6, + with length 1.2, + with height 0.05 + +bottom_cube = new Box at (0.6, 0, 0.83), + with color (0.2, 0.3, 0.8, 1), + with width 0.06, with length 0.06, with height 0.06 + +top_cube = new Box at (0.6, 0, 0.89) + with color (0.8, 0.2, 0.2, 1), + +bottle = new Bottle at (0.6, 0.3, 0.83) + +pickup_object = top_cube + +behavior StackLift(): + """Pick up the top cube from the stack and lift it.""" + do PickObject(pickup_object) + do LiftToHeight(1.05) + for _ in range(10): + if pickup_object.position.z > 1.0: + terminate simulation + wait +ego = new UR5e at (0, 0, 0), + with behavior StackLift() \ No newline at end of file diff --git a/examples/robosuite/stacked_cubes.scenic b/examples/robosuite/stacked_cubes.scenic new file mode 100644 index 000000000..7fc0ecd0f --- /dev/null +++ b/examples/robosuite/stacked_cubes.scenic @@ -0,0 +1,19 @@ +model scenic.simulators.robosuite.model + +# param camera_view = "sideview" + +work_table = new Table at (0.6, 0, 0.8), + with width 0.6, + with length 1.2, + with height 0.05 + +bottom_cube = new Box at (0.6, 0, 0.83), + with color (0.2, 0.3, 0.8, 1), + with width 0.06, with length 0.06, with height 0.06 + +top_cube = new Box at (0.6, 0, 0.89), + with color (0.8, 0.2, 0.2, 1), + +bottle = new Bottle at (0.6, 0.3, 0.83) + +ego = new UR5e at (0, 0, 0) \ No newline at end of file diff --git a/src/scenic/simulators/robosuite/__init__.py b/src/scenic/simulators/robosuite/__init__.py new file mode 100644 index 000000000..287e36abe --- /dev/null +++ b/src/scenic/simulators/robosuite/__init__.py @@ -0,0 +1,5 @@ +"""RoboSuite simulator interface for Scenic.""" + +from .simulator import RobosuiteSimulator, RobosuiteSimulation + +__all__ = ['RobosuiteSimulator', 'RobosuiteSimulation'] \ No newline at end of file diff --git a/src/scenic/simulators/robosuite/model.scenic b/src/scenic/simulators/robosuite/model.scenic new file mode 100644 index 000000000..0cf1a794e --- /dev/null +++ b/src/scenic/simulators/robosuite/model.scenic @@ -0,0 +1,253 @@ +# src/scenic/simulators/robosuite/model.scenic + +"""Scenic world model for RoboSuite - Custom Environments Only.""" + +from .simulator import RobosuiteSimulator, SetJointPositions, OSCPositionAction + +# Global parameters with defaults matching Robosuite's defaults +param env_config = {} +param controller_config = None +param camera_view = None +param render = True +param real_time = True +param speed = 1.0 +param lite_physics = None # None = use Robosuite default (True) + +# Simulator - no more use_environment parameter +simulator RobosuiteSimulator( + render=globalParameters.render, + real_time=globalParameters.real_time, + speed=globalParameters.speed, + env_config=globalParameters.env_config, + controller_config=globalParameters.controller_config, + camera_view=globalParameters.camera_view, + lite_physics=globalParameters.lite_physics +) + +# Default values dictionary +DEFAULTS = { + # Object properties + 'object_size': 0.03, + 'density': 1000, + 'friction': (1.0, 0.005, 0.0001), + 'solref': (0.02, 1.0), + 'solimp': (0.9, 0.95, 0.001, 0.5, 2.0), + 'default_color': (0.5, 0.5, 0.5, 1.0), + + # Arena properties + 'table_height': 0.8, + 'table_width': 0.8, + 'table_length': 0.8, + 'table_thickness': 0.05, + + # Robot properties + 'robot_width': 0.2, + 'robot_length': 0.2, + 'robot_height': 0.5, + + # Control parameters + 'control_gain': 3.0, + 'control_limit': 0.3, + 'position_tolerance': 0.02, + 'height_tolerance': 0.02, + 'gripper_open_steps': 20, + 'gripper_close_steps': 30, + 'max_control_steps': 100, + 'max_lift_steps': 200 +} + +# Base classes +class RoboSuiteObject(Object): + """Base class for all RoboSuite objects.""" + density: DEFAULTS['density'] + friction: DEFAULTS['friction'] + solref: DEFAULTS['solref'] + solimp: DEFAULTS['solimp'] + shape: BoxShape() + allowCollisions: True + +# Table for arena setup +class Table(RoboSuiteObject): + """Table in environment.""" + isTable: True + width: DEFAULTS['table_width'] + length: DEFAULTS['table_length'] + height: DEFAULTS['table_thickness'] + position: (0, 0, DEFAULTS['table_height']) + +# Base class for manipulable objects +class ManipulationObject(RoboSuiteObject): + """Base class for objects that can be manipulated.""" + color: DEFAULTS['default_color'] + +# Primitive shape objects (matching RoboSuite's naming) +class Box(ManipulationObject): + """Box object.""" + objectType: "Box" + width: DEFAULTS['object_size'] + length: DEFAULTS['object_size'] + height: DEFAULTS['object_size'] + +class Ball(ManipulationObject): + """Ball/sphere object.""" + objectType: "Ball" + radius: DEFAULTS['object_size'] + width: DEFAULTS['object_size'] * 2 + length: DEFAULTS['object_size'] * 2 + height: DEFAULTS['object_size'] * 2 + +class Cylinder(ManipulationObject): + """Cylinder object.""" + objectType: "Cylinder" + width: DEFAULTS['object_size'] * 2 + length: DEFAULTS['object_size'] * 2 + height: DEFAULTS['object_size'] * 4 + +class Capsule(ManipulationObject): + """Capsule object.""" + objectType: "Capsule" + width: DEFAULTS['object_size'] * 1.5 + length: DEFAULTS['object_size'] * 1.5 + height: DEFAULTS['object_size'] * 3 + +# Complex objects (matching RoboSuite's naming) +class Milk(ManipulationObject): + """Milk carton object.""" + objectType: "Milk" + +class Cereal(ManipulationObject): + """Cereal box object.""" + objectType: "Cereal" + +class Can(ManipulationObject): + """Can object.""" + objectType: "Can" + +class Bread(ManipulationObject): + """Bread object.""" + objectType: "Bread" + +class Bottle(ManipulationObject): + """Bottle object.""" + objectType: "Bottle" + +class Hammer(ManipulationObject): + """Hammer object.""" + objectType: "Hammer" + +class SquareNut(ManipulationObject): + """Square nut object.""" + objectType: "SquareNut" + +class RoundNut(ManipulationObject): + """Round nut object.""" + objectType: "RoundNut" + +# Robots (matching RoboSuite's naming) +class Robot(RoboSuiteObject): + """Base robot class.""" + robot_type: "Panda" + gripper_type: "default" + controller_config: None + initial_qpos: None + base_type: "default" + width: DEFAULTS['robot_width'] + length: DEFAULTS['robot_length'] + height: DEFAULTS['robot_height'] + + # Dynamic properties + joint_positions[dynamic]: [] + eef_pos[dynamic]: [0, 0, 0] + gripper_state[dynamic]: [0, 0] + +class Panda(Robot): + """Franka Emika Panda robot.""" + robot_type: "Panda" + gripper_type: "PandaGripper" + +class UR5e(Robot): + """Universal Robots UR5e.""" + robot_type: "UR5e" + gripper_type: "Robotiq85Gripper" + +class Sawyer(Robot): + """Rethink Robotics Sawyer.""" + robot_type: "Sawyer" + gripper_type: "RethinkGripper" + +class Jaco(Robot): + """Kinova Jaco robot.""" + robot_type: "Jaco" + gripper_type: "JacoThreeFingerGripper" + +class IIWA(Robot): + """KUKA IIWA robot.""" + robot_type: "IIWA" + gripper_type: "Robotiq140Gripper" + +# Behavior Library +behavior OpenGripper(steps=DEFAULTS['gripper_open_steps']): + """Open gripper over multiple steps.""" + for _ in range(steps): + take OSCPositionAction(gripper=-1) + +behavior CloseGripper(steps=DEFAULTS['gripper_close_steps']): + """Close gripper over multiple steps.""" + for _ in range(steps): + take OSCPositionAction(gripper=1) + +behavior MoveToPosition(target_pos, + tolerance=DEFAULTS['position_tolerance'], + max_steps=DEFAULTS['max_control_steps'], + gain=DEFAULTS['control_gain']): + """Move end-effector to target position.""" + for _ in range(max_steps): + eef_pos = self.eef_pos + error = [target_pos[i] - eef_pos[i] for i in range(3)] + + if sum(e**2 for e in error)**0.5 < tolerance: + return + + limit = DEFAULTS['control_limit'] + delta = [max(-limit, min(limit, e * gain)) for e in error] + take OSCPositionAction(position_delta=delta, gripper=-1) + +behavior MoveAboveObject(target_object, height_offset=0.1): + """Move above a tracked object.""" + target_pos = [target_object.position.x, + target_object.position.y, + target_object.position.z + height_offset] + do MoveToPosition(target_pos) + +behavior MoveToGrasp(target_object, grasp_offset=0.02): + """Move to grasp position for object.""" + target_pos = [target_object.position.x, + target_object.position.y, + target_object.position.z - grasp_offset] + do MoveToPosition(target_pos, tolerance=0.01) + +behavior LiftToHeight(target_height=1.0, max_steps=DEFAULTS['max_lift_steps']): + """Lift to absolute height.""" + for _ in range(max_steps): + eef_pos = self.eef_pos + error = [0, 0, target_height - eef_pos[2]] + + if abs(error[2]) < DEFAULTS['height_tolerance']: + return + + limit = DEFAULTS['control_limit'] + gain = DEFAULTS['control_gain'] + delta = [max(-limit, min(limit, e * gain)) for e in error] + take OSCPositionAction(position_delta=delta, gripper=1) + +behavior PickObject(target_object): + """Pick up a specific object.""" + do OpenGripper() + do MoveAboveObject(target_object) + do MoveToGrasp(target_object) + do CloseGripper() + +behavior PickAndLift(target_object, height=1.05): + """Complete pick and lift for specific object.""" + do PickObject(target_object) + do LiftToHeight(height) \ No newline at end of file diff --git a/src/scenic/simulators/robosuite/simulator.py b/src/scenic/simulators/robosuite/simulator.py new file mode 100644 index 000000000..ff9e8436f --- /dev/null +++ b/src/scenic/simulators/robosuite/simulator.py @@ -0,0 +1,566 @@ +# src/scenic/simulators/robosuite/simulator.py + +"""RoboSuite Simulator interface for Scenic - Custom Environment Only.""" + +from typing import Dict, List, Any, Optional, Union +import numpy as np +import mujoco + +try: + import robosuite as suite + from robosuite.environments.manipulation.manipulation_env import ManipulationEnv + from robosuite.models.arenas import EmptyArena, MultiTableArena + from robosuite.models.tasks import ManipulationTask + from robosuite.models.objects import ( + BallObject, BoxObject, CylinderObject, CapsuleObject, + MilkObject, CerealObject, CanObject, BreadObject, + HammerObject, SquareNutObject, RoundNutObject, BottleObject + ) + from robosuite.utils.observables import Observable, sensor +except ImportError as e: + suite = None + _import_error = e + +from scenic.core.simulators import Simulation, Simulator +from scenic.core.vectors import Vector +from scenic.core.dynamics import Action + +# Constants +DEFAULT_PHYSICS_TIMESTEP = 0.002 +DEFAULT_TIMESTEP = 0.01 +DEFAULT_ACTION_DIM = 7 +PHYSICS_SETTLE_STEPS = 10 + +# RoboSuite naming patterns +BODY_SUFFIXES = ['_main', '_body', ''] +JOINT_SUFFIX = '_joint0' + +# Camera view mapping +CAMERA_VIEWS = { + "frontview": 0, + "birdview": 1, + "agentview": 2, + "sideview": 3, + "robot0_robotview": 4, + "robot0_eye_in_hand": 5 +} + +# Object type mapping +OBJECT_FACTORIES = { + # Primitive objects with customizable size + 'Ball': lambda cfg: BallObject( + name=cfg['name'], + size=[cfg.get('radius', 0.02)], + rgba=cfg.get('color', [1, 0, 0, 1]) + ), + 'Box': lambda cfg: BoxObject( + name=cfg['name'], + size=cfg.get('size', [0.025, 0.025, 0.025]), + rgba=cfg.get('color', [1, 0, 0, 1]) + ), + 'Cylinder': lambda cfg: CylinderObject( + name=cfg['name'], + size=_extract_cylinder_size(cfg), + rgba=cfg.get('color', [1, 0, 0, 1]) + ), + 'Capsule': lambda cfg: CapsuleObject( + name=cfg['name'], + size=_extract_cylinder_size(cfg), + rgba=cfg.get('color', [1, 0, 0, 1]) + ), + # Complex objects with fixed sizes - no size parameter + 'Milk': lambda cfg: MilkObject(name=cfg['name']), + 'Cereal': lambda cfg: CerealObject(name=cfg['name']), + 'Can': lambda cfg: CanObject(name=cfg['name']), + 'Bread': lambda cfg: BreadObject(name=cfg['name']), + 'Hammer': lambda cfg: HammerObject(name=cfg['name']), + 'SquareNut': lambda cfg: SquareNutObject(name=cfg['name']), + 'RoundNut': lambda cfg: RoundNutObject(name=cfg['name']), + 'Bottle': lambda cfg: BottleObject(name=cfg['name']) +} + + +def _extract_cylinder_size(config: Dict) -> List[float]: + """Extract [radius, height] for cylinder-like objects.""" + if 'radius' in config and 'height' in config: + return [config['radius'], config['height']] + + size = config.get('size', [0.02, 0.02, 0.04]) + if len(size) == 3: # Convert from [width, length, height] + radius = (size[0] + size[1]) / 4 + return [radius, size[2]] + elif len(size) == 2: + return size + return [0.02, 0.04] # Default + + +class ScenicManipulationEnv(ManipulationEnv): + """Scenic-driven manipulation environment.""" + + def __init__(self, scenic_config: Dict, **kwargs): + self.scenic_config = scenic_config + self.scenic_objects = scenic_config.get('objects', []) + self.scenic_tables = scenic_config.get('tables', []) + super().__init__(**kwargs) + + def _load_model(self): + """Load models and create arena.""" + super()._load_model() + + self._position_robots() + mujoco_arena = self._create_arena() + mujoco_arena.set_origin([0, 0, 0]) + + self.mujoco_objects = self._create_objects() + + self.model = ManipulationTask( + mujoco_arena=mujoco_arena, + mujoco_robots=[robot.robot_model for robot in self.robots], + mujoco_objects=self.mujoco_objects + ) + + def _position_robots(self): + """Position robots based on scenic configuration.""" + for i, robot_config in enumerate(self.scenic_config.get('robots', [])): + if i < len(self.robots): + pos = robot_config.get('position', [0, 0, 0]) + self.robots[i].robot_model.set_base_xpos(pos) + + def _create_arena(self): + """Create arena based on table configuration.""" + if not self.scenic_tables: + # No tables - use empty arena (just floor) + return EmptyArena() + + # Use MultiTableArena for any number of tables + return MultiTableArena( + table_offsets=[t.get('position', [0, 0, 0.8]) for t in self.scenic_tables], + table_full_sizes=[t.get('size', (0.8, 0.8, 0.05)) for t in self.scenic_tables], + has_legs=[True] * len(self.scenic_tables) + ) + + def _create_objects(self) -> List: + """Create Robosuite objects from Scenic configuration.""" + mujoco_objects = [] + + for obj_config in self.scenic_objects: + obj_type = obj_config['type'] + factory = OBJECT_FACTORIES.get(obj_type, OBJECT_FACTORIES['Box']) + mj_obj = factory(obj_config) + mujoco_objects.append(mj_obj) + + return mujoco_objects + + def _setup_references(self): + """Setup references to simulation objects.""" + super()._setup_references() + self.obj_body_ids = { + obj.name: self.sim.model.body_name2id(obj.root_body) + for obj in self.mujoco_objects + } + + def _setup_observables(self) -> Dict: + """Setup observables for objects.""" + observables = super()._setup_observables() + + for obj in self.mujoco_objects: + @sensor(modality="object") + def obj_pos(obs_cache, name=obj.name): + return np.array(self.sim.data.body_xpos[self.obj_body_ids[name]]) + + obj_pos.__name__ = f"{obj.name}_pos" + observables[obj_pos.__name__] = Observable( + name=obj_pos.__name__, + sensor=obj_pos, + sampling_rate=self.control_freq + ) + + return observables + + def _reset_internal(self): + """Reset environment internals.""" + super()._reset_internal() + + # Set positions for all objects + for obj_config, mj_obj in zip(self.scenic_objects, self.mujoco_objects): + if 'position' in obj_config: + pos = obj_config['position'] + quat = obj_config.get('quaternion', [1, 0, 0, 0]) + self.sim.data.set_joint_qpos( + mj_obj.joints[0], + np.concatenate([np.array(pos), np.array(quat)]) + ) + + def reward(self, action=None) -> float: + """Compute reward.""" + return 0.0 + + def _check_success(self) -> bool: + """Check task success.""" + return False + + +class RobosuiteSimulator(Simulator): + """Simulator for RoboSuite custom environments.""" + + def __init__(self, render: bool = True, real_time: bool = True, speed: float = 1.0, + env_config: Optional[Dict] = None, controller_config: Optional[Dict] = None, + camera_view: Optional[str] = None, lite_physics: Optional[bool] = None): + super().__init__() + if suite is None: + raise RuntimeError(f"Unable to import RoboSuite: {_import_error}") + + self.render = render + self.real_time = real_time + self.speed = speed + self.env_config = env_config or {} + self.controller_config = controller_config + self.camera_view = camera_view + self.lite_physics = lite_physics + + def createSimulation(self, scene, **kwargs): + """Create a simulation instance.""" + return RobosuiteSimulation( + scene, self.render, self.real_time, self.speed, + self.env_config, self.controller_config, + self.camera_view, self.lite_physics, **kwargs + ) + + +class RobosuiteSimulation(Simulation): + """Simulation for RoboSuite custom environments.""" + + def __init__(self, scene, render: bool, real_time: bool, speed: float, + env_config: Optional[Dict], controller_config: Optional[Dict], + camera_view: Optional[str], lite_physics: Optional[bool], **kwargs): + self.render = render + self.real_time = real_time + self.speed = speed + self.env_config = env_config or {} + self.controller_config = controller_config + self.camera_view = camera_view + self.lite_physics = lite_physics + + # Environment state + self.robosuite_env = None + self.model = None + self.data = None + + # Object tracking + self.body_id_map = {} + self.object_name_map = {} + self.robots = [] + self.prev_positions = {} + self._current_obs = None + + # Action handling + self.pending_robot_action = None + self.action_dim = DEFAULT_ACTION_DIM + self.controller_type = None + + # Timing + self.timestep = kwargs.get('timestep') or DEFAULT_TIMESTEP + self.physics_timestep = kwargs.get('physics_timestep') or DEFAULT_PHYSICS_TIMESTEP + self.physics_steps = int(self.timestep / self.physics_timestep) + self.agents = [] + + super().__init__(scene, **kwargs) + + def setup(self): + """Initialize the RoboSuite environment.""" + super().setup() + + # Extract configuration from Scenic scene + scenic_config = self._extract_scenic_config() + + # Check for robots - at least one required + if not scenic_config['robots']: + raise ValueError("At least one robot is required in the scene") + + # Setup environment + robot_arg = self._get_robot_arg(scenic_config['robots']) + + env_kwargs = { + 'scenic_config': scenic_config, + 'robots': robot_arg, + 'has_renderer': self.render, + 'render_camera': self.camera_view, + 'camera_names': ["frontview", "robot0_eye_in_hand"], + 'controller_configs': self.controller_config, + **self.env_config + } + + if self.lite_physics is not None: + env_kwargs['lite_physics'] = self.lite_physics + + self.robosuite_env = ScenicManipulationEnv(**env_kwargs) + self._current_obs = self.robosuite_env.reset() + self._detect_controller_type(scenic_config['robots']) + + # Finalize setup + self._finalize_setup() + + def _finalize_setup(self): + """Common setup after environment creation.""" + self.model = self.robosuite_env.sim.model._model + self.data = self.robosuite_env.sim.data._data + + if self.render and self.camera_view is not None: + self._set_camera_view() + + self._setup_body_mapping() + + # Identify agents + self.agents = [obj for obj in self.objects + if hasattr(obj, 'behavior') and obj.behavior] + + def _set_camera_view(self): + """Set the camera view if specified.""" + camera_id = (CAMERA_VIEWS.get(self.camera_view.lower(), 0) + if isinstance(self.camera_view, str) + else self.camera_view) + if self.robosuite_env.viewer: + self.robosuite_env.viewer.set_camera(camera_id=camera_id) + + def _extract_scenic_config(self) -> Dict: + """Extract configuration from Scenic scene.""" + config = {'robots': [], 'tables': [], 'objects': []} + + for obj in self.objects: + if hasattr(obj, 'robot_type'): + self._add_robot_config(config['robots'], obj) + elif hasattr(obj, 'isTable') and obj.isTable: + self._add_table_config(config['tables'], obj) + elif hasattr(obj, 'objectType'): + self._add_object_config(config['objects'], obj) + + return config + + def _add_robot_config(self, robots: List, obj): + """Add robot configuration.""" + robots.append({ + 'type': obj.robot_type, + 'position': [obj.position.x, obj.position.y, obj.position.z] + }) + self.robots.append(obj) + + def _add_table_config(self, tables: List, obj): + """Add table configuration.""" + tables.append({ + 'position': [obj.position.x, obj.position.y, obj.position.z], + 'size': (obj.width, obj.length, obj.height) + }) + + def _add_object_config(self, objects: List, obj): + """Add object configuration.""" + # Use the Scenic object's name attribute directly + obj_name = getattr(obj, 'name', f"obj_{id(obj)}") + + config = { + 'type': obj.objectType, + 'name': obj_name, + 'position': [obj.position.x, obj.position.y, obj.position.z], + 'color': getattr(obj, 'color', [1, 0, 0, 1]) + } + + # Store mapping + self.object_name_map[obj] = obj_name + + if obj.objectType == 'Ball' and hasattr(obj, 'radius'): + config['radius'] = obj.radius + elif obj.objectType in ['Box', 'Cylinder', 'Capsule'] and hasattr(obj, 'width'): + config['size'] = [obj.width, obj.length, obj.height] + + objects.append(config) + + def _get_robot_arg(self, robots: List) -> Union[str, List[str]]: + """Get robot argument for environment creation.""" + if not robots: + raise ValueError("At least one robot is required") + return ([r['type'] for r in robots] if len(robots) > 1 + else robots[0]['type']) + + def _detect_controller_type(self, robot_configs: List): + """Detect controller type from first robot.""" + if robot_configs and self.robosuite_env.robots: + first_robot = self.robosuite_env.robots[0] + if hasattr(first_robot, 'controller'): + self.controller_type = type(first_robot.controller).__name__ + if hasattr(first_robot.controller, 'control_dim'): + self.action_dim = first_robot.controller.control_dim + + def _setup_body_mapping(self): + """Map environment objects to MuJoCo body IDs.""" + for obj in self.objects: + if obj in self.object_name_map: + self._map_object_body(obj) + + def _map_object_body(self, obj): + """Map single object to body ID.""" + obj_name = self.object_name_map.get(obj) + if not obj_name: + return + + for suffix in BODY_SUFFIXES: + body_name = f"{obj_name}{suffix}" + try: + body_id = mujoco.mj_name2id( + self.model, mujoco.mjtObj.mjOBJ_BODY, body_name + ) + if body_id != -1: + self.body_id_map[obj_name] = body_id + self.prev_positions[obj_name] = self.data.xpos[body_id].copy() + break + except: + continue + + def createObjectInSimulator(self, obj): + """Required by Scenic's Simulator interface.""" + # Objects are created during setup in ScenicManipulationEnv + pass + + def executeActions(self, allActions: Dict[Any, List]) -> None: + """Execute actions by calling their applyTo methods.""" + super().executeActions(allActions) + + self.pending_robot_action = None + + for agent in self.agents: + if agent in allActions and allActions[agent]: + for action in allActions[agent]: + if action and hasattr(action, 'applyTo'): + action.applyTo(agent, self) + break + + def step(self): + """Step the simulation forward one timestep.""" + if hasattr(self, '_done') and self._done: + return + + # Store previous positions + for name, body_id in self.body_id_map.items(): + self.prev_positions[name] = self.data.xpos[body_id].copy() + + action = (self.pending_robot_action if self.pending_robot_action is not None + else np.zeros(self.action_dim)) + + # Step simulation + for _ in range(self.physics_steps): + obs, reward, done, info = self.robosuite_env.step(action) + self._current_obs = obs + self._done = done + if self.render: + self.robosuite_env.render() + if done: + break + + self.pending_robot_action = None + + def getProperties(self, obj, properties: List[str]) -> Dict[str, Any]: + """Get current property values for an object.""" + values = {} + robot_idx = self.robots.index(obj) if obj in self.robots else None + + for prop in properties: + if prop == 'position': + values[prop] = self._get_object_position(obj) + elif prop == 'joint_positions' and robot_idx is not None: + values[prop] = self._get_robot_joints(robot_idx) + elif prop == 'eef_pos' and robot_idx is not None: + values[prop] = self._get_eef_position(robot_idx) + elif prop == 'gripper_state' and robot_idx is not None: + values[prop] = self._get_gripper_state(robot_idx) + else: + values[prop] = getattr(obj, prop, None) + + return values + + def _get_object_position(self, obj) -> Vector: + """Get object position.""" + obj_name = self.object_name_map.get(obj) + if obj_name and self._current_obs: + obs_key = f"{obj_name}_pos" + if obs_key in self._current_obs: + pos = self._current_obs[obs_key] + return Vector(pos[0], pos[1], pos[2]) + elif obj_name in self.body_id_map: + body_id = self.body_id_map[obj_name] + pos = self.data.xpos[body_id] + return Vector(pos[0], pos[1], pos[2]) + return obj.position + + def _get_robot_joints(self, robot_idx: int) -> List: + """Get robot joint positions.""" + if robot_idx < len(self.robosuite_env.robots): + return list(self.robosuite_env.robots[robot_idx]._joint_positions) + return [] + + def _get_eef_position(self, robot_idx: int) -> List: + """Get end-effector position.""" + if self._current_obs: + return list(self._current_obs.get(f'robot{robot_idx}_eef_pos', [0, 0, 0])) + return [0, 0, 0] + + def _get_gripper_state(self, robot_idx: int) -> List: + """Get gripper state.""" + if self._current_obs: + return list(self._current_obs.get(f'robot{robot_idx}_gripper_qpos', [0, 0])) + return [0, 0] + + def getCurrentObservation(self) -> Optional[Dict]: + """Get current observation dictionary.""" + return self._current_obs + + def checkSuccess(self) -> bool: + """Check if task is successfully completed.""" + if hasattr(self.robosuite_env, '_check_success'): + return self.robosuite_env._check_success() + return False + + def destroy(self): + """Clean up simulation resources.""" + if self.robosuite_env: + self.robosuite_env.close() + self.robosuite_env = None + + +# Actions +class SetJointPositions(Action): + """Set robot joint positions.""" + + def __init__(self, positions: List[float]): + self.positions = positions + + def applyTo(self, agent, sim: RobosuiteSimulation): + """Apply action to agent.""" + if hasattr(sim, 'robots') and agent in sim.robots: + robot_idx = sim.robots.index(agent) + if robot_idx < len(sim.robosuite_env.robots): + sim.pending_robot_action = np.array(self.positions) + + +class OSCPositionAction(Action): + """Operational Space Control for end-effector.""" + + def __init__(self, position_delta: Optional[List[float]] = None, + orientation_delta: Optional[List[float]] = None, + gripper: Optional[float] = None): + self.position_delta = position_delta or [0, 0, 0] + self.orientation_delta = orientation_delta or [0, 0, 0] + self.gripper = gripper if gripper is not None else 0 + + def applyTo(self, agent, sim: RobosuiteSimulation): + """Apply action to agent.""" + if hasattr(sim, 'robots') and agent in sim.robots: + robot_idx = sim.robots.index(agent) + if robot_idx < len(sim.robosuite_env.robots): + if sim.controller_type == 'JOINT_POSITION': + action = np.zeros(sim.action_dim) + action[:3] = self.position_delta + else: + action = np.zeros(7) + action[:3] = self.position_delta + action[3:6] = self.orientation_delta + action[6] = self.gripper + + sim.pending_robot_action = action \ No newline at end of file