Make encoder networks.

Creates params and inference function for the encoder agent.

Convert rotations given as rotation matrices to axis/angle. Args: matrix: Rotation matrices as tensor of shape (..., 3, 3). Ret

Converts rotation matrices to 6D rotation representation by Zhou et al. [1] by dropping the last row. Note that 6D representation is not uniq

Extended MSE Loss with GMOF.

Prepare raw data for the f'{idx'}-th data.

Check output folder or file. Args: output_path (str): could be folder or file. allowed_suffix (List[str], optional):

Construct identity 3D rotation with the same shape

Rotate a 3D vector with the 3D rotation

Apply the rotation given by a quaternion to a 3D point. Usual torch rules for broadcasting apply. Args: quaternion: Tensor of qu

Multiply two quaternions representing rotations, returning the quaternion representing their composition, i.e. the versor with nonnegative re

Generate a single random 3x3 rotation matrix. Args: dtype: Type to return device: Device of returned tensor. Default: if Non

Resets the environment to an initial state.

Resets the environment to an initial state.

Converts 6D rotation representation by Zhou et al. [1] to rotation matrix using Gram--Schmidt orthogonalization per Section B of [1]. Arg

Importance-sample timesteps for a batch. :param batch_size: the number of timesteps. :param device: the torch device to save

Run one timestep of the environment's dynamics.

Direct trajectory optimization training.

Train step function. Args: inputs (Tensor): Input Tensor. kwargs (dict): Args for ``mmcv.parallel.sca

The iteration step during training. This method defines an iteration step during training, except for the back propagation and optimiz

Update the reweighting using losses from a model. Call this method from each rank with a batch of timesteps and the correspon

Validation step function. Args: inputs (tuple): Input data. kwargs (dict): Args for ``scatter_kwargs``.

The iteration step during validation. This method shares the same signature as :func:`train_step`, but used during val epochs. Note th

Create a weighted version of a given loss function. To use this decorator, the loss function must have the signature like `loss_func(pred, tar

Init random seed for each worker.