Ensure that array or tensor @x has size @size in dim @dim. Args: x (np.ndarray or torch.Tensor): input array or tensor size

Load to memory.

Create a multi layer perceptron (MLP), which is a collection of fully-connected layers each followed by an activation function. :param i

Expand all tensors in nested dictionary or list or tuple at a single dimension @dim by @size. Args: x (dict or list or tuple): a

Expand a tensor at a single dimension @dim by @size Args: x (torch.Tensor): input tensor size (int): size to expand

Flatten all tensors in nested dictionary or list or tuple, from @begin_axis onwards. Args: x (dict or list or tuple): a possibly nes

Flatten a nested dict or list to a list. For example, given a dict { a: 1 b: { c: 2 } c: 3

Flatten a tensor in all dimensions from @begin_axis onwards. Args: x (torch.Tensor): tensor to flatten begin_axis (int): whi

Apply @gather_along_dim_with_dim_single to all tensors in a nested dictionary or list or tuple. Args: x (dict or list or tuple)

Compute the decay factor for the exponential moving average.

Common shapes T,D T,N,D T,H,W,C T,N,H,W,C

Added kwargs vs diffuser's original implementation Unified API to get any scheduler from its name. Args: name (`str` or `Schedu

Reduce tensor at a dimension by named reduction functions. Args: x (torch.Tensor): tensor to be reduced reduction (str): one

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

Reshape selected dimensions for all tensors in nested dictionary or list or tuple to a target dimension. Args: x (dict or l

Quick loading and saving for reserach, saves full state of the workspace. However, loading a snapshot assumes the code stays exactly

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

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

Convert tensor to one-hot representation, assuming a certain number of total class labels. Args: tensor (torch.Tensor): tensor conta

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 nest

Unsqueeze and expand a tensor at a dimension @dim by @size. Args: x (dict or list or tuple): a possibly nested dictionary or list or

output: obs: key: T, * action: T, Da

When file_path is None, don't record.

@crowsonkb's notes on EMA Warmup: If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for mod