robomimic.utils package#

Submodules#

robomimic.utils.dataset module#

This file contains Dataset classes that are used by torch dataloaders to fetch batches from hdf5 files.

class robomimic.utils.dataset.SequenceDataset(hdf5_path, obs_keys, dataset_keys, frame_stack=1, seq_length=1, pad_frame_stack=True, pad_seq_length=True, get_pad_mask=False, goal_mode=None, hdf5_cache_mode=None, hdf5_use_swmr=True, hdf5_normalize_obs=False, filter_by_attribute=None, load_next_obs=True)#

Bases: torch.utils.data.dataset.Dataset

close_and_delete_hdf5_handle()#: Maybe close the file handle.

get_dataset_for_ep(ep, key)#: Helper utility to get a dataset for a specific demonstration. Takes into account whether the dataset has been loaded into memory.

get_dataset_sampler()#: Return instance of torch.utils.data.Sampler or None. Allows for dataset to define custom sampling logic, such as re-weighting the probability of samples being drawn. See the train function in scripts/train.py, and torch DataLoader documentation, for more info.

get_dataset_sequence_from_demo(demo_id, index_in_demo, keys, num_frames_to_stack=0, seq_length=1)#

Extract a (sub)sequence of dataset items from a demo given the @keys of the items (e.g., states, actions).

Parameters:

demo_id (str) – id of the demo, e.g., demo_0
index_in_demo (int) – beginning index of the sequence wrt the demo
keys (tuple) – list of keys to extract
num_frames_to_stack (int) – numbers of frame to stack. Seq gets prepended with repeated items if out of range
seq_length (int) – sequence length to extract. Seq gets post-pended with repeated items if out of range

Returns:

a dictionary of extracted items.

get_item(index)#: Main implementation of getitem when not using cache.

get_obs_normalization_stats()#

Returns dictionary of mean and std for each observation key if using observation normalization, otherwise None.

Returns:

a dictionary for observation: normalization. This maps observation keys to dicts with a “mean” and “std” of shape (1, …) where … is the default shape for the observation.

Return type:

obs_normalization_stats (dict)

get_obs_sequence_from_demo(demo_id, index_in_demo, keys, num_frames_to_stack=0, seq_length=1, prefix='obs')#

Extract a (sub)sequence of observation items from a demo given the @keys of the items.

Parameters:

demo_id (str) – id of the demo, e.g., demo_0
index_in_demo (int) – beginning index of the sequence wrt the demo
keys (tuple) – list of keys to extract
num_frames_to_stack (int) – numbers of frame to stack. Seq gets prepended with repeated items if out of range
seq_length (int) – sequence length to extract. Seq gets post-pended with repeated items if out of range
prefix (str) – one of “obs”, “next_obs”

Returns:

a dictionary of extracted items.

get_sequence_from_demo(demo_id, index_in_demo, keys, num_frames_to_stack=0, seq_length=1)#

Extract a (sub)sequence of data items from a demo given the @keys of the items.

Parameters:

demo_id (str) – id of the demo, e.g., demo_0
index_in_demo (int) – beginning index of the sequence wrt the demo
keys (tuple) – list of keys to extract
num_frames_to_stack (int) – numbers of frame to stack. Seq gets prepended with repeated items if out of range
seq_length (int) – sequence length to extract. Seq gets post-pended with repeated items if out of range

Returns:

a dictionary of extracted items.

get_trajectory_at_index(index)#: Method provided as a utility to get an entire trajectory, given the corresponding @index.

property hdf5_file#: This property allows for a lazy hdf5 file open.

hdf5_file_opened()#: Convenient context manager to open the file on entering the scope and then close it on leaving.

load_dataset_in_memory(demo_list, hdf5_file, obs_keys, dataset_keys, load_next_obs)#

Loads the hdf5 dataset into memory, preserving the structure of the file. Note that this differs from self.getitem_cache, which, if active, actually caches the outputs of the getitem operation.

Parameters:

demo_list (list) – list of demo keys, e.g., ‘demo_0’
hdf5_file (h5py.File) – file handle to the hdf5 dataset.
obs_keys (list, tuple) – observation keys to fetch, e.g., ‘images’
dataset_keys (list, tuple) – dataset keys to fetch, e.g., ‘actions’
load_next_obs (bool) – whether to load next_obs from the dataset

Returns:

dictionary of loaded data.

Return type:

all_data (dict)

load_demo_info(filter_by_attribute=None, demos=None)#

Parameters:

filter_by_attribute (str) – if provided, use the provided filter key to select a subset of demonstration trajectories to load
demos (list) – list of demonstration keys to load from the hdf5 file. If omitted, all demos in the file (or under the @filter_by_attribute filter key) are used.

normalize_obs()#: Computes a dataset-wide mean and standard deviation for the observations (per dimension and per obs key) and returns it.

robomimic.utils.env_utils module#

This file contains several utility functions for working with environment wrappers provided by the repository, and with environment metadata saved in dataset files.

robomimic.utils.env_utils.check_env_type(type_to_check, env_meta=None, env_type=None, env=None)#

Checks whether the passed env_meta, env_type, or env is of type @type_to_check. Type corresponds to EB.EnvType.

Parameters:

type_to_check (int) – type to check equality against
env_meta (dict) –
environment metadata, which should be loaded from demonstration hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see @FileUtils.env_from_checkpoint). Contains 3 keys:

’env_name’:

name of environment

’type’:

type of environment, should be a value in EB.EnvType

’env_kwargs’:

dictionary of keyword arguments to pass to environment constructor
env_type (int) – the type of environment, which determines the env class that will be instantiated. Should be a value in EB.EnvType.
env (instance of EB.EnvBase) – environment instance

robomimic.utils.env_utils.check_env_version(env, env_meta)#

Checks whether the passed env and env_meta dictionary having matching environment versions. Logs warning if cannot find version or versions do not match.

Parameters:

env (instance of EB.EnvBase) – environment instance
env_meta (dict) –
environment metadata, which should be loaded from demonstration hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see @FileUtils.env_from_checkpoint). Contains following key:

’env_version’:

environment version, type str

robomimic.utils.env_utils.create_env(env_type, env_name, render=False, render_offscreen=False, use_image_obs=False, use_depth_obs=False, **kwargs)#

Create environment.

Parameters:

env_type (int) – the type of environment, which determines the env class that will be instantiated. Should be a value in EB.EnvType.
env_name (str) – name of environment
render (bool) – if True, environment supports on-screen rendering
render_offscreen (bool) – if True, environment supports off-screen rendering. This is forced to be True if @use_image_obs is True.
use_image_obs (bool) – if True, environment is expected to render rgb image observations on every env.step call. Set this to False for efficiency reasons, if image observations are not required.
use_depth_obs (bool) – if True, environment is expected to render depth image observations on every env.step call. Set this to False for efficiency reasons, if depth observations are not required.

robomimic.utils.env_utils.create_env_for_data_processing(env_meta, camera_names, camera_height, camera_width, reward_shaping, env_class=None, render=None, render_offscreen=None, use_image_obs=None, use_depth_obs=None)#

Creates environment for processing dataset observations and rewards.

Parameters:

env_meta (dict) –
environment metadata, which should be loaded from demonstration hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see @FileUtils.env_from_checkpoint). Contains 3 keys:

’env_name’:

name of environment

’type’:

type of environment, should be a value in EB.EnvType

’env_kwargs’:

dictionary of keyword arguments to pass to environment constructor
camera_names (list of st) – 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
render_offscreen (bool or None) – optionally override rendering behavior
use_image_obs (bool or None) – optionally override rendering behavior
use_depth_obs (bool or None) – optionally override rendering behavior

robomimic.utils.env_utils.create_env_from_metadata(env_meta, env_name=None, render=False, render_offscreen=False, use_image_obs=False, use_depth_obs=False)#

Create environment.

Parameters:

env_meta (dict) –
environment metadata, which should be loaded from demonstration hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see @FileUtils.env_from_checkpoint). Contains 3 keys:

’env_name’:

name of environment

’type’:

type of environment, should be a value in EB.EnvType

’env_kwargs’:

dictionary of keyword arguments to pass to environment constructor
env_name (str) – name of environment. Only needs to be provided if making a different environment from the one in @env_meta.
render (bool) – if True, environment supports on-screen rendering
render_offscreen (bool) – if True, environment supports off-screen rendering. This is forced to be True if @use_image_obs is True.
use_image_obs (bool) – if True, environment is expected to render rgb image observations on every env.step call. Set this to False for efficiency reasons, if image observations are not required.
use_depth_obs (bool) – if True, environment is expected to render depth image observations on every env.step call. Set this to False for efficiency reasons, if depth observations are not required.

robomimic.utils.env_utils.get_env_class(env_meta=None, env_type=None, env=None)#

Return env class from either env_meta, env_type, or env. Note the use of lazy imports - this ensures that modules are only imported when the corresponding env type is requested. This can be useful in practice. For example, a training run that only requires access to gym environments should not need to import robosuite.

Parameters:

env_meta (dict) –
environment metadata, which should be loaded from demonstration hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see @FileUtils.env_from_checkpoint). Contains 3 keys:

’env_name’:

name of environment

’type’:

type of environment, should be a value in EB.EnvType

’env_kwargs’:

dictionary of keyword arguments to pass to environment constructor
env_type (int) – the type of environment, which determines the env class that will be instantiated. Should be a value in EB.EnvType.
env (instance of EB.EnvBase) – environment instance

robomimic.utils.env_utils.get_env_type(env_meta=None, env_type=None, env=None)#

Helper function to get env_type from a variety of inputs.

Parameters:

env_meta (dict) –
environment metadata, which should be loaded from demonstration hdf5 with @FileUtils.get_env_metadata_from_dataset or from checkpoint (see @FileUtils.env_from_checkpoint). Contains 3 keys:

’env_name’:

name of environment

’type’:

type of environment, should be a value in EB.EnvType

’env_kwargs’:

dictionary of keyword arguments to pass to environment constructor
env_type (int) – the type of environment, which determines the env class that will be instantiated. Should be a value in EB.EnvType.
env (instance of EB.EnvBase) – environment instance

robomimic.utils.env_utils.is_robosuite_env(env_meta=None, env_type=None, env=None)#: Determines whether the environment is a robosuite environment. Accepts either env_meta, env_type, or env.

robomimic.utils.env_utils.set_env_specific_obs_processing(env_meta=None, env_type=None, env=None)#: Sets env-specific observation processing. As an example, robosuite depth observations correspond to raw depth and should not be normalized by default, while default depth processing normalizes and clips all values to [0, 1].

robomimic.utils.env_utils.wrap_env_from_config(env, config)#: Wraps environment using the provided Config object to determine which wrappers to use (if any).

robomimic.utils.file_utils module#

A collection of utility functions for working with files, such as reading metadata from demonstration datasets, loading model checkpoints, or downloading dataset files.

class robomimic.utils.file_utils.DownloadProgressBar(*_, **__)#

Bases: tqdm.std.tqdm

update_to(b=1, bsize=1, tsize=None)#

robomimic.utils.file_utils.algo_name_from_checkpoint(ckpt_path=None, ckpt_dict=None)#

Return algorithm name that was used to train a checkpoint or loaded model dictionary.

Parameters:

ckpt_path (str) – Path to checkpoint file. Only needed if not providing @ckpt_dict.
ckpt_dict (dict) – Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path.

Returns:

algorithm name

ckpt_dict (dict): loaded checkpoint dictionary (convenient to avoid: re-loading checkpoint from disk multiple times)

Return type:

algo_name (str)

robomimic.utils.file_utils.config_from_checkpoint(algo_name=None, ckpt_path=None, ckpt_dict=None, verbose=False)#

Helper function to restore config from a checkpoint file or loaded model dictionary.

Parameters:

algo_name (str) – Algorithm name.
ckpt_path (str) – Path to checkpoint file. Only needed if not providing @ckpt_dict.
ckpt_dict (dict) – Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path.
verbose (bool) – if True, include print statements

Returns:

Raw loaded configuration, without properties replaced.

ckpt_dict (dict): loaded checkpoint dictionary (convenient to avoid: re-loading checkpoint from disk multiple times)

Return type:

config (dict)

robomimic.utils.file_utils.create_hdf5_filter_key(hdf5_path, demo_keys, key_name)#

Creates a new hdf5 filter key in hdf5 file @hdf5_path with name @key_name that corresponds to the demonstrations @demo_keys. Filter keys are generally useful to create named subsets of the demonstrations in an hdf5, making it easy to train, test, or report statistics on a subset of the trajectories in a file.

Returns the list of episode lengths that correspond to the filtering.

Parameters:

hdf5_path (str) – path to hdf5 file
demo_keys ([str]) – list of demonstration keys which should correspond to this filter key. For example, [“demo_0”, “demo_1”].
key_name (str) – name of filter key to create

Returns:

list of episode lengths that corresponds to: each demonstration in the new filter key

Return type:

ep_lengths ([int])

robomimic.utils.file_utils.download_url(url, download_dir, check_overwrite=True)#

First checks that @url is reachable, then downloads the file at that url into the directory specified by @download_dir. Prints a progress bar during the download using tqdm.

Modified from https://github.com/tqdm/tqdm#hooks-and-callbacks, and https://stackoverflow.com/a/53877507.

Parameters:

url (str) – url string
download_dir (str) – path to directory where file should be downloaded
check_overwrite (bool) – if True, will sanity check the download fpath to make sure a file of that name doesn’t already exist there

robomimic.utils.file_utils.env_from_checkpoint(ckpt_path=None, ckpt_dict=None, env_name=None, render=False, render_offscreen=False, verbose=False)#

Creates an environment using the metadata saved in a checkpoint.

Parameters:

ckpt_path (str) – Path to checkpoint file. Only needed if not providing @ckpt_dict.
ckpt_dict (dict) – Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path.
env_name (str) – if provided, override environment name saved in checkpoint
render (bool) – if True, environment supports on-screen rendering
render_offscreen (bool) – if True, environment supports off-screen rendering. This is forced to be True if saved model uses image observations.

Returns:

environment created using checkpoint

ckpt_dict (dict): loaded checkpoint dictionary (convenient to avoid: re-loading checkpoint from disk multiple times)

Return type:

env (EnvBase instance)

robomimic.utils.file_utils.get_demos_for_filter_key(hdf5_path, filter_key)#

Gets demo keys that correspond to a particular filter key.

Parameters:

hdf5_path (str) – path to hdf5 file
filter_key (str) – name of filter key

Returns:

list of demonstration keys that: correspond to this filter key. For example, [“demo_0”, “demo_1”].

Return type:

demo_keys ([str])

robomimic.utils.file_utils.get_env_metadata_from_dataset(dataset_path, set_env_specific_obs_processors=True)#

Retrieves env metadata from dataset.

Parameters:

dataset_path (str) – path to dataset
set_env_specific_obs_processors (bool) – environment might have custom rules for how to process observations - if this flag is true, make sure ObsUtils will use these custom settings. This is a good place to do this operation to make sure it happens before loading data, running a trained model, etc.

Returns:

environment metadata. Contains 3 keys:

’env_name’:

name of environment

’type’:

type of environment, should be a value in EB.EnvType

’env_kwargs’:

dictionary of keyword arguments to pass to environment constructor

Return type:

env_meta (dict)

robomimic.utils.file_utils.get_shape_metadata_from_dataset(dataset_path, all_obs_keys=None, verbose=False)#

Retrieves shape metadata from dataset.

Parameters:

dataset_path (str) – path to dataset
all_obs_keys (list) – list of all modalities used by the model. If not provided, all modalities present in the file are used.
verbose (bool) – if True, include print statements

Returns:

shape metadata. Contains the following keys:

’ac_dim’:

action space dimension

’all_shapes’:

dictionary that maps observation key string to shape

’all_obs_keys’:

list of all observation modalities used

’use_images’:

bool, whether or not image modalities are present

’use_depths’:

bool, whether or not depth modalities are present

Return type:

shape_meta (dict)

robomimic.utils.file_utils.load_dict_from_checkpoint(ckpt_path)#

Load checkpoint dictionary from a checkpoint file.

Parameters:: ckpt_path (str) – Path to checkpoint file.
Returns:: Loaded checkpoint dictionary.
Return type:: ckpt_dict (dict)

robomimic.utils.file_utils.maybe_dict_from_checkpoint(ckpt_path=None, ckpt_dict=None)#

Utility function for the common use case where either an ckpt path or a ckpt_dict is provided. This is a no-op if ckpt_dict is not None, otherwise it loads the model dict from the ckpt path.

Parameters:

ckpt_path (str) – Path to checkpoint file. Only needed if not providing @ckpt_dict.
ckpt_dict (dict) – Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path.

Returns:

Loaded checkpoint dictionary.

Return type:

ckpt_dict (dict)

robomimic.utils.file_utils.policy_from_checkpoint(device=None, ckpt_path=None, ckpt_dict=None, verbose=False)#

This function restores a trained policy from a checkpoint file or loaded model dictionary.

Parameters:

device (torch.device) – if provided, put model on this device
ckpt_path (str) – Path to checkpoint file. Only needed if not providing @ckpt_dict.
ckpt_dict (dict) – Loaded model checkpoint dictionary. Only needed if not providing @ckpt_path.
verbose (bool) – if True, include print statements

Returns:

instance of Algo that has the saved weights from: the checkpoint file, and also acts as a policy that can easily interact with an environment in a training loop
ckpt_dict (dict): loaded checkpoint dictionary (convenient to avoid: re-loading checkpoint from disk multiple times)

Return type:

model (RolloutPolicy)

robomimic.utils.file_utils.update_config(cfg)#

Updates the config for backwards-compatibility if it uses outdated configurations.

See https://github.com/ARISE-Initiative/robomimic/releases/tag/v0.2.0 for more info.

Parameters:: cfg (dict) – Raw dictionary of config values

robomimic.utils.file_utils.url_is_alive(url)#

Checks that a given URL is reachable. From https://gist.github.com/dehowell/884204.

Parameters:: url (str) – url string
Returns:: True if url is reachable, False otherwise
Return type:: is_alive (bool)

robomimic.utils.hyperparam_utils module#

A collection of utility functions and classes for generating config jsons for hyperparameter sweeps.

class robomimic.utils.hyperparam_utils.ConfigGenerator(base_config_file, wandb_proj_name='debug', script_file=None, generated_config_dir=None)#

Bases: object

Useful class to keep track of hyperparameters to sweep, and to generate the json configs for each experiment run.

add_param(key, name, group, values, value_names=None)#

Add parameter to the hyperparameter sweep.

Parameters:

key (str) – location of parameter in the config, using hierarchical key format (ex. train/data = config.train.data)
name (str) – name, as it will appear in the experiment name
group (int) – group id - parameters with the same ID have their values swept together
values (list) – list of values to sweep over for this parameter
value_names ([str]) – if provided, strings to use in experiment name for each value, instead of the parameter value. This is helpful for parameters that may have long or large values (for example, dataset path).

generate()#: Generates json configs for the hyperparameter sweep using attributes @self.parameters, @self.base_config_file, and @self.script_file, all of which should have first been set externally by calling @add_param, @set_base_config_file, and @set_script_file.

robomimic.utils.hyperparam_utils.get_value_for_key(dic, k)#

Get value for nested dictionary with levels denoted by “/” or “.”. For example, if @k is “a/b”, then this function returns @dic[“a”][“b”].

Parameters:

dic (dict) – a nested dictionary
k (str) – a single string meant to index several levels down into the nested dictionary, where levels can be denoted by “/” or by “.”.

Returns:

the nested dictionary value for the provided key

Return type:

val

robomimic.utils.hyperparam_utils.load_json(json_file, verbose=True)#

Simple utility function to load a json file as a dict.

Parameters:

json_file (str) – path to json file to load
verbose (bool) – if True, pretty print the loaded json dictionary

Returns:

json dictionary

Return type:

config (dict)

robomimic.utils.hyperparam_utils.save_json(config, json_file)#

Simple utility function to save a dictionary to a json file on disk.

Parameters:

config (dict) – dictionary to save
json_file (str) – path to json file to write

robomimic.utils.hyperparam_utils.set_value_for_key(dic, k, v)#

Set value for hierarchical dictionary with levels denoted by “/” or “.”.

Parameters:

dic (dict) – a nested dictionary
k (str) – a single string meant to index several levels down into the nested dictionary, where levels can be denoted by “/” or by “.”.
v – the value to set at the provided key

robomimic.utils.log_utils module#

This file contains utility classes and functions for logging to stdout, stderr, and to tensorboard.

class robomimic.utils.log_utils.DataLogger(log_dir, config, log_tb=True, log_wandb=False)#

Bases: object

Logging class to log metrics to tensorboard and/or retrieve running statistics about logged data.

close()#: Run before terminating to make sure all logs are flushed

get_stats(k)#

Computes running statistics for a particular key. :param k: key string :type k: str

Returns:: dictionary of statistics
Return type:: stats (dict)

record(k, v, epoch, data_type='scalar', log_stats=False)#: Record data with logger. :param k: key string :type k: str :param v: value to store :type v: float or image :param epoch: current epoch number :param data_type: the type of data. either ‘scalar’ or ‘image’ :type data_type: str :param log_stats: whether to store the mean/max/min/std for all data logged so far with key k :type log_stats: bool

class robomimic.utils.log_utils.PrintLogger(log_file)#

Bases: object

This class redirects print statements to both console and a file.

flush()#

write(message)#

class robomimic.utils.log_utils.custom_tqdm(*_, **__)#

Bases: tqdm.std.tqdm

Small extension to tqdm to make a few changes from default behavior. By default tqdm writes to stderr. Instead, we change it to write to stdout.

robomimic.utils.log_utils.flush_warnings()#: This function flushes all warnings from the global warning buffer to the terminal and clears the global registry.

robomimic.utils.log_utils.log_warning(message, color='yellow', print_now=True)#

This function logs a warning message by recording it in a global warning buffer. The global registry will be maintained until @flush_warnings is called, at which point the warnings will get printed to the terminal.

Parameters:

message (str) – warning message to display
color (str) – color of message - defaults to “yellow”
print_now (bool) – if True (default), will print to terminal immediately, in addition to adding it to the global warning buffer

robomimic.utils.log_utils.silence_stdout()#

This contextmanager will redirect stdout so that nothing is printed to the terminal. Taken from the link below:

https://stackoverflow.com/questions/6735917/redirecting-stdout-to-nothing-in-python

robomimic.utils.loss_utils module#

This file contains a collection of useful loss functions for use with torch tensors.

robomimic.utils.loss_utils.KLD_0_1_loss(mu, logvar)#

KL divergence loss. Computes D_KL( N(mu, sigma) || N(0, 1) ). Note that this function averages across the batch dimension, but sums across dimension.

Parameters:

mu (torch.Tensor) – mean tensor of shape (B, D)
logvar (torch.Tensor) – logvar tensor of shape (B, D)

Returns:

KL divergence loss between the input gaussian distribution: and N(0, 1)

Return type:

loss (torch.Tensor)

robomimic.utils.loss_utils.KLD_gaussian_loss(mu_1, logvar_1, mu_2, logvar_2)#

KL divergence loss between two Gaussian distributions. This function computes the average loss across the batch.

Parameters:

mu_1 (torch.Tensor) – first means tensor of shape (B, D)
logvar_1 (torch.Tensor) – first logvars tensor of shape (B, D)
mu_2 (torch.Tensor) – second means tensor of shape (B, D)
logvar_2 (torch.Tensor) – second logvars tensor of shape (B, D)

Returns:

KL divergence loss between the two gaussian distributions

Return type:

loss (torch.Tensor)

robomimic.utils.loss_utils.cosine_loss(preds, labels)#

Cosine loss between two tensors.

Parameters:

preds (torch.Tensor) – torch tensor
labels (torch.Tensor) – torch tensor

Returns:

cosine loss

Return type:

loss (torch.Tensor)

robomimic.utils.loss_utils.log_mean_exp(x, dim)#

Compute the log(mean(exp(x), dim)) in a numerically stable manner. Adapted from CS 236 at Stanford.

Parameters:

x (torch.Tensor) – a tensor
dim (int) – dimension along which mean is computed

Returns:

log(mean(exp(x), dim))

Return type:

y (torch.Tensor)

robomimic.utils.loss_utils.log_normal(x, m, v)#

Log probability of tensor x under diagonal multivariate normal with mean m and variance v. The last dimension of the tensors is treated as the dimension of the Gaussian distribution - all other dimensions are treated as independent Gaussians. Adapted from CS 236 at Stanford.

Parameters:

x (torch.Tensor) – tensor with shape (B, …, D)
m (torch.Tensor) – means tensor with shape (B, …, D) or (1, …, D)
v (torch.Tensor) – variances tensor with shape (B, …, D) or (1, …, D)

Returns:

log probabilities of shape (B, …)

Return type:

log_prob (torch.Tensor)

robomimic.utils.loss_utils.log_normal_mixture(x, m, v, w=None, log_w=None)#

Log probability of tensor x under a uniform mixture of Gaussians. Adapted from CS 236 at Stanford.

Parameters:

x (torch.Tensor) – tensor with shape (B, D)
m (torch.Tensor) – means tensor with shape (B, M, D) or (1, M, D), where M is number of mixture components
v (torch.Tensor) – variances tensor with shape (B, M, D) or (1, M, D) where M is number of mixture components
w (torch.Tensor) – weights tensor - if provided, should be shape (B, M) or (1, M)
log_w (torch.Tensor) – log-weights tensor - if provided, should be shape (B, M) or (1, M)

Returns:

log probabilities of shape (B,)

Return type:

log_prob (torch.Tensor)

robomimic.utils.loss_utils.log_sum_exp(x, dim=0)#

Compute the log(sum(exp(x), dim)) in a numerically stable manner. Adapted from CS 236 at Stanford.

Parameters:

x (torch.Tensor) – a tensor
dim (int) – dimension along which sum is computed

Returns:

log(sum(exp(x), dim))

Return type:

y (torch.Tensor)

robomimic.utils.loss_utils.project_values_onto_atoms(values, probabilities, atoms)#

Project the categorical distribution given by @probabilities on the grid of values given by @values onto a grid of values given by @atoms. This is useful when computing a bellman backup where the backed up values from the original grid will not be in the original support, requiring L2 projection.

Each value in @values has a corresponding probability in @probabilities - this probability mass is shifted to the closest neighboring grid points in @atoms in proportion. For example, if the value in question is 0.2, and the neighboring atoms are 0 and 1, then 0.8 of the probability weight goes to atom 0 and 0.2 of the probability weight will go to 1.

Adapted from https://github.com/deepmind/acme/blob/master/acme/tf/losses/distributional.py#L42

Parameters:

values – value grid to project, of shape (batch_size, n_atoms)
probabilities – probabilities for categorical distribution on @values, shape (batch_size, n_atoms)
atoms – value grid to project onto, of shape (n_atoms,) or (1, n_atoms)

Returns:

new probability vectors that correspond to the L2 projection of the categorical distribution onto @atoms

robomimic.utils.obs_utils module#

A collection of utilities for working with observation dictionaries and different kinds of modalities such as images.

class robomimic.utils.obs_utils.DepthModality#

Bases: robomimic.utils.obs_utils.Modality

Modality for depth observations

name = 'depth'#

class robomimic.utils.obs_utils.ImageModality#

Bases: robomimic.utils.obs_utils.Modality

Modality for RGB image observations

name = 'rgb'#

class robomimic.utils.obs_utils.LowDimModality#

Bases: robomimic.utils.obs_utils.Modality

Modality for low dimensional observations

name = 'low_dim'#

class robomimic.utils.obs_utils.Modality#

Bases: object

Observation Modality class to encapsulate necessary functions needed to process observations of this modality

classmethod add_keys(keys)#

Adds the observation @keys associated with this modality to the current set of keys.

Parameters:: keys (list or set) – observation keys to add to associate with this modality

keys = {}#

name = None#

classmethod process_obs(obs)#

Prepares an observation @obs of this modality for network input.

Parameters:: obs (np.array or torch.Tensor) – raw observation, which may include a leading batch dimension
Returns:: processed observation
Return type:: np.array or torch.Tensor

classmethod process_obs_from_dict(obs_dict, inplace=True)#

Receives a dictionary of keyword mapped observations @obs_dict, and processes the observations with keys corresponding to this modality. A copy will be made of the received dictionary unless @inplace is True

Parameters:

obs_dict (dict) – Dictionary mapping observation keys to observations
inplace (bool) – If True, will modify @obs_dict in place, otherwise, will create a copy

Returns:

observation dictionary with processed observations corresponding to this modality

Return type:

dict

classmethod set_keys(keys)#

Sets the observation keys associated with this modality.

Parameters:: keys (list or set) – observation keys to associate with this modality

classmethod set_obs_processor(processor=None)#

Sets the processor for this observation modality. If @processor is set to None, then the obs processor will use the default one (self.process_obs(…)). Otherwise, @processor should be a function to process this corresponding observation modality.

Parameters:: processor (function or None) – If not None, should be function that takes in either a np.array or torch.Tensor and output the processed array / tensor. If None, will reset to the default processor (self.process_obs(…))

classmethod set_obs_unprocessor(unprocessor=None)#

Sets the unprocessor for this observation modality. If @unprocessor is set to None, then the obs unprocessor will use the default one (self.unprocess_obs(…)). Otherwise, @unprocessor should be a function to process this corresponding observation modality.

Parameters:: unprocessor (function or None) – If not None, should be function that takes in either a np.array or torch.Tensor and output the unprocessed array / tensor. If None, will reset to the default unprocessor (self.unprocess_obs(…))

classmethod unprocess_obs(obs)#

Prepares an observation @obs of this modality for deployment.

Parameters:: obs (np.array or torch.Tensor) – processed observation, which may include a leading batch dimension
Returns:: unprocessed observation
Return type:: np.array or torch.Tensor

class robomimic.utils.obs_utils.ObservationKeyToModalityDict#

Bases: dict

Custom dictionary class with the sole additional purpose of automatically registering new “keys” at runtime without breaking. This is mainly for backwards compatibility, where certain keys such as “latent”, “actions”, etc. are used automatically by certain models (e.g.: VAEs) but were never specified by the user externally in their config. Thus, this dictionary will automatically handle those keys by implicitly associating them with the low_dim modality.

class robomimic.utils.obs_utils.ScanModality#

Bases: robomimic.utils.obs_utils.Modality

Modality for scan observations

name = 'scan'#

robomimic.utils.obs_utils.batch_image_chw_to_hwc(im)#

Inverse of channel swap in @batch_image_hwc_to_chw.

Parameters:

im (np.array or torch.Tensor) – image of shape (batch, channel, height, width) or (channel, height, width)

Returns:

image of shape (batch, height, width, channel): or (height, width, channel)

Return type:

im (np.array or torch.Tensor)

robomimic.utils.obs_utils.batch_image_hwc_to_chw(im)#

Channel swap for images - useful for preparing images for torch training.

Parameters:

im (np.array or torch.Tensor) – image of shape (batch, height, width, channel) or (height, width, channel)

Returns:

image of shape (batch, channel, height, width): or (channel, height, width)

Return type:

im (np.array or torch.Tensor)

robomimic.utils.obs_utils.center_crop(im, t_h, t_w)#

Takes a center crop of an image.

Parameters:

im (np.array or torch.Tensor) – image of shape (…, height, width, channel)
t_h (int) – height of crop
t_w (int) – width of crop

Returns:

center cropped image

Return type:

im (np.array or torch.Tensor)

robomimic.utils.obs_utils.crop_image_from_indices(images, crop_indices, crop_height, crop_width)#

Crops images at the locations specified by @crop_indices. Crops will be taken across all channels.

Parameters:

images (torch.Tensor) – batch of images of shape […, C, H, W]
crop_indices (torch.Tensor) – batch of indices of shape […, N, 2] where N is the number of crops to take per image and each entry corresponds to the pixel height and width of where to take the crop. Note that the indices can also be of shape […, 2] if only 1 crop should be taken per image. Leading dimensions must be consistent with @images argument. Each index specifies the top left of the crop. Values must be in range [0, H - CH - 1] x [0, W - CW - 1] where H and W are the height and width of @images and CH and CW are @crop_height and @crop_width.
crop_height (int) – height of crop to take
crop_width (int) – width of crop to take

Returns:

cropped images of shape […, C, @crop_height, @crop_width]

Return type:

crops (torch.Tesnor)

robomimic.utils.obs_utils.get_processed_shape(obs_modality, input_shape)#

Given observation modality @obs_modality and expected inputs of shape @input_shape (excluding batch dimension), return the expected processed observation shape resulting from process_{obs_modality}.

Parameters:

obs_modality (str) – Observation modality to use (e.g.: low_dim, rgb, depth, etc…)
input_shape (list of int) – Expected input dimensions, excluding the batch dimension

Returns:

expected processed input shape

Return type:

list of int

robomimic.utils.obs_utils.has_modality(modality, obs_keys)#

Returns True if @modality is present in the list of observation keys @obs_keys.

Parameters:

modality (str) – modality to check for, e.g.: rgb, depth, etc.
obs_keys (list) – list of observation keys

robomimic.utils.obs_utils.initialize_default_obs_encoder(obs_encoder_config)#

Initializes the default observation encoder kwarg information to be used by all networks if no values are manually specified at runtime.

Parameters:: obs_encoder_config (Config) – Observation encoder config to use. Should be equivalent to config.observation.encoder

robomimic.utils.obs_utils.initialize_obs_modality_mapping_from_dict(modality_mapping)#

This function is an alternative to @initialize_obs_utils_with_obs_specs, that allows manually setting of modalities. NOTE: Only one of these should be called at runtime – not both! (Note that all training scripts that use a config)

automatically handle obs modality mapping, so using this function is usually unnecessary)

Parameters:: modality_mapping (dict) – Maps modality string names (e.g.: rgb, low_dim, etc.) to a list of observation keys that should belong to that modality

robomimic.utils.obs_utils.initialize_obs_utils_with_config(config)#

Utility function to parse config and call @initialize_obs_utils_with_obs_specs and @initialize_default_obs_encoder_kwargs with the correct arguments.

Parameters:: config (BaseConfig instance) – config object

robomimic.utils.obs_utils.initialize_obs_utils_with_obs_specs(obs_modality_specs)#

This function should be called before using any observation key-specific functions in this file, in order to make sure that all utility functions are aware of the observation modalities (e.g. which ones are low-dimensional, which ones are rgb, etc.).

It constructs two dictionaries: (1) that map observation modality (e.g. low_dim, rgb) to a list of observation keys under that modality, and (2) that maps the inverse, specific observation keys to their corresponding observation modality.

Input should be a nested dictionary (or list of such dicts) with the following structure:

obs_variant (str):
obs_modality (str): observation keys (list) …

…

Example

{

“obs”: {: “low_dim”: [“robot0_eef_pos”, “robot0_eef_quat”], “rgb”: [“agentview_image”, “robot0_eye_in_hand”],

} “goal”: {

“low_dim”: [“robot0_eef_pos”], “rgb”: [“agentview_image”]

}

In the example, raw observations consist of low-dim and rgb modalities, with the robot end effector pose under low-dim, and the agentview and wrist camera images under rgb, while goal observations also consist of low-dim and rgb modalities, with a subset of the raw observation keys per modality.

Parameters:: obs_modality_specs (dict or list) – A nested dictionary (see docstring above for an example) or a list of nested dictionaries. Accepting a list as input makes it convenient for situations where multiple modules may each have their own modality spec.

robomimic.utils.obs_utils.key_is_obs_modality(key, obs_modality)#

Check if observation key corresponds to modality @obs_modality.

Parameters:

key (str) – obs key name to check
obs_modality (str) – observation modality - e.g.: “low_dim”, “rgb”

robomimic.utils.obs_utils.normalize_obs(obs_dict, obs_normalization_stats)#

Normalize observations using the provided “mean” and “std” entries for each observation key. The observation dictionary will be modified in-place.

Parameters:

obs_dict (dict) – dictionary mapping observation key to np.array or torch.Tensor. Can have any number of leading batch dimensions.
obs_normalization_stats (dict) – this should map observation keys to dicts with a “mean” and “std” of shape (1, …) where … is the default shape for the observation.

Returns:

obs dict with normalized observation arrays

Return type:

obs_dict (dict)

robomimic.utils.obs_utils.obs_encoder_kwargs_from_config(obs_encoder_config)#

Generate a set of args used to create visual backbones for networks from the observation encoder config.

Parameters:: obs_encoder_config (Config) – Config object containing relevant encoder information. Should be equivalent to config.observation.encoder
Returns:: Processed encoder kwargs
Return type:: dict

robomimic.utils.obs_utils.process_frame(frame, channel_dim, scale)#

Given frame fetched from dataset, process for network input. Converts array to float (from uint8), normalizes pixels from range [0, @scale] to [0, 1], and channel swaps from (H, W, C) to (C, H, W).

Parameters:

frame (np.array or torch.Tensor) – frame array
channel_dim (int) – Number of channels to sanity check for
scale (float or None) – Value to normalize inputs by

Returns:

processed frame

Return type:

processed_frame (np.array or torch.Tensor)

robomimic.utils.obs_utils.process_obs(obs, obs_modality=None, obs_key=None)#

Process observation @obs corresponding to @obs_modality modality (or implicitly inferred from @obs_key) to prepare for network input.

Note that either obs_modality OR obs_key must be specified!

If both are specified, obs_key will override obs_modality

Parameters:

obs (np.array or torch.Tensor) – Observation to process. Leading batch dimension is optional
obs_modality (str) – Observation modality (e.g.: depth, image, low_dim, etc.)
obs_key (str) – Name of observation from which to infer @obs_modality

Returns:

processed observation

Return type:

processed_obs (np.array or torch.Tensor)

robomimic.utils.obs_utils.process_obs_dict(obs_dict)#

Process observations in observation dictionary to prepare for network input.

Parameters:: obs_dict (dict) – dictionary mapping observation keys to np.array or torch.Tensor. Leading batch dimensions are optional.
Returns:: dictionary where observation keys have been processed by their corresponding processors
Return type:: new_dict (dict)

robomimic.utils.obs_utils.register_encoder_core(target_class)#

robomimic.utils.obs_utils.register_obs_key(target_class)#

robomimic.utils.obs_utils.register_randomizer(target_class)#

robomimic.utils.obs_utils.repeat_and_stack_observation(obs_dict, n)#

Given an observation dictionary and a desired repeat value @n, this function will return a new observation dictionary where each modality is repeated @n times and the copies are stacked in the first dimension.

For example, if a batch of 3 observations comes in, and n is 2, the output will look like [ob1; ob1; ob2; ob2; ob3; ob3] in each modality.

Parameters:

obs_dict (dict) – dictionary mapping observation key to np.array or torch.Tensor. Leading batch dimensions are optional.
n (int) – number to repeat by

Returns:

repeated obs dict

Return type:

repeat_obs_dict (dict)

robomimic.utils.obs_utils.sample_random_image_crops(images, crop_height, crop_width, num_crops, pos_enc=False)#

For each image, randomly sample @num_crops crops of size (@crop_height, @crop_width), from @images.

Parameters:

images (torch.Tensor) – batch of images of shape […, C, H, W]
crop_height (int) – height of crop to take
crop_width (int) – width of crop to take
num_crops (n) – number of crops to sample
pos_enc (bool) – if True, also add 2 channels to the outputs that gives a spatial encoding of the original source pixel locations. This means that the output crops will contain information about where in the source image it was sampled from.

Returns:

crops of shape (…, @num_crops, C, @crop_height, @crop_width): if @pos_enc is False, otherwise (…, @num_crops, C + 2, @crop_height, @crop_width)

crop_inds (torch.Tensor): sampled crop indices of shape (…, N, 2)

Return type:

crops (torch.Tensor)

robomimic.utils.obs_utils.unprocess_frame(frame, channel_dim, scale)#

Given frame prepared for network input, prepare for saving to dataset. Inverse of @process_frame.

Parameters:

frame (np.array or torch.Tensor) – frame array
channel_dim (int) – What channel dimension should be (used for sanity check)
scale (float or None) – Scaling factor to apply during denormalization

Returns:

frame passed through: inverse operation of @process_frame

Return type:

unprocessed_frame (np.array or torch.Tensor)

robomimic.utils.obs_utils.unprocess_obs(obs, obs_modality=None, obs_key=None)#

Prepare observation @obs corresponding to @obs_modality modality (or implicitly inferred from @obs_key) to prepare for deployment.

Note that either obs_modality OR obs_key must be specified!

If both are specified, obs_key will override obs_modality

Parameters:

obs (np.array or torch.Tensor) – Observation to unprocess. Leading batch dimension is optional
obs_modality (str) – Observation modality (e.g.: depth, image, low_dim, etc.)
obs_key (str) – Name of observation from which to infer @obs_modality

Returns:

unprocessed observation

Return type:

unprocessed_obs (np.array or torch.Tensor)

robomimic.utils.obs_utils.unprocess_obs_dict(obs_dict)#

Prepare processed observation dictionary for saving to dataset. Inverse of @process_obs.

Parameters:

obs_dict (dict) – dictionary mapping observation keys to np.array or torch.Tensor. Leading batch dimensions are optional.

Returns:

dictionary where observation keys have been unprocessed by: their respective unprocessor methods

Return type:

new_dict (dict)

robomimic.utils.python_utils module#

Set of general purpose utility functions for easier interfacing with Python API

robomimic.utils.python_utils.extract_class_init_kwargs_from_dict(cls, dic, copy=False, verbose=False)#

Helper function to return a dictionary of key-values that specifically correspond to @cls class’s __init__ constructor method, from @dic which may or may not contain additional, irrelevant kwargs.

Note that @dic may possibly be missing certain kwargs as specified by cls.__init__. No error will be raised.

Parameters:

cls (object) – Class from which to grab __init__ kwargs that will be be used as filtering keys for @dic
dic (dict) – Dictionary containing multiple key-values
copy (bool) – If True, will deepcopy all values corresponding to the specified @keys
verbose (bool) – If True (or if macro DEBUG is True), then will print out mismatched keys

Returns:

Extracted subset dictionary possibly containing only the specified keys from cls.__init__ and their: corresponding values

Return type:

dict

robomimic.utils.python_utils.extract_subset_dict(dic, keys, copy=False)#

Helper function to extract a subset of dictionary key-values from a current dictionary. Optionally (deep)copies the values extracted from the original @dic if @copy is True.

Parameters:

dic (dict) – Dictionary containing multiple key-values
keys (Iterable) – Specific keys to extract from @dic. If the key doesn’t exist in @dic, then the key is skipped
copy (bool) – If True, will deepcopy all values corresponding to the specified @keys

Returns:

Extracted subset dictionary containing only the specified @keys and their corresponding values

Return type:

dict

robomimic.utils.python_utils.get_class_init_kwargs(cls)#

Helper function to return a list of all valid keyword arguments (excluding “self”) for the given @cls class.

Parameters:: cls (object) – Class from which to grab __init__ kwargs
Returns:: All keyword arguments (excluding “self”) specified by @cls __init__ constructor method
Return type:: list

robomimic.utils.tensor_utils module#

A collection of utilities for working with nested tensor structures consisting of numpy arrays and torch tensors.

robomimic.utils.tensor_utils.assert_size_at_dim(x, size, dim, msg)#

Ensure that arrays and tensors in nested dictionary or list or tuple have size @size in dim @dim.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
size (int) – size that tensors should have at @dim
dim (int) – dimension to check

robomimic.utils.tensor_utils.assert_size_at_dim_single(x, size, dim, msg)#

Ensure that array or tensor @x has size @size in dim @dim.

Parameters:

x (np.ndarray or torch.Tensor) – input array or tensor
size (int) – size that tensors should have at @dim
dim (int) – dimension to check
msg (str) – text to display if assertion fails

robomimic.utils.tensor_utils.clone(x)#

Clones all torch tensors and numpy arrays in nested dictionary or list or tuple and returns a new nested structure.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.contiguous(x)#

Makes all torch tensors and numpy arrays contiguous in nested dictionary or list or tuple and returns a new nested structure.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.detach(x)#

Detaches all torch tensors in nested dictionary or list or tuple and returns a new nested structure.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.expand_at(x, size, dim)#

Expand all tensors in nested dictionary or list or tuple at a single dimension @dim by @size.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
size (int) – size to expand
dim (int) – dimension to expand

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.expand_at_single(x, size, dim)#

Expand a tensor at a single dimension @dim by @size

Parameters:

x (torch.Tensor) – input tensor
size (int) – size to expand
dim (int) – dimension to expand

Returns:

expanded tensor

Return type:

y (torch.Tensor)

robomimic.utils.tensor_utils.flatten(x, begin_axis=1)#

Flatten all tensors in nested dictionary or list or tuple, from @begin_axis onwards.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
begin_axis (int) – which axis to flatten from

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.flatten_nested_dict_list(d, parent_key='', sep='_', item_key='')#

Flatten a nested dict or list to a list.

For example, given a dict {

a: 1 b: {

c: 2

} c: 3

}

the function would return [(a, 1), (b_c, 2), (c, 3)]

Parameters:

d (dict, list) – a nested dict or list to be flattened
parent_key (str) – recursion helper
sep (str) – separator for nesting keys
item_key (str) – recursion helper

Returns:

a list of (key, value) tuples

Return type:

list

robomimic.utils.tensor_utils.flatten_single(x, begin_axis=1)#

Flatten a tensor in all dimensions from @begin_axis onwards.

Parameters:

x (torch.Tensor) – tensor to flatten
begin_axis (int) – which axis to flatten from

Returns:

flattened tensor

Return type:

y (torch.Tensor)

robomimic.utils.tensor_utils.gather_along_dim_with_dim(x, target_dim, source_dim, indices)#

Apply @gather_along_dim_with_dim_single to all tensors in a nested dictionary or list or tuple.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
target_dim (int) – dimension to gather values along
source_dim (int) – dimension to hold constant and use for gathering values from the other dimensions
indices (torch.Tensor) – flat index tensor with same shape as tensor @x along @source_dim

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.gather_along_dim_with_dim_single(x, target_dim, source_dim, indices)#

This function indexes out a target dimension of a tensor in a structured way, by allowing a different value to be selected for each member of a flat index tensor (@indices) corresponding to a source dimension. This can be interpreted as moving along the source dimension, using the corresponding index value in @indices to select values for all other dimensions outside of the source and target dimensions. A common use case is to gather values in target dimension 1 for each batch member (target dimension 0).

Parameters:

x (torch.Tensor) – tensor to gather values for
target_dim (int) – dimension to gather values along
source_dim (int) – dimension to hold constant and use for gathering values from the other dimensions
indices (torch.Tensor) – flat index tensor with same shape as tensor @x along @source_dim

Returns:

gathered tensor, with dimension @target_dim indexed out

Return type:

y (torch.Tensor)

robomimic.utils.tensor_utils.gather_sequence(seq, indices)#

Given a nested dictionary or list or tuple, gathers an element from each sequence of the batch for tensors with leading dimensions [B, T, …].

Parameters:

seq (dict or list or tuple) – a possibly nested dictionary or list or tuple with tensors of leading dimensions [B, T, …]
indices (torch.Tensor) – tensor indices of shape [B]

Returns:

new nested dict-list-tuple with tensors of shape [B, …]

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.gather_sequence_single(seq, indices)#

Given a tensor with leading dimensions [B, T, …], gather an element from each sequence in the batch given an index for each sequence.

Parameters:

seq (torch.Tensor) – tensor with leading dimensions [B, T, …]
indices (torch.Tensor) – tensor indices of shape [B]

Returns:

indexed tensor of shape [B, ….]

Return type:

y (torch.Tensor)

robomimic.utils.tensor_utils.get_shape(x)#

Get all shapes of arrays and tensors in nested dictionary or list or tuple.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple

Returns:

new nested dict-list-tuple that contains each array or: tensor’s shape

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.index_at_time(x, ind)#

Indexes all torch tensors and numpy arrays in dimension 1 with index @ind in nested dictionary or list or tuple and returns a new nested structure.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
ind (int) – index

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.join_dimensions(x, begin_axis, end_axis)#

Joins all dimensions between dimensions (@begin_axis, @end_axis) into a flat dimension, for all tensors in nested dictionary or list or tuple.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
begin_axis (int) – begin dimension
end_axis (int) – end dimension

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.list_of_flat_dict_to_dict_of_list(list_of_dict)#

Helper function to go from a list of flat dictionaries to a dictionary of lists. By “flat” we mean that none of the values are dictionaries, but are numpy arrays, floats, etc.

Parameters:: list_of_dict (list) – list of flat dictionaries
Returns:: dictionary of lists
Return type:: dict_of_list (dict)

robomimic.utils.tensor_utils.map_ndarray(x, func)#

Apply function @func to np.ndarray objects in a nested dictionary or list or tuple.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
func (function) – function to apply to each array

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.map_tensor(x, func)#

Apply function @func to torch.Tensor objects in a nested dictionary or list or tuple.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
func (function) – function to apply to each tensor

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.map_tensor_ndarray(x, tensor_func, ndarray_func)#

Apply function @tensor_func to torch.Tensor objects and @ndarray_func to np.ndarray objects in a nested dictionary or list or tuple.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
tensor_func (function) – function to apply to each tensor
ndarray_Func (function) – function to apply to each array

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.named_reduce(x, reduction, dim)#

Reduces all tensors in nested dictionary or list or tuple at a dimension using a named reduction function.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
reduction (str) – one of [“sum”, “max”, “mean”, “flatten”]
dim (int) – dimension to be reduced (or begin axis for flatten)

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.named_reduce_single(x, reduction, dim)#

Reduce tensor at a dimension by named reduction functions.

Parameters:

x (torch.Tensor) – tensor to be reduced
reduction (str) – one of [“sum”, “max”, “mean”, “flatten”]
dim (int) – dimension to be reduced (or begin axis for flatten)

Returns:

reduced tensor

Return type:

y (torch.Tensor)

robomimic.utils.tensor_utils.pad_sequence(seq, padding, batched=False, pad_same=True, pad_values=None)#

Pad a nested dictionary or list or tuple of sequence tensors in the time dimension (dimension 1).

Parameters:

seq (dict or list or tuple) – a possibly nested dictionary or list or tuple with tensors of leading dimensions [B, T, …]
padding (tuple) – begin and end padding, e.g. [1, 1] pads both begin and end of the sequence by 1
batched (bool) – if sequence has the batch dimension
pad_same (bool) – if pad by duplicating
pad_values (scalar or (ndarray, Tensor)) – values to be padded if not pad_same

Returns:

padded sequence (dict or list or tuple)

robomimic.utils.tensor_utils.pad_sequence_single(seq, padding, batched=False, pad_same=True, pad_values=None)#

Pad input tensor or array @seq in the time dimension (dimension 1).

Parameters:

seq (np.ndarray or torch.Tensor) – sequence to be padded
padding (tuple) – begin and end padding, e.g. [1, 1] pads both begin and end of the sequence by 1
batched (bool) – if sequence has the batch dimension
pad_same (bool) – if pad by duplicating
pad_values (scalar or (ndarray, Tensor)) – values to be padded if not pad_same

Returns:

padded sequence (np.ndarray or torch.Tensor)

robomimic.utils.tensor_utils.recursive_dict_list_tuple_apply(x, type_func_dict)#

Recursively apply functions to a nested dictionary or list or tuple, given a dictionary of {data_type: function_to_apply}.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
type_func_dict (dict) – a mapping from data types to the functions to be applied for each data type.

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.repeat_by_expand_at(x, repeats, dim)#

Repeat a dimension by combining expand and reshape operations.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
repeats (int) – number of times to repeat the target dimension
dim (int) – dimension to repeat on

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.reshape_dimensions(x, begin_axis, end_axis, target_dims)#

Reshape selected dimensions for all tensors in nested dictionary or list or tuple to a target dimension.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
begin_axis (int) – begin dimension
end_axis (int) – end dimension (inclusive)
target_dims (tuple or list) – target shape for the range of dimensions (@begin_axis, @end_axis)

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.reshape_dimensions_single(x, begin_axis, end_axis, target_dims)#

Reshape selected dimensions in a tensor to a target dimension.

Parameters:

x (torch.Tensor) – tensor to reshape
begin_axis (int) – begin dimension
end_axis (int) – end dimension (inclusive)
target_dims (tuple or list) – target shape for the range of dimensions (@begin_axis, @end_axis)

Returns:

reshaped tensor

Return type:

y (torch.Tensor)

robomimic.utils.tensor_utils.time_distributed(inputs, op, activation=None, inputs_as_kwargs=False, inputs_as_args=False, **kwargs)#

Apply function @op to all tensors in nested dictionary or list or tuple @inputs in both the batch (B) and time (T) dimension, where the tensors are expected to have shape [B, T, …]. Will do this by reshaping tensors to [B * T, …], passing through the op, and then reshaping outputs to [B, T, …].

Parameters:

inputs (list or tuple or dict) – a possibly nested dictionary or list or tuple with tensors of leading dimensions [B, T, …]
op – a layer op that accepts inputs
activation – activation to apply at the output
inputs_as_kwargs (bool) – whether to feed input as a kwargs dict to the op
inputs_as_args (bool) –
kwargs (dict) – other kwargs to supply to the op

Returns:

new nested dict-list-tuple with tensors of leading dimension [B, T].

Return type:

outputs (dict or list or tuple)

robomimic.utils.tensor_utils.to_batch(x)#

Introduces a leading batch dimension of 1 for all torch tensors and numpy arrays in nested dictionary or list or tuple and returns a new nested structure.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.to_device(x, device)#

Sends all torch tensors in nested dictionary or list or tuple to device @device, and returns a new nested structure.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
device (torch.Device) – device to send tensors to

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.to_float(x)#

Converts all torch tensors and numpy arrays in nested dictionary or list or tuple to float type entries, and returns a new nested structure.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.to_list(x)#

Converts all torch tensors and numpy arrays in nested dictionary or list or tuple to a list, and returns a new nested structure. Useful for json encoding.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.to_numpy(x)#

Converts all torch tensors in nested dictionary or list or tuple to numpy (and leaves existing numpy arrays as-is), and returns a new nested structure.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.to_one_hot(tensor, num_class)#

Convert all tensors in nested dictionary or list or tuple to one-hot representation, assuming a certain number of total class labels.

Parameters:

tensor (dict or list or tuple) – a possibly nested dictionary or list or tuple
num_class (int) – number of classes

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.to_one_hot_single(tensor, num_class)#

Convert tensor to one-hot representation, assuming a certain number of total class labels.

Parameters:

tensor (torch.Tensor) – tensor containing integer labels
num_class (int) – number of classes

Returns:

tensor containing one-hot representation of labels

Return type:

x (torch.Tensor)

robomimic.utils.tensor_utils.to_sequence(x)#

Introduces a time dimension of 1 at dimension 1 for all torch tensors and numpy arrays in nested dictionary or list or tuple and returns a new nested structure.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.to_tensor(x)#

Converts all numpy arrays in nested dictionary or list or tuple to torch tensors (and leaves existing torch Tensors as-is), and returns a new nested structure.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.to_torch(x, device)#

Converts all numpy arrays and torch tensors in nested dictionary or list or tuple to torch tensors on device @device and returns a new nested structure.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
device (torch.Device) – device to send tensors to

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.to_uint8(x)#

Converts all torch tensors and numpy arrays in nested dictionary or list or tuple to uint8 type entries, and returns a new nested structure.

Parameters:: x (dict or list or tuple) – a possibly nested dictionary or list or tuple
Returns:: new nested dict-list-tuple
Return type:: y (dict or list or tuple)

robomimic.utils.tensor_utils.unsqueeze(x, dim)#

Adds dimension of size 1 at dimension @dim in all torch tensors and numpy arrays in nested dictionary or list or tuple and returns a new nested structure.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
dim (int) – dimension

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.tensor_utils.unsqueeze_expand_at(x, size, dim)#

Unsqueeze and expand a tensor at a dimension @dim by @size.

Parameters:

x (dict or list or tuple) – a possibly nested dictionary or list or tuple
size (int) – size to expand
dim (int) – dimension to unsqueeze and expand

Returns:

new nested dict-list-tuple

Return type:

y (dict or list or tuple)

robomimic.utils.test_utils module#

Utilities for testing algorithm implementations - used mainly by scripts in tests directory.

robomimic.utils.test_utils.checkpoint_path_from_test_run()#: Helper function that gets the path of a model checkpoint after a test training run is finished.

robomimic.utils.test_utils.config_from_modifier(base_config, config_modifier)#

Helper function to load a base config, modify it using the passed @config modifier function, and finalize it for training.

Parameters:

base_config (BaseConfig instance) – starting config object that is loaded (to change algorithm config defaults), and then modified with @config_modifier
config_modifier (function) – function that takes a config object as input, and modifies it

robomimic.utils.test_utils.example_dataset_path()#: Path to dataset to use for testing and example purposes. It should exist under the tests/assets directory, and will be downloaded from a server if it does not exist.

robomimic.utils.test_utils.example_momart_dataset_path()#: Path to momart dataset to use for testing and example purposes. It should exist under the tests/assets directory, and will be downloaded from a server if it does not exist.

robomimic.utils.test_utils.get_base_config(algo_name)#

Base config for testing algorithms.

Parameters:: algo_name (str) – name of algorithm - loads the corresponding json from the config templates directory

robomimic.utils.test_utils.maybe_remove_dir(dir_to_remove)#

Remove directory if it exists.

Parameters:: dir_to_remove (str) – path to directory to remove

robomimic.utils.test_utils.maybe_remove_file(file_to_remove)#

Remove file if it exists.

Parameters:: file_to_remove (str) – path to file to remove

robomimic.utils.test_utils.temp_dataset_path()#: Defines default dataset path to write to for testing.

robomimic.utils.test_utils.temp_model_dir_path()#: Path to a temporary model directory to write to for testing and example purposes.

robomimic.utils.test_utils.temp_video_path()#: Defines default video path to write to for testing.

robomimic.utils.test_utils.test_eval_agent_from_checkpoint(ckpt_path, device)#

Test loading a model from checkpoint and running a rollout with the trained agent for a small number of steps.

Parameters:

ckpt_path (str) – path to a checkpoint pth file
device (torch.Device) – torch device

robomimic.utils.test_utils.test_run(base_config, config_modifier)#

Takes a base_config and config_modifier (function that modifies a passed Config object) and runs training as a test. It also takes the trained checkpoint, tries to load the policy and environment from the checkpoint, and run an evaluation rollout. Returns a string that is colored green if the run finished successfully without any issues, and colored red if an error occurred. If an error occurs, the traceback is included in the string.

Parameters:

base_config (BaseConfig instance) – starting config object that is loaded (to change algorithm config defaults), and then modified with @config_modifier
config_modifier (function) – function that takes a config object as input, and modifies it

Returns:

a green “passed!” string, or a red “failed with error” string that contains: the traceback

Return type:

ret (str)

robomimic.utils.torch_utils module#

This file contains some PyTorch utilities.

robomimic.utils.torch_utils.backprop_for_loss(net, optim, loss, max_grad_norm=None, retain_graph=False)#

Backpropagate loss and update parameters for network with name @name.

Parameters:

net (torch.nn.Module) – network to update
optim (torch.optim.Optimizer) – optimizer to use
loss (torch.Tensor) – loss to use for backpropagation
max_grad_norm (float) – if provided, used to clip gradients
retain_graph (bool) – if True, graph is not freed after backward call

Returns:

average gradient norms from backpropagation

Return type:

grad_norms (float)

class robomimic.utils.torch_utils.dummy_context_mgr#

Bases: object

A dummy context manager - useful for having conditional scopes (such as @maybe_no_grad). Nothing happens in this scope.

robomimic.utils.torch_utils.get_torch_device(try_to_use_cuda)#

Return torch device. If using cuda (GPU), will also set cudnn.benchmark to True to optimize CNNs.

Parameters:: try_to_use_cuda (bool) – if True and cuda is available, will use GPU
Returns:: device to use for models
Return type:: device (torch.Device)

robomimic.utils.torch_utils.hard_update(source, target)#

Hard update @target parameters to match @source.

Parameters:

source (torch.nn.Module) – source network to provide parameters
target (torch.nn.Module) – target network to update parameters for

robomimic.utils.torch_utils.lr_scheduler_from_optim_params(net_optim_params, net, optimizer)#

Helper function to return a LRScheduler from the optim_params section of the config for a particular network. Returns None if a scheduler is not needed.

Parameters:

optim_params (Config) – optim_params part of algo_config corresponding to @net. This determines whether a learning rate scheduler is created.
net (torch.nn.Module) – module whose parameters this optimizer will be responsible
optimizer (torch.optim.Optimizer) – optimizer for this net

Returns:

learning rate scheduler

Return type:

lr_scheduler (torch.optim.lr_scheduler or None)

robomimic.utils.torch_utils.maybe_no_grad(no_grad)#

Parameters:: no_grad (bool) – if True, the returned context will be torch.no_grad(), otherwise it will be a dummy context

robomimic.utils.torch_utils.optimizer_from_optim_params(net_optim_params, net)#

Helper function to return a torch Optimizer from the optim_params section of the config for a particular network.

Parameters:

optim_params (Config) – optim_params part of algo_config corresponding to @net. This determines the optimizer that is created.
net (torch.nn.Module) – module whose parameters this optimizer will be responsible

Returns:

optimizer

Return type:

optimizer (torch.optim.Optimizer)

robomimic.utils.torch_utils.reparameterize(mu, logvar)#

Reparameterize for the backpropagation of z instead of q. This makes it so that we can backpropagate through the sampling of z from our encoder when feeding the sampled variable to the decoder.

(See “The reparameterization trick” section of https://arxiv.org/abs/1312.6114)

Parameters:

mu (torch.Tensor) – batch of means from the encoder distribution
logvar (torch.Tensor) – batch of log variances from the encoder distribution

Returns:

batch of sampled latents from the encoder distribution that: support backpropagation

Return type:

z (torch.Tensor)

robomimic.utils.torch_utils.soft_update(source, target, tau)#

Soft update from the parameters of a @source torch module to a @target torch module with strength @tau. The update follows target = target * (1 - tau) + source * tau.

Parameters:

source (torch.nn.Module) – source network to push target network parameters towards
target (torch.nn.Module) – target network to update

robomimic.utils.train_utils module#

This file contains several utility functions used to define the main training loop. It mainly consists of functions to assist with logging, rollouts, and the @run_epoch function, which is the core training logic for models in this repository.

robomimic.utils.train_utils.dataset_factory(config, obs_keys, filter_by_attribute=None, dataset_path=None)#

Create a SequenceDataset instance to pass to a torch DataLoader.

Parameters:

config (BaseConfig instance) – config object
obs_keys (list) – list of observation modalities that are required for training (this will inform the dataloader on what modalities to load)
filter_by_attribute (str) – if provided, use the provided filter key to select a subset of demonstration trajectories to load
dataset_path (str) – if provided, the SequenceDataset instance should load data from this dataset path. Defaults to config.train.data.

Returns:

dataset object

Return type:

dataset (SequenceDataset instance)

robomimic.utils.train_utils.get_exp_dir(config, auto_remove_exp_dir=False)#

Create experiment directory from config. If an identical experiment directory exists and @auto_remove_exp_dir is False (default), the function will prompt the user on whether to remove and replace it, or keep the existing one and add a new subdirectory with the new timestamp for the current run.

Parameters:

auto_remove_exp_dir (bool) – if True, automatically remove the existing experiment folder if it exists at the same path.

Returns:

path to created log directory (sub-folder in experiment directory) output_dir (str): path to created models directory (sub-folder in experiment directory)

to store model checkpoints

video_dir (str): path to video directory (sub-folder in experiment directory): to store rollout videos

Return type:

log_dir (str)

robomimic.utils.train_utils.is_every_n_steps(interval, current_step, skip_zero=False)#

Convenient function to check whether current_step is at the interval. Returns True if current_step % interval == 0 and asserts a few corner cases (e.g., interval <= 0)

Parameters:

interval (int) – target interval
current_step (int) – current step
skip_zero (bool) – whether to skip 0 (return False at 0)

Returns:

whether current_step is at the interval

Return type:

is_at_interval (bool)

robomimic.utils.train_utils.load_data_for_training(config, obs_keys)#

Data loading at the start of an algorithm.

Parameters:

config (BaseConfig instance) – config object
obs_keys (list) – list of observation modalities that are required for training (this will inform the dataloader on what modalities to load)

Returns:

train dataset object valid_dataset (SequenceDataset instance): valid dataset object (only if using validation)

Return type:

train_dataset (SequenceDataset instance)

robomimic.utils.train_utils.rollout_with_stats(policy, envs, horizon, use_goals=False, num_episodes=None, render=False, video_dir=None, video_path=None, epoch=None, video_skip=5, terminate_on_success=False, verbose=False)#

A helper function used in the train loop to conduct evaluation rollouts per environment and summarize the results.

Can specify @video_dir (to dump a video per environment) or @video_path (to dump a single video for all environments).

Parameters:

policy (RolloutPolicy instance) – policy to use for rollouts.
envs (dict) – dictionary that maps env_name (str) to EnvBase instance. The policy will be rolled out in each env.
horizon (int) – maximum number of steps to roll the agent out for
use_goals (bool) – if True, agent is goal-conditioned, so provide goal observations from env
num_episodes (int) – number of rollout episodes per environment
render (bool) – if True, render the rollout to the screen
video_dir (str) – if not None, dump rollout videos to this directory (one per environment)
video_path (str) – if not None, dump a single rollout video for all environments
epoch (int) – epoch number (used for video naming)
video_skip (int) – how often to write video frame
terminate_on_success (bool) – if True, terminate episode early as soon as a success is encountered
verbose (bool) – if True, print results of each rollout

Returns:

dictionary of rollout statistics (e.g. return, success rate, …): averaged across all rollouts

video_paths (dict): path to rollout videos for each environment

Return type:

all_rollout_logs (dict)

robomimic.utils.train_utils.run_epoch(model, data_loader, epoch, validate=False, num_steps=None, obs_normalization_stats=None)#

Run an epoch of training or validation.

Parameters:

model (Algo instance) – model to train
data_loader (DataLoader instance) – data loader that will be used to serve batches of data to the model
epoch (int) – epoch number
validate (bool) – whether this is a training epoch or validation epoch. This tells the model whether to do gradient steps or purely do forward passes.
num_steps (int) – if provided, this epoch lasts for a fixed number of batches (gradient steps), otherwise the epoch is a complete pass through the training dataset
obs_normalization_stats (dict or None) – if provided, this should map observation keys to dicts with a “mean” and “std” of shape (1, …) where … is the default shape for the observation.

Returns:

dictionary of logged training metrics averaged across all batches

Return type:

step_log_all (dict)

robomimic.utils.train_utils.run_rollout(policy, env, horizon, use_goals=False, render=False, video_writer=None, video_skip=5, terminate_on_success=False)#

Runs a rollout in an environment with the current network parameters.

Parameters:

policy (RolloutPolicy instance) – policy to use for rollouts.
env (EnvBase instance) – environment to use for rollouts.
horizon (int) – maximum number of steps to roll the agent out for
use_goals (bool) – if True, agent is goal-conditioned, so provide goal observations from env
render (bool) – if True, render the rollout to the screen
video_writer (imageio Writer instance) – if not None, use video writer object to append frames at rate given by @video_skip
video_skip (int) – how often to write video frame
terminate_on_success (bool) – if True, terminate episode early as soon as a success is encountered

Returns:

dictionary containing return, success rate, etc.

Return type:

results (dict)

robomimic.utils.train_utils.save_model(model, config, env_meta, shape_meta, ckpt_path, obs_normalization_stats=None)#

Save model to a torch pth file.

Parameters:

model (Algo instance) – model to save
config (BaseConfig instance) – config to save
env_meta (dict) – env metadata for this training run
shape_meta (dict) – shape metdata for this training run
ckpt_path (str) – writes model checkpoint to this path
obs_normalization_stats (dict) – optionally pass a dictionary for observation normalization. This should map observation keys to dicts with a “mean” and “std” of shape (1, …) where … is the default shape for the observation.

robomimic.utils.train_utils.should_save_from_rollout_logs(all_rollout_logs, best_return, best_success_rate, epoch_ckpt_name, save_on_best_rollout_return, save_on_best_rollout_success_rate)#

Helper function used during training to determine whether checkpoints and videos should be saved. It will modify input attributes appropriately (such as updating the best returns and success rates seen and modifying the epoch ckpt name), and returns a dict with the updated statistics.

Parameters:

all_rollout_logs (dict) – dictionary of rollout results that should be consistent with the output of @rollout_with_stats
best_return (dict) – dictionary that stores the best average rollout return seen so far during training, for each environment
best_success_rate (dict) – dictionary that stores the best average success rate seen so far during training, for each environment
epoch_ckpt_name (str) – what to name the checkpoint file - this name might be modified by this function
save_on_best_rollout_return (bool) – if True, should save checkpoints that achieve a new best rollout return
save_on_best_rollout_success_rate (bool) – if True, should save checkpoints that achieve a new best rollout success rate

Returns:

dictionary that contains updated input attributes @best_return,: @best_success_rate, @epoch_ckpt_name, along with two additional attributes @should_save_ckpt (True if should save this checkpoint), and @ckpt_reason (string that contains the reason for saving the checkpoint)

Return type:

save_info (dict)

robomimic.utils.vis_utils module#

This file contains utility functions for visualizing image observations in the training pipeline. These functions can be a useful debugging tool.

robomimic.utils.vis_utils.depth_to_rgb(depth_map, depth_min=None, depth_max=None)#: Convert depth map to rgb array by computing normalized depth values in [0, 1].

robomimic.utils.vis_utils.image_tensor_to_disk(image, fname)#

Writes an image tensor to disk. Any leading batch dimensions are indexed out with the first element.

Parameters:

image (torch.Tensor) – image of shape […, C, H, W]. All leading dimensions will be indexed out with the first element
fname (str) – path to save image to

robomimic.utils.vis_utils.image_tensor_to_numpy(image)#

Converts processed image tensors to numpy so that they can be saved to disk or video. A useful utility function for visualizing images in the middle of training.

Parameters:: image (torch.Tensor) – images of shape […, C, H, W]
Returns:: converted images of shape […, H, W, C] and type uint8
Return type:: image (np.array)

robomimic.utils.vis_utils.image_to_disk(image, fname)#

Writes an image to disk.

Parameters:

image (np.array) – image of shape [H, W, 3]
fname (str) – path to save image to

robomimic.utils.vis_utils.visualize_image_randomizer(original_image, randomized_image, randomizer_name=None)#

A function that visualizes the before and after of an image-based input randomizer :param original_image: batch of original image shaped [B, H, W, 3] :param randomized_image: randomized image shaped [B, N, H, W, 3]. N is the number of randomization per input sample :param randomizer_name: (Optional) name of the randomizer

Returns:: None

robomimic 0.3 documentation

robomimic.utils package

Contents

robomimic.utils package#

Submodules#

robomimic.utils.dataset module#

robomimic.utils.env_utils module#

robomimic.utils.file_utils module#

robomimic.utils.hyperparam_utils module#

robomimic.utils.log_utils module#

robomimic.utils.loss_utils module#

robomimic.utils.obs_utils module#

robomimic.utils.python_utils module#

robomimic.utils.tensor_utils module#

robomimic.utils.test_utils module#

robomimic.utils.torch_utils module#

robomimic.utils.train_utils module#

robomimic.utils.vis_utils module#

Module contents#