conditional_action_unet

conditional_action_unet

Conditional U-Net Decoder for action generation. Reference implementation: Diffusion Policy (https://arxiv.org/abs/2303.04137)

ConditionalActionUNet

ConditionalActionUNet(input_keys, action_space, action_heads, observation_space, observation_horizon, prediction_horizon, device, embedding_dimension=256, down_dimensions=None, kernel_size=5, num_groups=8, condition_predict_scale=False)

Bases: ActionDecoder

Conditional U-Net decoder for generative action generation.

This architecture: - Uses FiLM (Feature-wise Linear Modulation) for conditioning - Accepts global conditioning from concatenated observation features - Optionally supports local (sequence-aligned) conditioning - Designed for use with Diffusion or Flow Matching algorithms

The decoder expects: - Noisy actions as input (via actions parameter during forward) - Timesteps injected by the algorithm (via features["timestep"]) - Observation features for global conditioning (via features dict)

Note: This decoder is specifically designed for diffusion/flow matching algorithms and expects the algorithm to handle noise scheduling and timestep injection.

Initialize Conditional U-Net decoder.

Parameters:

Name	Type	Description	Default
input_keys	list[str]	List of feature keys expected from encoder pipeline	required
action_space	ActionSpace	Action space configuration	required
action_heads	dict[str, ActionHead]	Must be empty. The U-Net owns its final convolution.	required
observation_space	ObservationSpace	Observation space configuration	required
observation_horizon	int	Number of observation timesteps (for history)	required
prediction_horizon	int	Number of actions to predict (horizon)	required
device	str	Device to run the model on	required
embedding_dimension	int	Diffusion timestep embedding dimension	256
down_dimensions	list[int] | None	List of channel dimensions for downsampling layers	None
kernel_size	int	Kernel size for convolutions in residual blocks	5
num_groups	int	Number of groups for group normalization	8
condition_predict_scale	bool	If True, conditions predict scaling factors in FiLM	False

Source code in src/versatil/models/decoding/decoders/factory/conditional_action_unet.py

def __init__(
    self,
    input_keys: list[str],
    action_space: ActionSpace,
    action_heads: dict[str, ActionHead],
    observation_space: ObservationSpace,
    observation_horizon: int,
    prediction_horizon: int,
    device: str,
    embedding_dimension: int = 256,
    down_dimensions: list[int] | None = None,
    kernel_size: int = 5,
    num_groups: int = 8,
    condition_predict_scale: bool = False,
) -> None:
    """Initialize Conditional U-Net decoder.

    Args:
        input_keys: List of feature keys expected from encoder pipeline
        action_space: Action space configuration
        action_heads: Must be empty. The U-Net owns its final convolution.
        observation_space: Observation space configuration
        observation_horizon: Number of observation timesteps (for history)
        prediction_horizon: Number of actions to predict (horizon)
        device: Device to run the model on
        embedding_dimension: Diffusion timestep embedding dimension
        down_dimensions: List of channel dimensions for downsampling layers
        kernel_size: Kernel size for convolutions in residual blocks
        num_groups: Number of groups for group normalization
        condition_predict_scale: If True, conditions predict scaling factors in FiLM

    """
    decoder_input = DecoderInput(
        keys=input_keys,
        raises_for_types=[FeatureType.SPATIAL.value],
        requires_actions=True,
    )
    super().__init__(
        decoder_input=decoder_input,
        action_space=action_space,
        action_heads=action_heads,
        observation_space=observation_space,
        prediction_horizon=prediction_horizon,
        observation_horizon=observation_horizon,
        device=device,
    )

    if down_dimensions is None:
        down_dimensions = [256, 512, 1024]

    self.embedding_dimension = embedding_dimension
    self.down_dimensions = down_dimensions
    self.kernel_size = kernel_size
    self.num_groups = num_groups
    self.condition_predict_scale = condition_predict_scale

    self._global_conditioning_dimension: int | None = None
    self._feature_projections: nn.ModuleDict | None = None
    self.unet_conditioning_builder = UNetInputBuilder(
        embedding_dimension=embedding_dimension,
    )

    # U-Net will be lazily initialized on first forward pass
    # (once we know the global conditioning dimension)
    self._unet: ConditionalUnet1D | None = None

forward

forward(features, actions=None)

Forward pass through the conditional U-Net.

This method is called by the decoding algorithm (Diffusion, FlowMatching) which provides: - Noisy actions - Observation features dictionary containing the timestep key

Parameters:

Name	Type	Description	Default
features	dict[str, Tensor]	Dictionary of encoded features from the encoding pipeline.	required
actions	dict[str, Tensor] | None	Dictionary of noise-injected actions (provided by algorithm during training)	None

Returns:

Type	Description
dict[str, Tensor]	Dictionary containing denoised predictions for each action head

Raises:

Type	Description
ValueError	If timesteps or actions are missing.

Source code in src/versatil/models/decoding/decoders/factory/conditional_action_unet.py

def forward(
    self,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Forward pass through the conditional U-Net.

    This method is called by the decoding algorithm (Diffusion, FlowMatching)
    which provides:
    - Noisy actions
    - Observation features dictionary containing the timestep key

    Args:
        features: Dictionary of encoded features from the encoding pipeline.
        actions: Dictionary of noise-injected actions (provided by algorithm during training)

    Returns:
        Dictionary containing denoised predictions for each action head

    Raises:
        ValueError: If timesteps or actions are missing.
    """
    if actions is None:
        raise ValueError(
            "ConditionalActionUNet requires 'actions' parameter. "
            "The algorithm should provide noisy actions during forward pass."
        )

    noisy_actions = self.action_space.concatenate_action_tensors(
        actions=actions,
        prediction_horizon=self.prediction_horizon,
        owner_name=self.__class__.__name__,
    )
    timesteps = extract_timestep_conditioning(
        features=features,
        batch_size=noisy_actions.shape[0],
        action_device=noisy_actions.device,
    )  # (B,)
    observation_features = filter_timestep_feature(features=features)

    # Prepare global conditioning
    global_conditioning = self._prepare_global_conditioning(
        observation_features
    )  # (B, global_conditioning_dimension)

    # Run U-Net denoising
    if self._unet is None:
        raise RuntimeError(
            "U-Net should be initialized by now. "
            "Call _prepare_global_conditioning before running the U-Net."
        )
    denoised = self._unet(
        noisy_input=noisy_actions,
        timesteps=timesteps,
        local_conditioning=None,  # Not using local conditioning
        global_conditioning=global_conditioning,
    )  # (B, T, action_dimension)

    return self.action_space.split_action_tensor(
        action_tensor=denoised,
        owner_name=self.__class__.__name__,
    )