toddlerbot.manipulation package

Subpackages

• toddlerbot.manipulation.datasets package
  • Submodules
  • toddlerbot.manipulation.datasets.get_pose_data module
    • process_raw_dataset()
  • toddlerbot.manipulation.datasets.process_raw_data module
    • process_raw_dataset()
  • toddlerbot.manipulation.datasets.teleop_dataset module
    • TeleopImageDataset
  • Module contents
• toddlerbot.manipulation.models package
  • Submodules
  • toddlerbot.manipulation.models.diffusion_model module
    • ConditionalResidualBlock1D
      • ConditionalResidualBlock1D.forward()
    • ConditionalUnet1D
      • ConditionalUnet1D.forward()
    • Conv1dBlock
      • Conv1dBlock.forward()
    • Downsample1d
      • Downsample1d.forward()
    • SinusoidalPosEmb
      • SinusoidalPosEmb.forward()
    • Upsample1d
      • Upsample1d.forward()
  • Module contents
• toddlerbot.manipulation.utils package
  • Submodules
  • toddlerbot.manipulation.utils.dataset_utils module
    • create_sample_indices()
    • create_video_grid()
    • get_data_stats()
    • normalize_data()
    • sample_sequence()
    • unnormalize_data()
  • toddlerbot.manipulation.utils.model_utils module
    • get_resnet()
    • replace_bn_with_gn()
    • replace_submodules()
  • Module contents

Submodules

toddlerbot.manipulation.inference_class module

Diffusion policy inference class for manipulation tasks.

class toddlerbot.manipulation.inference_class.DPModel(ckpt_path, stats=None)

Bases: object

Diffusion Policy model for manipulation inference.

get_action_from_obs(obs_deque)

Generate an action based on a sequence of observations.

This function processes a deque of observations to prepare inputs, extracts features using a vision encoder, and performs inference to generate an action. The action is derived from denoised samples obtained through a diffusion model.

Parameters:

obs_deque (collections.deque) – A deque containing the sequence of observations.

Returns:

The generated action in the required format.

Return type:

torch.Tensor

inference_ddim(obs_cond, nsteps=10, naction=None)

Performs DDIM (Denoising Diffusion Implicit Models) inference to generate actions based on observed conditions.

Parameters:

• obs_cond (torch.Tensor) – The observed conditions used as input to condition the action generation.

• nsteps (int, optional) – The number of diffusion steps to perform. Defaults to 10.

• naction (torch.Tensor, optional) – Initial noisy actions. If None, actions are initialized from Gaussian noise.

Returns:

The denoised actions after performing the specified number of diffusion steps.

Return type:

torch.Tensor

inference_ddpm(obs_cond, nsteps, naction=None)

Performs inference using the Denoising Diffusion Probabilistic Model (DDPM) to generate actions.

This function initializes a noisy action from Gaussian noise and iteratively refines it using a noise prediction network and a noise scheduler. The process involves predicting noise at each timestep and removing it to obtain a cleaner action sample.

Parameters:

• obs_cond (Tensor) – The observation condition tensor used as global conditioning input for the noise prediction network.

• nsteps (int) – The number of diffusion steps to perform during the inference process.

• naction (Tensor, optional) – Initial noisy action tensor. If not provided, it is initialized from Gaussian noise.

Returns:

The refined action tensor after performing the inverse diffusion process.

Return type:

Tensor

load_model(ckpt_path, stats=None)

Loads a pre-trained model from a checkpoint file and initializes the network components.

Parameters:

• ckpt_path (str) – Path to the checkpoint file containing the model’s state dictionary.

• stats (dict, optional) – Pre-computed statistics for the model. If not provided, statistics will be loaded from the checkpoint.

Loads the model’s state dictionary and statistics from the checkpoint, or uses provided statistics. Sets the model to evaluation mode.

prediction_to_action(naction)

Converts a normalized action prediction to a denormalized action sequence.

This method takes a normalized action prediction tensor, detaches it from the computation graph, and converts it to a NumPy array. It then denormalizes the action data using predefined statistics and extracts a sequence of actions based on the specified action horizon.

Parameters:

naction (torch.Tensor) – A tensor containing the normalized action predictions with shape (B, pred_horizon, action_dim).

Returns:

A denormalized action sequence with shape (action_horizon, action_dim).

Return type:

numpy.ndarray

prepare_inputs(obs_deque)

Prepares and normalizes input data for model processing.

This function stacks and normalizes the last set of observations from a deque, transferring them to the specified device for further processing.

Parameters:

obs_deque (collections.deque) – A deque containing the most recent observations, where each observation is a dictionary with keys “image” and “agent_pos”.

Returns:

A tuple containing:

• nimages (torch.Tensor): A tensor of stacked and normalized images, transferred to the specified device.

• nagent_poses (torch.Tensor): A tensor of normalized agent positions, transferred to the specified device.

Return type:

tuple

toddlerbot.manipulation.train module

Train diffusion policy models for robot manipulation tasks.

toddlerbot.manipulation.train.train(dataset_path_list: List[str], exp_folder_path: str, weights: str, pred_horizon: int, obs_horizon: int, action_horizon: int, action_dim: int = 16, vision_feature_dim: int = 512, num_diffusion_iters: int = 100, num_epochs: int = 1000, early_stopping_patience: int = 100, train_split_ratio: float = 0.8)

Trains a neural network model using a dataset of teleoperation images and actions.

Parameters:

• dataset_path_list (List[str]) – List of paths to the datasets.

• exp_folder_path (str) – Path to the folder where experiment outputs will be saved.

• weights (str) – Pre-trained weights for the vision encoder.

• pred_horizon (int) – Prediction horizon for the model.

• obs_horizon (int) – Observation horizon for the model.

• action_horizon (int) – Action horizon for the model.

• action_dim (int, optional) – Dimensionality of the action space. Defaults to 16.

• vision_feature_dim (int, optional) – Dimensionality of the vision feature space. Defaults to 512.

• num_diffusion_iters (int, optional) – Number of diffusion iterations. Defaults to 100.

• num_epochs (int, optional) – Number of training epochs. Defaults to 1000.

• early_stopping_patience (int, optional) – Number of epochs to wait for improvement before stopping early. Defaults to 100.

• train_split_ratio (float, optional) – Ratio of the dataset to use for training. Defaults to 0.8.

Module contents

Manipulation and grasping capabilities for ToddlerBot.

This package provides functionality for robot manipulation tasks, including:

• Teleoperation data collection and processing

• Diffusion-based manipulation policy training

• Dataset utilities for manipulation learning

• Model architectures for action prediction

• Inference and deployment tools for manipulation policies

The manipulation framework supports learning from demonstration through teleoperation, with tools for data collection, processing, and neural network training using diffusion models for robust manipulation policies.

Key components include dataset management, model training pipelines, and inference classes for real-time manipulation control.