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