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: robosuite.wrappers.wrapper.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: robosuite.wrappers.wrapper.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: robosuite.wrappers.wrapper.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)

Bases: robosuite.wrappers.wrapper.Wrapper, gym.core.Env

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.

Raises

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

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

reset()

Extends env reset method to return flattened observation instead of normal OrderedDict.

Returns

Flattened environment observation space after reset occurs

Return type

np.array

seed(seed=None)

Utility function to set numpy seed

Parameters

seed (None or int) – If specified, numpy seed to set

Raises

TypeError – [Seed must be integer]

step(action)

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

Parameters

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

Returns

  • (np.array) flattened 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.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: robosuite.wrappers.wrapper.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