import numpy as np from robosuite.environments.manipulation.single_arm_env import SingleArmEnv from robosuite.models.arenas import PhantomTableArena from robosuite.models.tasks import ManipulationTask class Phantom(SingleArmEnv): """ This class corresponds to the stacking task for a single robot arm. Args: robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms) Note: Must be a single single-arm robot! env_configuration (str): Specifies how to position the robots within the environment (default is "default"). For most single arm environments, this argument has no impact on the robot setup. controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a custom controller. Else, uses the default controller for this specific task. Should either be single dict if same controller is to be used for all robots or else it should be a list of the same length as "robots" param gripper_types (str or list of str): type of gripper, used to instantiate gripper models from gripper factory. Default is "default", which is the default grippers(s) associated with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model overrides the default gripper. Should either be single str if same gripper type is to be used for all robots or else it should be a list of the same length as "robots" param initialization_noise (dict or list of dict): Dict containing the initialization noise parameters. The expected keys and corresponding value types are specified below: :`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial joint positions. Setting this value to `None` or 0.0 results in no noise being applied. If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied, If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range :`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform" Should either be single dict if same noise value is to be used for all robots or else it should be a list of the same length as "robots" param :Note: Specifying "default" will automatically use the default noise settings. Specifying None will automatically create the required dict with "magnitude" set to 0.0. table_full_size (3-tuple): x, y, and z dimensions of the table. table_friction (3-tuple): the three mujoco friction parameters for the table. use_camera_obs (bool): if True, every observation includes rendered image(s) use_object_obs (bool): if True, include object (cube) information in the observation. reward_scale (None or float): Scales the normalized reward function by the amount specified. If None, environment reward remains unnormalized reward_shaping (bool): if True, use dense rewards. placement_initializer (ObjectPositionSampler): if provided, will be used to place objects on every reset, else a UniformRandomSampler is used by default. has_renderer (bool): If true, render the simulation state in a viewer instead of headless mode. has_offscreen_renderer (bool): True if using off-screen rendering render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None' will result in the default angle being applied, which is useful as it can be dragged / panned by the user using the mouse render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise. render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise. render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering. Defaults to -1, in which case the device will be inferred from environment variables (GPUS or CUDA_VISIBLE_DEVICES). control_freq (float): how many control signals to receive in every second. This sets the amount of simulation time that passes between every action input. horizon (int): Every episode lasts for exactly @horizon timesteps. ignore_done (bool): True if never terminating the environment (ignore @horizon). hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else, only calls sim.reset and resets all robosuite-internal variables camera_names (str or list of str): name of camera to be rendered. Should either be single str if same name is to be used for all cameras' rendering or else it should be a list of cameras to render. :Note: At least one camera must be specified if @use_camera_obs is True. :Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each robot's camera list). camera_heights (int or list of int): height of camera frame. Should either be single int if same height is to be used for all cameras' frames or else it should be a list of the same length as "camera names" param. camera_widths (int or list of int): width of camera frame. Should either be single int if same width is to be used for all cameras' frames or else it should be a list of the same length as "camera names" param. camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single bool if same depth setting is to be used for all cameras or else it should be a list of the same length as "camera names" param. camera_segmentations (None or str or list of str or list of list of str): Camera segmentation(s) to use for each camera. Valid options are: `None`: no segmentation sensor used `'instance'`: segmentation at the class-instance level `'class'`: segmentation at the class level `'element'`: segmentation at the per-geom level If not None, multiple types of segmentations can be specified. A [list of str / str or None] specifies [multiple / a single] segmentation(s) to use for all cameras. A list of list of str specifies per-camera segmentation setting(s) to use. Raises: AssertionError: [Invalid number of robots specified] """ def __init__( self, robots, env_configuration="default", controller_configs=None, gripper_types="default", initialization_noise="default", table_full_size=(0.8, 0.8, 0.05), table_friction=(1.0, 5e-3, 1e-4), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, placement_initializer=None, has_renderer=False, has_offscreen_renderer=True, render_camera="frontview", render_collision_mesh=False, render_visual_mesh=True, render_gpu_device_id=-1, control_freq=20, horizon=1000, ignore_done=False, hard_reset=True, camera_names="frontview", camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, # {None, instance, class, element} renderer="mujoco", renderer_config=None, object_placements=None, direct_gripper_control=False, camera_pos=None, camera_quat_wxyz=None, camera_fov=None, camera_sensorsize=None, camera_principalpixel=None, camera_focalpixel=None, ): # settings for table top self.table_full_size = table_full_size self.table_friction = table_friction self.table_offset = np.array((0, 0, 0.8)) # reward configuration self.reward_scale = reward_scale self.reward_shaping = reward_shaping # whether to use ground-truth object states self.use_object_obs = use_object_obs # object placement initializer self.placement_initializer = placement_initializer self.object_placements = object_placements self.camera_pos = camera_pos self.camera_quat_wxyz = camera_quat_wxyz self.camera_fov = camera_fov self.camera_sensorsize = camera_sensorsize self.camera_principalpixel = camera_principalpixel self.camera_focalpixel = camera_focalpixel # pdb.set_trace() super().__init__( robots=robots, env_configuration=env_configuration, controller_configs=controller_configs, mount_types="default", gripper_types=gripper_types, initialization_noise=initialization_noise, use_camera_obs=use_camera_obs, has_renderer=has_renderer, has_offscreen_renderer=has_offscreen_renderer, render_camera=render_camera, render_collision_mesh=render_collision_mesh, render_visual_mesh=render_visual_mesh, render_gpu_device_id=render_gpu_device_id, control_freq=control_freq, horizon=horizon, ignore_done=ignore_done, hard_reset=hard_reset, camera_names=camera_names, camera_heights=camera_heights, camera_widths=camera_widths, camera_depths=camera_depths, camera_segmentations=camera_segmentations, renderer=renderer, renderer_config=renderer_config, direct_gripper_control=direct_gripper_control, ) def reset(self, object_placements=None): self.object_placements = object_placements return super().reset() def _load_model(self): """ Loads an xml model, puts it in self.model """ super()._load_model() # Adjust base pose accordingly xpos = self.robots[0].robot_model.base_xpos_offset["table"](self.table_full_size[0]) self.robots[0].robot_model.set_base_xpos(xpos) # load model for table top workspace mujoco_arena = PhantomTableArena( table_full_size=self.table_full_size, table_friction=self.table_friction, table_offset=self.table_offset, ) # Arena always gets set to zero origin mujoco_arena.set_origin([0, 0, 0]) # task includes arena, robot, and objects of interest self.model = ManipulationTask( mujoco_arena=mujoco_arena, mujoco_robots=[robot.robot_model for robot in self.robots], ) # Modify default frontview camera if self.camera_pos is not None: robot_base_pos = np.array([-0.56, 0, 0.912]) mujoco_arena.set_camera( camera_name="frontview", pos=self.camera_pos + robot_base_pos, quat=self.camera_quat_wxyz, camera_attribs={"sensorsize": np.array2string(self.camera_sensorsize)[1:-1], "resolution": f"{self.camera_widths[0]} {self.camera_heights[0]}", "principalpixel": np.array2string(self.camera_principalpixel)[1:-1], "focalpixel": np.array2string(self.camera_focalpixel)[1:-1],} ) def _setup_references(self): """ Sets up references to important components. A reference is typically an index or a list of indices that point to the corresponding elements in a flatten array, which is how MuJoCo stores physical simulation data. """ super()._setup_references() def _reset_internal(self): """ Resets simulation internal configurations. """ super()._reset_internal() def visualize(self, vis_settings): """ In addition to super call, visualize gripper site proportional to the distance to the cube. Args: vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific component should be visualized. Should have "grippers" keyword as well as any other relevant options specified. """ # Run superclass method first super().visualize(vis_settings=vis_settings) # # Color the gripper visualization site according to its distance to the cube # if vis_settings["grippers"]: # self._visualize_gripper_to_target(gripper=self.robots[0].gripper, target=self.cubeA) def reward(self, action): return 0.0