Utilities for training and sampling diffusion models. Ported directly from here, and then adapted over time to further experimentation.

Contains training state for the learner.

Compose a data pipeline with a sequence of transforms. Args: transforms (list[dict | callable]): Either config dicts of trans

A wrapper of concatenated dataset. Same as :obj:`torch.utils.data.dataset.ConcatDataset`, but add `get_cat_ids` function. Args: da

A DistributedDataParallel wrapper for models in 3D mesh estimation task. In 3D mesh estimation task, there is a need to wrap different modules i

A wrapper of repeated dataset. The length of repeated dataset will be `times` larger than the original dataset. This is useful when the data l

A diffusion process which can skip steps in a base diffusion process. :param use_timesteps: a collection (sequence or set) of timesteps from

Automatically resets Brax envs that are done.

Automatically resets Brax envs that are done.

Base class for mmhuman3d architecture.

Base motion dataset. Args: data_prefix (str): the prefix of data path. pipeline (list): a list of dict, where each element represe

Collect data from the loader relevant to the specific task. This is usually the last stage of the data loader pipeline. Args: keys

Contrastive loss function. Based on: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf

r"""Crop motion sequences. Args: crop_size (int): The size of the cropped motion sequence.

Maintains episode step count and sets done at episode end.

Maintains episode step count and sets done at episode end.

Brax env with eval metrics.

Brax env with eval metrics.

Class to run evaluations.

State of file existence.

Trains a humanoid to mimic reference motion.

Trains a humanoid to mimic reference motion.

Trains a humanoid to mimic reference motion.

MSELoss. Args: reduction (str, optional): The method that reduces the loss to a scalar. Options are "none", "mean" and "sum".

Which type of output the model predicts.

What is used as the model's output variance. The LEARNED_RANGE option has been added to allow the model to predict values between FIXED_

Normalize motion sequences. Args: mean_path (str): Path of mean file. std_path (str): Path of std file.

r"""Random crop motion sequences. Each sequence will be padded with zeros to the maximum length. Args: min_size (int or None): The mi

A distribution over timesteps in the diffusion process, intended to reduce variance of the objective. By default, samplers perform unbias

TextMotion dataset. Args: text_dir (str): Path to the directory containing the text files.

Vectorizes Brax env.

Vectorizes Brax env.

Wrap fields of the data dictionary into lists for evaluation. This class can be used as a last step of a test or validation pipeline for sing