toddlerbot.manipulation.models package

Submodules

toddlerbot.manipulation.models.diffusion_model module

Diffusion model architecture for 1D temporal sequences.

This module implements a 1D UNet architecture for noise prediction in diffusion models.

Components: - SinusoidalPosEmb: Positional encoding for diffusion timesteps - Downsample1d/Upsample1d: Temporal resolution change layers - Conv1dBlock: Conv1d -> GroupNorm -> Mish activation - ConditionalResidualBlock1D: Residual block with FiLM conditioning

class toddlerbot.manipulation.models.diffusion_model.ConditionalResidualBlock1D(in_channels, out_channels, cond_dim, kernel_size=3, n_groups=8)

Bases: Module

forward(x, cond)

x : [ batch_size x in_channels x horizon ] cond : [ batch_size x cond_dim]

returns: out : [ batch_size x out_channels x horizon ]

class toddlerbot.manipulation.models.diffusion_model.ConditionalUnet1D(input_dim, global_cond_dim, diffusion_step_embed_dim=256, down_dims=[256, 512, 1024], kernel_size=5, n_groups=8)

Bases: Module

forward(sample: Tensor, timestep: Tensor | float | int, global_cond=None)

x: (B,T,input_dim) timestep: (B,) or int, diffusion step global_cond: (B,global_cond_dim) output: (B,T,input_dim)

class toddlerbot.manipulation.models.diffusion_model.Conv1dBlock(inp_channels, out_channels, kernel_size, n_groups=8)

Bases: Module

Conv1d –> GroupNorm –> Mish

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class toddlerbot.manipulation.models.diffusion_model.Downsample1d(dim)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class toddlerbot.manipulation.models.diffusion_model.SinusoidalPosEmb(dim)

Bases: Module

Sinusoidal positional embedding for diffusion timesteps.

forward(x)

Generate sinusoidal positional embeddings.

class toddlerbot.manipulation.models.diffusion_model.Upsample1d(dim)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Module contents

Manipulation models module for ToddlerBot.