robosuite.wrappers package#

Submodules#

robosuite.wrappers.data_collection_wrapper module#

This file implements a wrapper for saving simulation states to disk. This data collection wrapper is useful for collecting demonstrations.

class robosuite.wrappers.data_collection_wrapper.DataCollectionWrapper(env, directory, collect_freq=1, flush_freq=100)#

Bases: Wrapper

close()#

Override close method in order to flush left over data

reset()#

Extends vanilla reset() function call to accommodate data collection

Returns:

Environment observation space after reset occurs

Return type:

OrderedDict

step(action)#

Extends vanilla step() function call to accommodate data collection

Parameters:

action (np.array) – Action to take in environment

Returns:

  • (OrderedDict) observations from the environment

  • (float) reward from the environment

  • (bool) whether the current episode is completed or not

  • (dict) misc information

Return type:

4-tuple

robosuite.wrappers.demo_sampler_wrapper module#

This file contains a wrapper for sampling environment states from a set of demonstrations on every reset. The main use case is for altering the start state distribution of training episodes for learning RL policies.

class robosuite.wrappers.demo_sampler_wrapper.DemoSamplerWrapper(env, demo_path, need_xml=False, num_traj=-1, sampling_schemes=('uniform', 'random'), scheme_ratios=(0.9, 0.1), open_loop_increment_freq=100, open_loop_initial_window_width=25, open_loop_window_increment=25)#

Bases: Wrapper

Initializes a wrapper that provides support for resetting the environment state to one from a demonstration. It also supports curriculums for altering how often to sample from demonstration vs. sampling a reset state from the environment.

Parameters:
  • env (MujocoEnv) – The environment to wrap.

  • demo_path (str) – The path to the folder containing the demonstrations. There should be a demo.hdf5 file and a folder named models with all of the stored model xml files from the demonstrations.

  • need_xml (bool) – If True, the mujoco model needs to be reloaded when sampling a state from a demonstration. This could be because every demonstration was taken under varied object properties, for example. In this case, every sampled state comes with a corresponding xml to be used for the environment reset.

  • num_traj (int) – If provided, subsample @number demonstrations from the provided set of demonstrations instead of using all of them.

  • sampling_schemes (list of str) –

    A list of sampling schemes to be used. The following strings are valid schemes:

    ’random’: sample a reset state directly from the wrapped environment

    ’uniform’: sample a state from a demonstration uniformly at random

    ’forward’: sample a state from a window that grows progressively from

    the start of demonstrations

    ’reverse’: sample a state from a window that grows progressively from

    the end of demonstrations

  • scheme_ratios (list of float --> np.array) – A list of probability values to assign to each member of @sampling_schemes. Must be non-negative and sum to 1.

  • open_loop_increment_freq (int) – How frequently to increase the window size in open loop schemes (“forward” and “reverse”). The window size will increase by @open_loop_window_increment every @open_loop_increment_freq samples. Only samples that are generated by open loop schemes contribute to this count.

  • open_loop_initial_window_width (int) – The width of the initial sampling window, in terms of number of demonstration time steps, for open loop schemes.

  • open_loop_window_increment (int) – The window size will increase by @open_loop_window_increment every @open_loop_increment_freq samples. This number is in terms of number of demonstration time steps.

Raises:
  • AssertionError – [Incompatible envs]

  • AssertionError – [Invalid sampling scheme]

  • AssertionError – [Invalid scheme ratio]

reset()#

Logic for sampling a state from the demonstration and resetting the simulation to that state.

Returns:

Environment observation space after reset occurs

Return type:

OrderedDict

sample()#

This is the core sampling method. Samples a state from a demonstration, in accordance with the configuration.

Returns:

If np.array, is the state sampled from a demo file. If 2-tuple, additionally

includes the model xml file

Return type:

None or np.array or 2-tuple

robosuite.wrappers.domain_randomization_wrapper module#

This file implements a wrapper for facilitating domain randomization over robosuite environments.

class robosuite.wrappers.domain_randomization_wrapper.DomainRandomizationWrapper(env, seed=None, randomize_color=True, randomize_camera=True, randomize_lighting=True, randomize_dynamics=True, color_randomization_args={'geom_names': None, 'local_material_interpolation': 0.3, 'local_rgb_interpolation': 0.2, 'randomize_local': True, 'randomize_material': True, 'randomize_skybox': True, 'texture_variations': ['rgb', 'checker', 'noise', 'gradient']}, camera_randomization_args={'camera_names': None, 'fovy_perturbation_size': 5.0, 'position_perturbation_size': 0.01, 'randomize_fovy': True, 'randomize_position': True, 'randomize_rotation': True, 'rotation_perturbation_size': 0.087}, lighting_randomization_args={'ambient_perturbation_size': 0.1, 'diffuse_perturbation_size': 0.1, 'direction_perturbation_size': 0.35, 'light_names': None, 'position_perturbation_size': 0.1, 'randomize_active': True, 'randomize_ambient': True, 'randomize_diffuse': True, 'randomize_direction': True, 'randomize_position': True, 'randomize_specular': True, 'specular_perturbation_size': 0.1}, dynamics_randomization_args={'armature_perturbation_size': 0.01, 'body_names': None, 'damping_perturbation_size': 0.01, 'density_perturbation_ratio': 0.1, 'friction_perturbation_ratio': 0.1, 'frictionloss_perturbation_size': 0.05, 'geom_names': None, 'inertia_perturbation_ratio': 0.02, 'joint_names': None, 'mass_perturbation_ratio': 0.02, 'position_perturbation_size': 0.0015, 'quaternion_perturbation_size': 0.003, 'randomize_armature': True, 'randomize_damping': True, 'randomize_density': True, 'randomize_friction': True, 'randomize_frictionloss': True, 'randomize_inertia': True, 'randomize_mass': True, 'randomize_position': True, 'randomize_quaternion': True, 'randomize_solimp': True, 'randomize_solref': True, 'randomize_stiffness': True, 'randomize_viscosity': True, 'solimp_perturbation_ratio': 0.1, 'solref_perturbation_ratio': 0.1, 'stiffness_perturbation_ratio': 0.1, 'viscosity_perturbation_ratio': 0.1}, randomize_on_reset=True, randomize_every_n_steps=1)#

Bases: Wrapper

Wrapper that allows for domain randomization mid-simulation.

Parameters:
  • env (MujocoEnv) – The environment to wrap.

  • seed (int) – Integer used to seed all randomizations from this wrapper. It is used to create a np.random.RandomState instance to make sure samples here are isolated from sampling occurring elsewhere in the code. If not provided, will default to using global random state.

  • randomize_color (bool) – if True, randomize geom colors and texture colors

  • randomize_camera (bool) – if True, randomize camera locations and parameters

  • randomize_lighting (bool) – if True, randomize light locations and properties

  • randomize_dyanmics (bool) – if True, randomize dynamics parameters

  • color_randomization_args (dict) – Color-specific randomization arguments

  • camera_randomization_args (dict) – Camera-specific randomization arguments

  • lighting_randomization_args (dict) – Lighting-specific randomization arguments

  • dynamics_randomization_args (dict) – Dyanmics-specific randomization arguments

  • randomize_on_reset (bool) – if True, randomize on every call to @reset. This, in conjunction with setting @randomize_every_n_steps to 0, is useful to generate a new domain per episode.

  • randomize_every_n_steps (int) – determines how often randomization should occur. Set to 0 if randomization should happen manually (by calling @randomize_domain)

randomize_domain()#

Runs domain randomization over the environment.

reset()#

Extends superclass method to reset the domain randomizer.

Returns:

Environment observation space after reset occurs

Return type:

OrderedDict

restore_default_domain()#

Restores the simulation model parameters saved in the last call to @save_default_domain.

save_default_domain()#

Saves the current simulation model parameters so that they can be restored later.

step(action)#

Extends vanilla step() function call to accommodate domain randomization

Returns:

  • (OrderedDict) observations from the environment

  • (float) reward from the environment

  • (bool) whether the current episode is completed or not

  • (dict) misc information

Return type:

4-tuple

step_randomization()#

Steps the internal randomization state

robosuite.wrappers.gym_wrapper module#

This file implements a wrapper for facilitating compatibility with OpenAI gym. This is useful when using these environments with code that assumes a gym-like interface.

class robosuite.wrappers.gym_wrapper.GymWrapper(env, keys=None, flatten_obs=True)#

Bases: Wrapper, Env

close()#

wrapper for calling underlying env close function

compute_reward(achieved_goal, desired_goal, info)#

Dummy function to be compatible with gym interface that simply returns environment reward

Parameters:
  • achieved_goal – [NOT USED]

  • desired_goal – [NOT USED]

  • info – [NOT USED]

Returns:

environment reward

Return type:

float

metadata: dict[str, Any] = None#
render_mode: str | None = None#

Initializes the Gym wrapper. Mimics many of the required functionalities of the Wrapper class found in the gym.core module

Parameters:
  • env (MujocoEnv) – The environment to wrap.

  • keys (None or list of str) – If provided, each observation will consist of concatenated keys from the wrapped environment’s observation dictionary. Defaults to proprio-state and object-state.

  • flatten_obs (bool) – Whether to flatten the observation dictionary into a 1d array. Defaults to True.

Raises:

AssertionError – [Object observations must be enabled if no keys]

reset(seed=None, options=None)#

Extends env reset method to return observation instead of normal OrderedDict and optionally resets seed

Returns:

  • (np.array) observations from the environment

  • (dict) an empty dictionary, as part of the standard return format

Return type:

2-tuple

step(action)#

Extends vanilla step() function call to return observation instead of normal OrderedDict.

Parameters:

action (np.array) – Action to take in environment

Returns:

  • (np.array) observations from the environment

  • (float) reward from the environment

  • (bool) episode ending after reaching an env terminal state

  • (bool) episode ending after an externally defined condition

  • (dict) misc information

Return type:

4-tuple

robosuite.wrappers.visualization_wrapper module#

This file implements a wrapper for visualizing important sites in a given environment.

By default, this visualizes all sites possible for the environment. Visualization options for a given environment can be found by calling get_visualization_settings(), and can be set individually by calling set_visualization_setting(setting, visible).

class robosuite.wrappers.visualization_wrapper.VisualizationWrapper(env, indicator_configs=None)#

Bases: Wrapper

get_indicator_names()#

Gets all indicator object names for this environment.

Returns:

Indicator names for this environment.

Return type:

list

get_visualization_settings()#

Gets all settings for visualizing this environment

Returns:

Visualization keywords for this environment.

Return type:

list

reset()#

Extends vanilla reset() function call to accommodate visualization

Returns:

Environment observation space after reset occurs

Return type:

OrderedDict

set_indicator_pos(indicator, pos)#

Sets the specified @indicator to the desired position @pos

Parameters:
  • indicator (str) – Name of the indicator to set

  • pos (3-array) – (x, y, z) Cartesian world coordinates to set the specified indicator to

set_visualization_setting(setting, visible)#

Sets the specified @setting to have visibility = @visible.

Parameters:
  • setting (str) – Visualization keyword to set

  • visible (bool) – True if setting should be visualized.

step(action)#

Extends vanilla step() function call to accommodate visualization

Parameters:

action (np.array) – Action to take in environment

Returns:

  • (OrderedDict) observations from the environment

  • (float) reward from the environment

  • (bool) whether the current episode is completed or not

  • (dict) misc information

Return type:

4-tuple

robosuite.wrappers.wrapper module#

This file contains the base wrapper class for Mujoco environments. Wrappers are useful for data collection and logging. Highly recommended.

class robosuite.wrappers.wrapper.Wrapper(env)#

Bases: object

Base class for all wrappers in robosuite.

Parameters:

env (MujocoEnv) – The environment to wrap.

property action_dim#

By default, grabs the normal environment action_dim

Returns:

Action space dimension

Return type:

int

property action_spec#

By default, grabs the normal environment action_spec

Returns:

  • (np.array) minimum (low) action values

  • (np.array) maximum (high) action values

Return type:

2-tuple

classmethod class_name()#
observation_spec()#

By default, grabs the normal environment observation_spec

Returns:

Observations from the environment

Return type:

OrderedDict

render(**kwargs)#

By default, run the normal environment render() function

Parameters:

**kwargs (dict) – Any args to pass to environment render function

reset()#

By default, run the normal environment reset() function

Returns:

Environment observation space after reset occurs

Return type:

OrderedDict

step(action)#

By default, run the normal environment step() function

Parameters:

action (np.array) – action to take in environment

Returns:

  • (OrderedDict) observations from the environment

  • (float) reward from the environment

  • (bool) whether the current episode is completed or not

  • (dict) misc information

Return type:

4-tuple

property unwrapped#

Grabs unwrapped environment

Returns:

Unwrapped environment

Return type:

env (MujocoEnv)

Module contents#