213 lines
10 KiB
Python
213 lines
10 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2025 Nur Muhammad Mahi Shafiullah,
|
|
# and The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
from dataclasses import dataclass, field
|
|
|
|
from lerobot.common.optim.optimizers import AdamConfig
|
|
from lerobot.common.optim.schedulers import DiffuserSchedulerConfig
|
|
from lerobot.configs.policies import PreTrainedConfig
|
|
from lerobot.configs.types import NormalizationMode
|
|
|
|
|
|
@PreTrainedConfig.register_subclass("ditflow")
|
|
@dataclass
|
|
class DiTFlowConfig(PreTrainedConfig):
|
|
"""Configuration class for DiTFlowPolicy.
|
|
|
|
Defaults are configured for training with PushT providing proprioceptive and single camera observations.
|
|
|
|
The parameters you will most likely need to change are the ones which depend on the environment / sensors.
|
|
Those are: `input_shapes` and `output_shapes`.
|
|
|
|
Notes on the inputs and outputs:
|
|
- "observation.state" is required as an input key.
|
|
- Either:
|
|
- At least one key starting with "observation.image is required as an input.
|
|
AND/OR
|
|
- The key "observation.environment_state" is required as input.
|
|
- If there are multiple keys beginning with "observation.image" they are treated as multiple camera
|
|
views. Right now we only support all images having the same shape.
|
|
- "action" is required as an output key.
|
|
|
|
Args:
|
|
n_obs_steps: Number of environment steps worth of observations to pass to the policy (takes the
|
|
current step and additional steps going back).
|
|
horizon: DiT-flow model action prediction size as detailed in `DiTFlowPolicy.select_action`.
|
|
n_action_steps: The number of action steps to run in the environment for one invocation of the policy.
|
|
See `DiTFlowPolicy.select_action` for more details.
|
|
input_shapes: A dictionary defining the shapes of the input data for the policy. The key represents
|
|
the input data name, and the value is a list indicating the dimensions of the corresponding data.
|
|
For example, "observation.image" refers to an input from a camera with dimensions [3, 96, 96],
|
|
indicating it has three color channels and 96x96 resolution. Importantly, `input_shapes` doesn't
|
|
include batch dimension or temporal dimension.
|
|
output_shapes: A dictionary defining the shapes of the output data for the policy. The key represents
|
|
the output data name, and the value is a list indicating the dimensions of the corresponding data.
|
|
For example, "action" refers to an output shape of [14], indicating 14-dimensional actions.
|
|
Importantly, `output_shapes` doesn't include batch dimension or temporal dimension.
|
|
input_normalization_modes: A dictionary with key representing the modality (e.g. "observation.state"),
|
|
and the value specifies the normalization mode to apply. The two available modes are "mean_std"
|
|
which subtracts the mean and divides by the standard deviation and "min_max" which rescale in a
|
|
[-1, 1] range.
|
|
output_normalization_modes: Similar dictionary as `normalize_input_modes`, but to unnormalize to the
|
|
original scale. Note that this is also used for normalizing the training targets.
|
|
vision_backbone: Name of the torchvision resnet backbone to use for encoding images.
|
|
crop_shape: (H, W) shape to crop images to as a preprocessing step for the vision backbone. Must fit
|
|
within the image size. If None, no cropping is done.
|
|
crop_is_random: Whether the crop should be random at training time (it's always a center crop in eval
|
|
mode).
|
|
pretrained_backbone_weights: Pretrained weights from torchvision to initalize the backbone.
|
|
`None` means no pretrained weights.
|
|
use_group_norm: Whether to replace batch normalization with group normalization in the backbone.
|
|
The group sizes are set to be about 16 (to be precise, feature_dim // 16).
|
|
spatial_softmax_num_keypoints: Number of keypoints for SpatialSoftmax.
|
|
use_separate_rgb_encoders_per_camera: Whether to use a separate RGB encoder for each camera view.
|
|
|
|
frequency_embedding_dim: The embedding dimension for the time value embedding in the flow model.
|
|
num_blocks: The number of transformer blocks in the DiT flow model.
|
|
hidden_dim: The hidden dimension for the transformer blocks in the DiT flow model.
|
|
num_heads: The number of attention heads in the transformer blocks.
|
|
dropout: The dropout rate used inside the transformer blocks.
|
|
dim_feedforward: The expanded feedforward dimension in the MLPs used in the transformer block.
|
|
activation: The activation function used in the transformer blocks.
|
|
clip_sample: Whether to clip the sample to [-`clip_sample_range`, +`clip_sample_range`] for each
|
|
denoising step at inference time. WARNING: you will need to make sure your action-space is
|
|
normalized to fit within this range.
|
|
clip_sample_range: The magnitude of the clipping range as described above.
|
|
num_inference_steps: Number of reverse diffusion steps to use at inference time (steps are evenly
|
|
spaced).
|
|
do_mask_loss_for_padding: Whether to mask the loss when there are copy-padded actions. See
|
|
`LeRobotDataset` and `load_previous_and_future_frames` for mor information. Note, this defaults
|
|
to False as the original Diffusion Policy implementation does the same.
|
|
"""
|
|
|
|
# Inputs / output structure.
|
|
n_obs_steps: int = 2
|
|
horizon: int = 16
|
|
n_action_steps: int = 8
|
|
|
|
normalization_mapping: dict[str, NormalizationMode] = field(
|
|
default_factory=lambda: {
|
|
"VISUAL": NormalizationMode.MEAN_STD,
|
|
"STATE": NormalizationMode.MIN_MAX,
|
|
"ACTION": NormalizationMode.MIN_MAX,
|
|
}
|
|
)
|
|
|
|
# The original implementation doesn't sample frames for the last 7 steps,
|
|
# which avoids excessive padding and leads to improved training results.
|
|
drop_n_last_frames: int = 7 # horizon - n_action_steps - n_obs_steps + 1
|
|
|
|
# Architecture / modeling.
|
|
# Vision backbone.
|
|
vision_backbone: str = "resnet18"
|
|
crop_shape: tuple[int, int] | None = (84, 84)
|
|
crop_is_random: bool = True
|
|
pretrained_backbone_weights: str | None = None
|
|
use_group_norm: bool = True
|
|
spatial_softmax_num_keypoints: int = 32
|
|
use_separate_rgb_encoder_per_camera: bool = False
|
|
|
|
# Diffusion Transformer (DiT) parameters.
|
|
frequency_embedding_dim: int = 256
|
|
hidden_dim: int = 512
|
|
num_blocks: int = 6
|
|
num_heads: int = 16
|
|
dropout: float = 0.1
|
|
dim_feedforward: int = 4096
|
|
activation: str = "gelu"
|
|
|
|
# Noise scheduler.
|
|
training_noise_sampling: str = (
|
|
"uniform" # "uniform" or "beta", from pi0 https://www.physicalintelligence.company/download/pi0.pdf
|
|
)
|
|
clip_sample: bool = True
|
|
clip_sample_range: float = 1.0
|
|
|
|
# Inference
|
|
num_inference_steps: int | None = 100
|
|
|
|
# Loss computation
|
|
do_mask_loss_for_padding: bool = False
|
|
|
|
# Training presets
|
|
optimizer_lr: float = 1e-4
|
|
optimizer_betas: tuple = (0.95, 0.999)
|
|
optimizer_eps: float = 1e-8
|
|
optimizer_weight_decay: float = 1e-6
|
|
scheduler_name: str = "cosine"
|
|
scheduler_warmup_steps: int = 500
|
|
|
|
def __post_init__(self):
|
|
super().__post_init__()
|
|
|
|
"""Input validation (not exhaustive)."""
|
|
if not self.vision_backbone.startswith("resnet"):
|
|
raise ValueError(
|
|
f"`vision_backbone` must be one of the ResNet variants. Got {self.vision_backbone}."
|
|
)
|
|
|
|
if self.training_noise_sampling not in ("uniform", "beta"):
|
|
raise ValueError(
|
|
f"`training_noise_sampling` must be either 'uniform' or 'beta'. Got {self.training_noise_sampling}."
|
|
)
|
|
|
|
def get_optimizer_preset(self) -> AdamConfig:
|
|
return AdamConfig(
|
|
lr=self.optimizer_lr,
|
|
betas=self.optimizer_betas,
|
|
eps=self.optimizer_eps,
|
|
weight_decay=self.optimizer_weight_decay,
|
|
)
|
|
|
|
def get_scheduler_preset(self) -> DiffuserSchedulerConfig:
|
|
return DiffuserSchedulerConfig(
|
|
name=self.scheduler_name,
|
|
num_warmup_steps=self.scheduler_warmup_steps,
|
|
)
|
|
|
|
def validate_features(self) -> None:
|
|
if len(self.image_features) == 0 and self.env_state_feature is None:
|
|
raise ValueError("You must provide at least one image or the environment state among the inputs.")
|
|
|
|
if self.crop_shape is not None:
|
|
for key, image_ft in self.image_features.items():
|
|
if self.crop_shape[0] > image_ft.shape[1] or self.crop_shape[1] > image_ft.shape[2]:
|
|
raise ValueError(
|
|
f"`crop_shape` should fit within the images shapes. Got {self.crop_shape} "
|
|
f"for `crop_shape` and {image_ft.shape} for "
|
|
f"`{key}`."
|
|
)
|
|
|
|
# Check that all input images have the same shape.
|
|
first_image_key, first_image_ft = next(iter(self.image_features.items()))
|
|
for key, image_ft in self.image_features.items():
|
|
if image_ft.shape != first_image_ft.shape:
|
|
raise ValueError(
|
|
f"`{key}` does not match `{first_image_key}`, but we expect all image shapes to match."
|
|
)
|
|
|
|
@property
|
|
def observation_delta_indices(self) -> list:
|
|
return list(range(1 - self.n_obs_steps, 1))
|
|
|
|
@property
|
|
def action_delta_indices(self) -> list:
|
|
return list(range(1 - self.n_obs_steps, 1 - self.n_obs_steps + self.horizon))
|
|
|
|
@property
|
|
def reward_delta_indices(self) -> None:
|
|
return None
|