robomimic.envs package#

Submodules#

robomimic.envs.env_base module#

This file contains the base class for environment wrappers that are used to provide a standardized environment API for training policies and interacting with metadata present in datasets.

class robomimic.envs.env_base.EnvBase(env_name, render=False, render_offscreen=False, use_image_obs=False, use_depth_obs=False, postprocess_visual_obs=True, **kwargs)#

Bases: abc.ABC

A base class method for environments used by this repo.

abstract property action_dimension#

Returns dimension of actions (int).

abstract property base_env#

Grabs base simulation environment.

abstract classmethod create_for_data_processing(camera_names, camera_height, camera_width, reward_shaping, render=None, render_offscreen=None, use_image_obs=None, use_depth_obs=None, **kwargs)#

Create environment for processing datasets, which includes extracting observations, labeling dense / sparse rewards, and annotating dones in transitions.

Parameters:

  • camera_names ([str]) – list of camera names that correspond to image observations

  • camera_height (int) – camera height for all cameras

  • camera_width (int) – camera width for all cameras

  • reward_shaping (bool) – if True, use shaped environment rewards, else use sparse task completion rewards

  • render (bool or None) – optionally override rendering behavior. Defaults to False.

  • render_offscreen (bool or None) – optionally override rendering behavior. The default value is True if @camera_names is non-empty, False otherwise.

  • use_image_obs (bool or None) – optionally override rendering behavior. The default value is True if @camera_names is non-empty, False otherwise.

  • use_depth_obs (bool) – if True, use depth observations

Returns:

env (EnvBase instance)

abstract get_goal()#

Get goal observation. Not all environments support this.

abstract get_observation()#

Get environment observation

abstract get_reward()#

Get current reward.

abstract get_state()#

Get environment simulator state, compatible with @reset_to

abstract is_done()#

Check if the task is done (not necessarily successful).

abstract is_success()#

Check if the task condition(s) is reached. Should return a dictionary { str: bool } with at least a “task” key for the overall task success, and additional optional keys corresponding to other task criteria.

abstract property name#

Returns name of environment name (str).

abstract render(mode='human', height=None, width=None, camera_name=None)#

Render

abstract reset()#

Reset environment.

Returns:

initial observation dictionary.

Return type:

observation (dict)

abstract reset_to(state)#

Reset to a specific simulator state.

Parameters:

state (dict) – current simulator state

Returns:

observation dictionary after setting the simulator state

Return type:

observation (dict)

abstract property rollout_exceptions#

Return tuple of exceptions to except when doing rollouts. This is useful to ensure that the entire training run doesn’t crash because of a bad policy that causes unstable simulation computations.

abstract serialize()#

Save all information needed to re-instantiate this environment in a dictionary. This is the same as @env_meta - environment metadata stored in hdf5 datasets, and used in utils/env_utils.py.

abstract set_goal(**kwargs)#

Set goal observation with external specification. Not all environments support this.

abstract step(action)#

Step in the environment with an action.

Parameters:

action (np.array) – action to take

Returns:

new observation dictionary reward (float): reward for this step done (bool): whether the task is done info (dict): extra information

Return type:

observation (dict)

abstract property type#

Returns environment type (int) for this kind of environment. This helps identify this env class.

property version#

Returns version of environment (str). This is not an abstract method, some subclasses do not implement it

class robomimic.envs.env_base.EnvType#

Bases: object

Holds environment types - one per environment class. These act as identifiers for different environments.

GYM_TYPE = 2#
IG_MOMART_TYPE = 3#
ROBOSUITE_TYPE = 1#

robomimic.envs.env_gym module#

robomimic.envs.env_ig_momart module#

robomimic.envs.env_robosuite module#

This file contains the robosuite environment wrapper that is used to provide a standardized environment API for training policies and interacting with metadata present in datasets.

class robomimic.envs.env_robosuite.EnvRobosuite(env_name, render=False, render_offscreen=False, use_image_obs=False, use_depth_obs=False, postprocess_visual_obs=True, **kwargs)#

Bases: robomimic.envs.env_base.EnvBase

Wrapper class for robosuite environments (https://github.com/ARISE-Initiative/robosuite)

property action_dimension#

Returns dimension of actions (int).

property base_env#

Grabs base simulation environment.

classmethod create_for_data_processing(env_name, camera_names, camera_height, camera_width, reward_shaping, render=None, render_offscreen=None, use_image_obs=None, use_depth_obs=None, **kwargs)#

Create environment for processing datasets, which includes extracting observations, labeling dense / sparse rewards, and annotating dones in transitions.

Parameters:

  • env_name (str) – name of environment

  • camera_names (list of str) – list of camera names that correspond to image observations

  • camera_height (int) – camera height for all cameras

  • camera_width (int) – camera width for all cameras

  • reward_shaping (bool) – if True, use shaped environment rewards, else use sparse task completion rewards

  • render (bool or None) – optionally override rendering behavior. Defaults to False.

  • render_offscreen (bool or None) – optionally override rendering behavior. The default value is True if @camera_names is non-empty, False otherwise.

  • use_image_obs (bool or None) – optionally override rendering behavior. The default value is True if @camera_names is non-empty, False otherwise.

  • use_depth_obs (bool) – if True, use depth observations

get_camera_extrinsic_matrix(camera_name)#

Returns a 4x4 homogenous matrix corresponding to the camera pose in the world frame. MuJoCo has a weird convention for how it sets up the camera body axis, so we also apply a correction so that the x and y axis are along the camera view and the z axis points along the viewpoint. Normal camera convention: https://docs.opencv.org/2.4/modules/calib3d/doc/camera_calibration_and_3d_reconstruction.html :param camera_name: name of camera :type camera_name: str

Returns:

4x4 camera extrinsic matrix

Return type:

R (np.array)

get_camera_intrinsic_matrix(camera_name, camera_height, camera_width)#

Obtains camera intrinsic matrix. :param camera_name: name of camera :type camera_name: str :param camera_height: height of camera images in pixels :type camera_height: int :param camera_width: width of camera images in pixels :type camera_width: int

Returns:

3x3 camera matrix

Return type:

K (np.array)

get_camera_transform_matrix(camera_name, camera_height, camera_width)#

Camera transform matrix to project from world coordinates to pixel coordinates. :param camera_name: name of camera :type camera_name: str :param camera_height: height of camera images in pixels :type camera_height: int :param camera_width: width of camera images in pixels :type camera_width: int

Returns:

4x4 camera matrix to project from world coordinates to pixel coordinates

Return type:

K (np.array)

get_goal()#

Get goal observation. Not all environments support this.

get_observation(di=None)#

Get current environment observation dictionary.

Parameters:

di (dict) – current raw observation dictionary from robosuite to wrap and provide as a dictionary. If not provided, will be queried from robosuite.

get_real_depth_map(depth_map)#

Reproduced from https://github.com/ARISE-Initiative/robosuite/blob/c57e282553a4f42378f2635b9a3cbc4afba270fd/robosuite/utils/camera_utils.py#L106 since older versions of robosuite do not have this conversion from normalized depth values returned by MuJoCo to real depth values.

get_reward()#

Get current reward.

get_state()#

Get current environment simulator state as a dictionary. Should be compatible with @reset_to.

is_done()#

Check if the task is done (not necessarily successful).

is_success()#

Check if the task condition(s) is reached. Should return a dictionary { str: bool } with at least a “task” key for the overall task success, and additional optional keys corresponding to other task criteria.

property name#

Returns name of environment name (str).

render(mode='human', height=None, width=None, camera_name='agentview')#

Render from simulation to either an on-screen window or off-screen to RGB array.

Parameters:

  • mode (str) – pass “human” for on-screen rendering or “rgb_array” for off-screen rendering

  • height (int) – height of image to render - only used if mode is “rgb_array”

  • width (int) – width of image to render - only used if mode is “rgb_array”

  • camera_name (str) – camera name to use for rendering

reset()#

Reset environment.

Returns:

initial observation dictionary.

Return type:

observation (dict)

reset_to(state)#

Reset to a specific simulator state.

Parameters:

state (dict) – current simulator state that contains one or more of: - states (np.ndarray): initial state of the mujoco environment - model (str): mujoco scene xml

Returns:

observation dictionary after setting the simulator state (only

if “states” is in @state)

Return type:

observation (dict)

property rollout_exceptions#

Return tuple of exceptions to except when doing rollouts. This is useful to ensure that the entire training run doesn’t crash because of a bad policy that causes unstable simulation computations.

serialize()#

Save all information needed to re-instantiate this environment in a dictionary. This is the same as @env_meta - environment metadata stored in hdf5 datasets, and used in utils/env_utils.py.

set_goal(**kwargs)#

Set goal observation with external specification. Not all environments support this.

step(action)#

Step in the environment with an action.

Parameters:

action (np.array) – action to take

Returns:

new observation dictionary reward (float): reward for this step done (bool): whether the task is done info (dict): extra information

Return type:

observation (dict)

property type#

Returns environment type (int) for this kind of environment. This helps identify this env class.

property version#

Returns version of robosuite used for this environment, eg. 1.2.0

robomimic.envs.wrappers module#

A collection of useful environment wrappers.

class robomimic.envs.wrappers.EnvWrapper(env)#

Bases: object

Base class for all environment wrappers in robomimic.

classmethod class_name()#
property unwrapped#

Grabs unwrapped environment

Returns:

Unwrapped environment

Return type:

env (EnvBase instance)

class robomimic.envs.wrappers.FrameStackWrapper(env, num_frames)#

Bases: robomimic.envs.wrappers.EnvWrapper

Wrapper for frame stacking observations during rollouts. The agent receives a sequence of past observations instead of a single observation when it calls @env.reset, @env.reset_to, or @env.step in the rollout loop.

cache_obs_history()#
reset()#

Modify to return frame stacked observation which is @self.num_frames copies of the initial observation.

Returns:

each observation key in original observation now has

leading shape @self.num_frames and consists of the previous @self.num_frames observations

Return type:

obs_stacked (dict)

reset_to(state)#

Modify to return frame stacked observation which is @self.num_frames copies of the initial observation.

Returns:

each observation key in original observation now has

leading shape @self.num_frames and consists of the previous @self.num_frames observations

Return type:

obs_stacked (dict)

step(action)#

Modify to update the internal frame history and return frame stacked observation, which will have leading dimension @self.num_frames for each key.

Parameters:

action (np.array) – action to take

Returns:

each observation key in original observation now has

leading shape @self.num_frames and consists of the previous @self.num_frames observations

reward (float): reward for this step done (bool): whether the task is done info (dict): extra information

Return type:

obs_stacked (dict)

uncache_obs_history()#
update_obs(obs, action=None, reset=False)#

Module contents#