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, **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#

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#