feat: MolmoAct2 first party implementation

application/backend/pyproject.toml

@@ -124,7 +124,7 @@ explicit = false
 
 [tool.uv.sources]
 physicalai-train = { path = "../../library", editable = true }
-physicalai = { git = "https://github.com/openvinotoolkit/physicalai.git", rev = "a3f522f9554034d2d35e97bd2b0f72feae8b3067" }
+physicalai = { git = "https://github.com/openvinotoolkit/physicalai.git", rev = "de49e1fda4e38e0a5342b0a59a4cad11a1ee8de0" }
 # Keep torchao on PyPI for all extras.
 # The XPU index publishes `torchao` local-version wheels (e.g. `+xpu`) that can
 # be selected via unsafe-best-match even for non-XPU installs (like `--extra cuda`),

library/configs/physicalai/molmoact2/pusht.yaml

@@ -0,0 +1,79 @@
+# MolmoAct2 PushT Training Configuration
+#
+# Usage:
+#   physicalai fit --config configs/physicalai/molmoact2/pusht.yaml
+model:
+  class_path: physicalai.policies.MolmoAct2
+  init_args:
+    # Initialize weights, tokenizer, and image processor from the released base checkpoint.
+    pretrained_name_or_path: allenai/MolmoAct2
+
+    # Use one observation to predict and execute a 10-step action chunk.
+    n_obs_steps: 1
+    chunk_size: 10
+    n_action_steps: 10
+
+    # Start each generated action trajectory from fresh Gaussian noise.
+    use_random_input_noise: true
+
+    # Add low-rank adapters to the VLM while fully training the action expert.
+    lora_enabled: true
+    gradient_checkpointing: true
+
+    # PushT uses the base embodiment frame and dataset-owned normalization.
+    adapt_to_so101: false
+    convert_pretrained_so101_stats: false
+    preserve_pretrained_normalization_in_training: false
+
+    # Text inserted into the model prompt to describe the embodiment and action space.
+    setup_type: "single 2D point-mass pusher robot in the PushT block-pushing simulation"
+    control_mode: "absolute planar end-effector position"
+
+    # Independent learning rates for the text backbone, vision encoder, image
+    # connector, and action expert respectively.
+    optimizer_lr: 5e-5
+    optimizer_vit_lr: 5e-5
+    optimizer_connector_lr: 5e-5
+    optimizer_action_expert_lr: 5e-5
+    optimizer_betas: [0.9, 0.95]
+    optimizer_eps: 1e-6
+    optimizer_weight_decay: 0.0
+    optimizer_grad_clip_norm: 1.0
+
+    scheduler_warmup_steps: 200
+    scheduler_decay_steps: 15000
+    scheduler_decay_lr: 1e-6
+
+data:
+  class_path: physicalai.data.lerobot.LeRobotDataModule
+  init_args:
+    # Convert LeRobot samples into PhysicalAI Observation/Action objects.
+    repo_id: lerobot/pusht
+    train_batch_size: 8
+    data_format: physicalai
+    val_gym:
+      class_path: physicalai.gyms.pusht.PushTGym
+      init_args:
+        task_description: "Push the T-shaped block onto the T-shaped target."
+    num_rollouts_val: 10
+    # Worker processes decode and prepare samples in parallel.
+    num_workers: 0
+
+trainer:
+  max_epochs: 30
+  accelerator: gpu
+  devices: 1
+  # Mixed bfloat16 reduces memory use while retaining float32-like range.
+  precision: bf16-mixed
+  log_every_n_steps: 20
+  # Run gym rollouts every 500 optimizer steps, including across epoch boundaries.
+  val_check_interval: 500
+  check_val_every_n_epoch: null
+  callbacks:
+    - class_path: physicalai.train.RolloutVideoRecorderCallback
+      init_args:
+        output_dir: val_videos_pushT
+        fps: 10
+        record_mode: all
+        frame_key: top
+        caption: "Push the T-shaped block onto the T-shaped target."

library/configs/physicalai/molmoact2/so101.yaml

@@ -0,0 +1,72 @@
+# MolmoAct2 SO101 Training Configuration
+#
+# Usage:
+#   physicalai fit --config configs/physicalai/molmoact2/so101.yaml
+model:
+  class_path: physicalai.policies.MolmoAct2
+  init_args:
+    # Initialize from the released SO100/SO101 checkpoint. norm_tag is deliberately
+    # omitted so setup("fit") adopts this dataset's feature contract and quantiles.
+    pretrained_name_or_path: allenai/MolmoAct2-SO100_101
+
+    # Use one observation to predict and execute the full 30-step action chunk.
+    n_obs_steps: 1
+    chunk_size: 30
+    n_action_steps: 30
+
+    # Start each generated action trajectory from fresh Gaussian noise.
+    use_random_input_noise: true
+
+    # Add low-rank adapters to the VLM while fully training the action expert.
+    lora_enabled: true
+
+    # Recompute intermediate activations during backward passes to save GPU memory.
+    gradient_checkpointing: true
+
+    # Convert dataset samples and quantiles into the checkpoint's legacy SO101
+    # coordinate frame, without importing or retaining checkpoint normalization stats.
+    adapt_to_so101: true
+    convert_pretrained_so101_stats: false
+    preserve_pretrained_normalization_in_training: false
+
+    # Text inserted into the model prompt to describe the embodiment and action space.
+    setup_type: "single so100/so101 robotic arm in molmoact2"
+    control_mode: "absolute joint pose"
+
+    # Use the shared optimizer recipe validated on PushT.
+    optimizer_lr: 5e-5
+    optimizer_vit_lr: 5e-5
+    optimizer_connector_lr: 5e-5
+    optimizer_action_expert_lr: 5e-5
+    optimizer_betas: [0.9, 0.95]
+    optimizer_eps: 1e-6
+    optimizer_weight_decay: 0.0
+    optimizer_grad_clip_norm: 1.0
+    scheduler_warmup_steps: 200
+    scheduler_decay_steps: 15000
+    scheduler_decay_lr: 1e-6
+
+data:
+  class_path: physicalai.data.lerobot.LeRobotDataModule
+  init_args:
+    # Convert the SO101 LeRobot dataset into PhysicalAI Observation/Action objects.
+    repo_id: Daankrol/pick-and-place-multi-obj
+    train_batch_size: 16
+    val_batch_size: 16
+    # Hold out 10% of episodes reproducibly for validation-loss evaluation.
+    val_split: 0.1
+    val_split_seed: 42
+    data_format: physicalai
+    # Worker processes decode and prepare samples in parallel.
+    num_workers: 8
+
+trainer:
+  max_epochs: 8
+  accelerator: gpu
+  devices: 1
+  # Mixed bfloat16 reduces memory use while retaining float32-like range.
+  precision: bf16-mixed
+  log_every_n_steps: 20
+  # Validate every 500 optimizer steps, including across epoch boundaries.
+  val_check_interval: 500
+  check_val_every_n_epoch: null

library/pyproject.toml

@@ -104,6 +104,14 @@ smolvla = [
     "safetensors>=0.4.3,<1.0.0",
 ]
 
+# MolmoAct2
+# Supports LoRA fine-tuning for efficient adaptation
+molmoact2 = [
+    "transformers>=5.15.1,<5.16.0",
+    "peft>=0.13.0,<1.0.0",
+    "safetensors>=0.4.3,<1.0.0",
+]
+
 # RLWRLD RLDX-1
 rldx1 = [
     "transformers>=5.15.1,<5.16.0",
@@ -129,6 +137,7 @@ policies = [
     "physicalai-train[pi05]",
     "physicalai-train[groot]",
     "physicalai-train[smolvla]",
+    "physicalai-train[molmoact2]",
     "physicalai-train[rldx1]",
 ]
 
@@ -168,6 +177,7 @@ dev = [
     "pyrefly>=0.18.0",
     "ruff>=0.14.14",
     "diffusers>=0.38.0",
+    "peft>=0.13.0,<1.0.0",
 ]
 
 
@@ -219,7 +229,7 @@ explicit = false
 
 # PyTorch sources configuration
 [tool.uv.sources]
-physicalai = { git = "https://github.com/openvinotoolkit/physicalai.git", rev = "a3f522f9554034d2d35e97bd2b0f72feae8b3067" }
+physicalai = { git = "https://github.com/openvinotoolkit/physicalai.git", rev = "de49e1fda4e38e0a5342b0a59a4cad11a1ee8de0" }
 torch = [
     { index = "pytorch-cpu", extra = "cpu" },
     { index = "pytorch-cu128", extra = "cu128" },

library/src/physicalai/data/__init__.py

@@ -6,7 +6,7 @@
 from .datamodules import DataModule
 from .dataset import Dataset
 from .lerobot import LeRobotDataModule
-from .observation import Feature, FeatureType, NormalizationParameters, Observation
+from .observation import Feature, FeatureType, NormalizationParameters, NormalizationValue, Observation
 
 __all__ = [
     "DataModule",
@@ -15,5 +15,6 @@
     "FeatureType",
     "LeRobotDataModule",
     "NormalizationParameters",
+    "NormalizationValue",
     "Observation",
 ]

library/src/physicalai/data/observation.py

@@ -12,6 +12,8 @@
 import numpy as np
 import torch
 
+type NormalizationValue = float | list[float] | list[list[float]] | list[list[list[float]]] | None
+
 
 @dataclass
 class Observation:
@@ -487,14 +489,27 @@ class Feature:
 
 @dataclass(frozen=True)
 class NormalizationParameters:
-    """Parameters for normalizing a tensor."""
-
-    mean: list[float] | float | None = None
-    std: list[float] | float | None = None
-    min: list[float] | float | None = None
-    max: list[float] | float | None = None
-    q01: list[float] | float | None = None
-    q99: list[float] | float | None = None
+    """Statistics and optional selection mask for feature normalization.
+
+    Each statistic may be a scalar or a sequence matching the feature's
+    normalized dimensions.
+    """
+
+    #: Arithmetic mean used by mean/std normalization.
+    mean: NormalizationValue = None
+    #: Standard deviation used by mean/std normalization.
+    std: NormalizationValue = None
+    #: Lower bound used by min/max normalization.
+    min: NormalizationValue = None
+    #: Upper bound used by min/max normalization.
+    max: NormalizationValue = None
+    #: First-percentile bound used by quantile normalization.
+    q01: NormalizationValue = None
+    #: Ninety-ninth-percentile bound used by quantile normalization.
+    q99: NormalizationValue = None
+    #: Optional per-dimension selector: ``True`` dimensions are normalized and
+    #: ``False`` dimensions pass through unchanged for all statistical modes.
+    mask: list[bool] | None = None
 
 
 # Module-level constants for convenient dict access

library/src/physicalai/eval/rollout/functional.py

@@ -330,7 +330,8 @@ def run_rollout_loop(  # noqa: PLR0914
 
         # Record frame for video (happens DURING the loop, not after)
         if video_recorder is not None:
-            frame = _collect_frame(observation, frame_key)
+            video_frame_key = video_recorder.frame_key if video_recorder.frame_key is not None else frame_key
+            frame = _collect_frame(observation, video_frame_key)
             if frame is not None:
                 video_recorder.record_frame(frame)
 
@@ -340,7 +341,11 @@ def run_rollout_loop(  # noqa: PLR0914
             # InferenceModel doesn't have eval() but doesn't need it
             if hasattr(policy, "eval") and callable(policy.eval):
                 policy.eval()
-            action = policy.select_action(observation)  # shape: (B, action_dim)
+            # Move the observation onto the policy device (no-op for exported
+            # InferenceModel policies, which expose no `device`).
+            policy_device = getattr(policy, "device", None)
+            policy_observation = observation.to(policy_device) if policy_device is not None else observation
+            action = policy.select_action(policy_observation)  # shape: (B, action_dim)
 
         # For non-vectorized envs (batch_size=1), squeeze the batch dimension
         # LiberoGym and similar envs expect action shape (action_dim,) not (1, action_dim)

library/src/physicalai/eval/video.py

@@ -31,6 +31,8 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, Self
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     import numpy as np
 
 logger = logging.getLogger(__name__)
@@ -143,6 +145,7 @@ def __init__(
         codec: str = "h264",
         record_mode: RecordMode = "all",
         caption: str | None = None,
+        frame_key: str | Sequence[str] | None = None,
     ) -> None:
         """Initialize video recorder.
 
@@ -152,6 +155,8 @@ def __init__(
             codec: Video codec for encoding.
             record_mode: When to save videos.
             caption: Optional text burned into a bar at the bottom of every frame.
+            frame_key: Observation image key or keys to record. Uses the rollout's
+                frame key when omitted.
         """
         _check_imageio_available()
 
@@ -161,6 +166,7 @@ def __init__(
         self.codec = codec
         self.record_mode = record_mode
         self.caption = caption
+        self.frame_key = frame_key
 
         self._frames: list[np.ndarray] = []
         self._current_episode_name: str | None = None

library/src/physicalai/gyms/libero.py

@@ -398,11 +398,10 @@ def step(
         Raises:
             ValueError: If action has wrong dimensions.
         """
-        # Convert tensor to numpy if needed
+        # Convert tensor to numpy if needed.
+        # NumPy does not support torch.bfloat16 directly, so cast to float32 first.
         if isinstance(action, torch.Tensor):
-            if action.dtype == torch.bfloat16:
-                action = action.float()
-            action = action.cpu().numpy()
+            action = action.detach().to(dtype=torch.float32).cpu().numpy()
 
         # Validate action shape
         if action.ndim != 1: