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

Compute the average of the largest K scalar values recorded. Returns: - avg: Average of the largest K scalars.

Apply function @func to torch.Tensor objects in a nested dictionary or list or tuple. Args: x (dict or list or tuple): a possibl

Record a scalar value. Parameters: - scalar: The scalar value to be recorded.

Copy a on-disk zarr to in-memory compressed. Recommended

obs_dict: must include "obs" key result: must include "action" key

Reshape selected dimensions in a tensor to a target dimension. Args: x (torch.Tensor): tensor to reshape begin_axis (int): b

Print the number of parameters in each part of the model.

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

Output (n_steps,) + obs_shape

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

return the underlying mujoco sim

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

Pad input tensor or array @seq in the time dimension (dimension 1). Args: seq (np.ndarray or torch.Tensor): sequence to be padded

support different point cloud sampling methods point_cloud: (N, 6), xyz+rgb or (N, 3), xyz