mimicgen.envs.robosuite package#

Submodules#

mimicgen.envs.robosuite.coffee module#

class mimicgen.envs.robosuite.coffee.Coffee(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG

This class corresponds to the coffee task for a single robot arm.

Parameters

  • 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.

  • 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]

reward(action=None)#

Reward function for the task.

The sparse reward only consists of the threading component.

Note that the final reward is normalized and scaled by reward_scale / 2.0 as well so that the max score is equal to reward_scale

Parameters

action (np array) – [NOT USED]

Returns

reward value

Return type

float

visualize(vis_settings)#

In addition to super call, visualize gripper site proportional to the distance to the coffee machine.

Parameters

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.

class mimicgen.envs.robosuite.coffee.CoffeePreparation(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.coffee.Coffee

Harder coffee task where the task starts with materials in drawer and coffee machine closed. The robot needs to retrieve the coffee pod and mug from the drawer, open the coffee machine, place the pod and mug in the machine, and then close the lid.

class mimicgen.envs.robosuite.coffee.CoffeePreparation_D0(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.coffee.CoffeePreparation

Rename base class for convenience.

class mimicgen.envs.robosuite.coffee.CoffeePreparation_D1(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.coffee.CoffeePreparation_D0

Broader initialization for mug (whole right side of table, with rotation) and modest movement for coffee machine (some translation and rotation).

class mimicgen.envs.robosuite.coffee.Coffee_D0(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.coffee.Coffee

Rename base class for convenience.

class mimicgen.envs.robosuite.coffee.Coffee_D1(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.coffee.Coffee_D0

Wider initialization for pod and coffee machine.

class mimicgen.envs.robosuite.coffee.Coffee_D2(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.coffee.Coffee_D1

Similar to Coffee_D1, but put pod on the left, and machine on the right. Had to also move machine closer to robot (in x) to get kinematics to work out.

mimicgen.envs.robosuite.hammer_cleanup module#

Contains environments for BUDS hammer place task from robosuite task zoo repo. (https://github.com/ARISE-Initiative/robosuite-task-zoo)

class mimicgen.envs.robosuite.hammer_cleanup.HammerCleanup_D0(robot_init_qpos=None, **kwargs)#

Bases: robosuite_task_zoo.environments.manipulation.hammer_place.HammerPlaceEnv, mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG

Augment BUDS hammer place task for mimicgen.

edit_model_xml(xml_str)#

This function edits the model xml with custom changes, including resolving relative paths, applying changes retroactively to existing demonstration files, and other custom scripts. Environment subclasses should modify this function to add environment-specific xml editing features. :param xml_str: Mujoco sim demonstration XML file as string :type xml_str: str

Returns

Edited xml file as string

Return type

str

class mimicgen.envs.robosuite.hammer_cleanup.HammerCleanup_D1(robot_init_qpos=None, **kwargs)#

Bases: mimicgen.envs.robosuite.hammer_cleanup.HammerCleanup_D0

Move object and drawer with wide initialization. Note we had to make some objects movable that were fixtures before.

mimicgen.envs.robosuite.kitchen module#

Contains environments for BUDS kitchen task from robosuite task zoo repo. ((https://github.com/ARISE-Initiative/robosuite-task-zoo))

class mimicgen.envs.robosuite.kitchen.ButtonObjectNew(name, friction=None, damping=None)#

Bases: robosuite_task_zoo.models.kitchen.button_object.ButtonObject

Override some offsets for placement sampler.

property horizontal_radius#

Returns maximum distance from model root body to any radial point of the model.

Helps us put models programmatically without them flying away due to a huge initial contact force. Must be defined by subclass.

Returns

radius

Return type

float

class mimicgen.envs.robosuite.kitchen.Kitchen_D0(**kwargs)#

Bases: robosuite_task_zoo.environments.manipulation.kitchen.KitchenEnv, mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG

Augment BUDS kitchen task for mimicgen.

edit_model_xml(xml_str)#

This function edits the model xml with custom changes, including resolving relative paths, applying changes retroactively to existing demonstration files, and other custom scripts. Environment subclasses should modify this function to add environment-specific xml editing features. :param xml_str: Mujoco sim demonstration XML file as string :type xml_str: str

Returns

Edited xml file as string

Return type

str

class mimicgen.envs.robosuite.kitchen.Kitchen_D1(**kwargs)#

Bases: mimicgen.envs.robosuite.kitchen.Kitchen_D0

Specify wider distribution for objects including objects that didn’t move before. We also had to make some objects movable that were fixtures before.

visualize(vis_settings)#

Update site visualization to make sure stove object site visualization is kept up to date.

class mimicgen.envs.robosuite.kitchen.ServingRegionObjectNew(name, joints=None)#

Bases: robosuite.models.objects.objects.MujocoXMLObject

Override some offsets for placement sampler, and also turn the site into a visual-only geom so that it shows up in the first env step.

property horizontal_radius#

Returns maximum distance from model root body to any radial point of the model.

Helps us put models programmatically without them flying away due to a huge initial contact force. Must be defined by subclass.

Returns

radius

Return type

float

class mimicgen.envs.robosuite.kitchen.StoveObjectNew(name, joints=None)#

Bases: robosuite_task_zoo.models.kitchen.stove_object.StoveObject

Override some offsets for placement sampler.

property bottom_offset#

Returns vector from model root body to model bottom. Useful for, e.g. placing models on a surface. Must be defined by subclass.

Returns

(dx, dy, dz) offset vector

Return type

np.array

property horizontal_radius#

Returns maximum distance from model root body to any radial point of the model.

Helps us put models programmatically without them flying away due to a huge initial contact force. Must be defined by subclass.

Returns

radius

Return type

float

property top_offset#

Returns vector from model root body to model top. Useful for, e.g. placing models on a surface. Must be defined by subclass.

Returns

(dx, dy, dz) offset vector

Return type

np.array

mimicgen.envs.robosuite.mug_cleanup module#

Simpler object cleanup task (inspired by BUDS Hammer Place, see https://github.com/ARISE-Initiative/robosuite-task-zoo) where a single object needs to be packed away into a drawer. The default task is to cleanup a particular mug.

class mimicgen.envs.robosuite.mug_cleanup.MugCleanup(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None, shapenet_id='3143a4ac', shapenet_scale=0.8)#

Bases: mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG

This class corresponds to the object cleanup task for a single robot arm.

Parameters

  • 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.

  • 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]

reward(action=None)#

Reward function for the task.

The sparse reward only consists of the threading component.

Note that the final reward is normalized and scaled by reward_scale / 2.0 as well so that the max score is equal to reward_scale

Parameters

action (np array) – [NOT USED]

Returns

reward value

Return type

float

visualize(vis_settings)#

In addition to super call, visualize gripper site proportional to the distance to the coffee machine.

Parameters

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.

class mimicgen.envs.robosuite.mug_cleanup.MugCleanup_D0(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None, shapenet_id='3143a4ac', shapenet_scale=0.8)#

Bases: mimicgen.envs.robosuite.mug_cleanup.MugCleanup

Rename base class for convenience.

class mimicgen.envs.robosuite.mug_cleanup.MugCleanup_D1(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None, shapenet_id='3143a4ac', shapenet_scale=0.8)#

Bases: mimicgen.envs.robosuite.mug_cleanup.MugCleanup_D0

Wider initialization for both drawer and object.

class mimicgen.envs.robosuite.mug_cleanup.MugCleanup_O1(**kwargs)#

Bases: mimicgen.envs.robosuite.mug_cleanup.MugCleanup_D0

Use different mug.

class mimicgen.envs.robosuite.mug_cleanup.MugCleanup_O2(**kwargs)#

Bases: mimicgen.envs.robosuite.mug_cleanup.MugCleanup_D0

Use a random mug on each episode reset.

mimicgen.envs.robosuite.nut_assembly module#

class mimicgen.envs.robosuite.nut_assembly.NutAssembly_D0(**kwargs)#

Bases: robosuite.environments.manipulation.nut_assembly.NutAssembly, mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG

Augment robosuite nut assembly task for mimicgen.

edit_model_xml(xml_str)#

This function edits the model xml with custom changes, including resolving relative paths, applying changes retroactively to existing demonstration files, and other custom scripts. Environment subclasses should modify this function to add environment-specific xml editing features. :param xml_str: Mujoco sim demonstration XML file as string :type xml_str: str

Returns

Edited xml file as string

Return type

str

class mimicgen.envs.robosuite.nut_assembly.Square_D0(**kwargs)#

Bases: robosuite.environments.manipulation.nut_assembly.NutAssemblySquare, mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG

Augment robosuite nut assembly square task for mimicgen.

edit_model_xml(xml_str)#

This function edits the model xml with custom changes, including resolving relative paths, applying changes retroactively to existing demonstration files, and other custom scripts. Environment subclasses should modify this function to add environment-specific xml editing features. :param xml_str: Mujoco sim demonstration XML file as string :type xml_str: str

Returns

Edited xml file as string

Return type

str

class mimicgen.envs.robosuite.nut_assembly.Square_D1(**kwargs)#

Bases: mimicgen.envs.robosuite.nut_assembly.Square_D0

Specifies a different placement initializer for the pegs where it is initialized with a broader x-range and broader y-range.

class mimicgen.envs.robosuite.nut_assembly.Square_D2(**kwargs)#

Bases: mimicgen.envs.robosuite.nut_assembly.Square_D1

Even broader range for everything, and z-rotation randomization for peg.

mimicgen.envs.robosuite.pick_place module#

Slight variant of pick place task.

class mimicgen.envs.robosuite.pick_place.PickPlace_D0(**kwargs)#

Bases: robosuite.environments.manipulation.pick_place.PickPlace

Slightly easier task where we limit z-rotation to 0 to 90 degrees for all object initializations (instead of full 360).

mimicgen.envs.robosuite.single_arm_env_mg module#

class mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG(robots, env_configuration='default', controller_configs=None, mount_types='default', gripper_types='default', initialization_noise=None, use_camera_obs=True, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: robosuite.environments.manipulation.single_arm_env.SingleArmEnv

Custom version of base class for single arm robosuite tasks for mimicgen.

edit_model_xml(xml_str)#

This function edits the model xml with custom changes, including resolving relative paths, applying changes retroactively to existing demonstration files, and other custom scripts. Environment subclasses should modify this function to add environment-specific xml editing features. :param xml_str: Mujoco sim demonstration XML file as string :type xml_str: str

Returns

Edited xml file as string

Return type

str

mimicgen.envs.robosuite.stack module#

class mimicgen.envs.robosuite.stack.StackThree(**kwargs)#

Bases: mimicgen.envs.robosuite.stack.Stack_D0

Stack three cubes instead of two.

reward(action=None)#

We only return sparse rewards here.

staged_rewards()#

Helper function to calculate staged rewards based on current physical states. :returns: - (float): reward for reaching and grasping

  • (float): reward for lifting and aligning

  • (float): reward for stacking

Return type

3-tuple

class mimicgen.envs.robosuite.stack.StackThree_D0(**kwargs)#

Bases: mimicgen.envs.robosuite.stack.StackThree

Rename base class for convenience.

class mimicgen.envs.robosuite.stack.StackThree_D1(**kwargs)#

Bases: mimicgen.envs.robosuite.stack.StackThree_D0

Less z-rotation (for easier datagen) and much wider initialization bounds.

class mimicgen.envs.robosuite.stack.Stack_D0(**kwargs)#

Bases: robosuite.environments.manipulation.stack.Stack, mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG

Augment robosuite stack task for mimicgen.

edit_model_xml(xml_str)#

This function edits the model xml with custom changes, including resolving relative paths, applying changes retroactively to existing demonstration files, and other custom scripts. Environment subclasses should modify this function to add environment-specific xml editing features. :param xml_str: Mujoco sim demonstration XML file as string :type xml_str: str

Returns

Edited xml file as string

Return type

str

reward(action=None)#

Reward function for the task.

Sparse un-normalized reward:

  • a discrete reward of 2.0 is provided if the red block is stacked on the green block

Un-normalized components if using reward shaping:

  • Reaching: in [0, 0.25], to encourage the arm to reach the cube

  • Grasping: in {0, 0.25}, non-zero if arm is grasping the cube

  • Lifting: in {0, 1}, non-zero if arm has lifted the cube

  • Aligning: in [0, 0.5], encourages aligning one cube over the other

  • Stacking: in {0, 2}, non-zero if cube is stacked on other cube

The reward is max over the following:

  • Reaching + Grasping

  • Lifting + Aligning

  • Stacking

The sparse reward only consists of the stacking component.

Note that the final reward is normalized and scaled by reward_scale / 2.0 as well so that the max score is equal to reward_scale

Parameters

action (np array) – [NOT USED]

Returns

reward value

Return type

float

class mimicgen.envs.robosuite.stack.Stack_D1(**kwargs)#

Bases: mimicgen.envs.robosuite.stack.Stack_D0

Much wider initialization bounds.

mimicgen.envs.robosuite.threading module#

class mimicgen.envs.robosuite.threading.Threading(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG

This class corresponds to the threading task for a single robot arm.

Parameters

  • 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.

  • 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]

reward(action=None)#

Reward function for the task.

The sparse reward only consists of the threading component.

Note that the final reward is normalized and scaled by reward_scale / 2.0 as well so that the max score is equal to reward_scale

Parameters

action (np array) – [NOT USED]

Returns

reward value

Return type

float

visualize(vis_settings)#

In addition to super call, visualize gripper site proportional to the distance to the needle.

Parameters

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.

class mimicgen.envs.robosuite.threading.Threading_D0(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.threading.Threading

Rename base class for convenience.

class mimicgen.envs.robosuite.threading.Threading_D1(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.threading.Threading_D0

Specifies a wider placement initializer for the needle and the tripod.

class mimicgen.envs.robosuite.threading.Threading_D2(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.threading.Threading_D1

Same as Target 1, but put needle on left and tripod on right.

mimicgen.envs.robosuite.three_piece_assembly module#

class mimicgen.envs.robosuite.three_piece_assembly.ThreePieceAssembly(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.single_arm_env_mg.SingleArmEnv_MG

This class corresponds to the three piece assembly task for a single robot arm.

Parameters

  • 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.

  • 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]

reward(action=None)#

Reward function for the task.

The sparse reward only consists of the threading component.

Note that the final reward is normalized and scaled by reward_scale / 2.0 as well so that the max score is equal to reward_scale

Parameters

action (np array) – [NOT USED]

Returns

reward value

Return type

float

visualize(vis_settings)#

In addition to super call, visualize gripper site proportional to the distance to the needle.

Parameters

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.

class mimicgen.envs.robosuite.three_piece_assembly.ThreePieceAssembly_D0(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.three_piece_assembly.ThreePieceAssembly

Rename base class for convenience.

class mimicgen.envs.robosuite.three_piece_assembly.ThreePieceAssembly_D1(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.three_piece_assembly.ThreePieceAssembly_D0

All pieces still have fixed z-rotation, as in original task, but base can be placed anywhere.

class mimicgen.envs.robosuite.three_piece_assembly.ThreePieceAssembly_D2(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, 0.005, 0.0001), use_camera_obs=True, use_object_obs=True, reward_scale=1.0, reward_shaping=False, 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='agentview', camera_heights=256, camera_widths=256, camera_depths=False, camera_segmentations=None, renderer='mujoco', renderer_config=None)#

Bases: mimicgen.envs.robosuite.three_piece_assembly.ThreePieceAssembly_D1

All pieces can move anywhere on table and in z-rotation.

Module contents#