feat(train): align native SmolVLA recipe

Introduce SmolVLA-style VLM prefix encoder backbone and IMF-AttnRes
conditioned agent. Add episode-level train/val split, action MSE
validation in training, and headless eval support.

docs/superpowers/plans/2026-04-23-sim-air-insert-ring-bar.md

@@ -0,0 +1,311 @@
+# sim_air_insert_ring_bar Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Add an independent dual-Diana MuJoCo task `sim_air_insert_ring_bar` with a square ring block, a square bar block, staged rewards, strict finite-geometry in-air insertion success detection, and a task-specific scripted policy.
+
+**Architecture:** Reuse the current dual-Diana EE-control stack and environment factory, but add a task-specific scene XML, robot asset entrypoint, sampling helpers, and a new task-specific environment module. Keep `sim_transfer` untouched while introducing pure-Python geometry helpers and focused tests so reward/success behavior can be regression tested without requiring a full MuJoCo rollout in every test.
+
+**Tech Stack:** Python, unittest, MuJoCo XML assets, existing dual-Diana environment classes, Hydra-compatible task naming/config patterns.
+
+---
+
+## File Structure / Responsibilities
+
+- **Create:** `roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml`
+  - Defines the rigid ring body and bar body, each with a free joint and stable box-based geoms.
+- **Create:** `roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml`
+  - Scene entrypoint that includes the shared world/table/robot assets plus the new object XML.
+- **Modify:** `roboimi/assets/robots/diana_med.py`
+  - Add a task-specific robot asset class for the new scene XML without changing existing `BiDianaMed` behavior.
+- **Modify:** `roboimi/utils/act_ex_utils.py`
+  - Add deterministic helpers to sample left/right planar placement regions for ring and bar objects.
+- **Modify:** `roboimi/utils/constants.py`
+  - Register the new task name and default metadata.
+- **Create:** `roboimi/envs/double_air_insert_env.py`
+  - New task-specific environment, finite-geometry success helpers, reset logic, reward logic, and task factory branch.
+- **Modify:** `roboimi/envs/double_pos_ctrl_env.py`
+  - Route `make_sim_env()` to the new task-specific environment while keeping current `sim_transfer` logic unchanged.
+- **Create:** `roboimi/demos/diana_air_insert_policy.py`
+  - Task-specific waypoint/open-loop scripted policy for grasp-lift-align-insert.
+- **Modify:** `roboimi/demos/vla_scripts/eval_vla.py`
+  - Reset the new task with the correct sampled task state instead of assuming a single transfer box pose.
+- **Create:** `tests/test_air_insert_env.py`
+  - Focused unit tests for sampling, reset helpers, reward progression, and strict success detection.
+- **Modify:** `tests/test_eval_vla_headless.py`
+  - Add coverage that headless evaluation dispatches the correct reset sampler for the new task.
+- **Modify:** `tests/test_robot_asset_paths.py`
+  - Verify the new robot asset class resolves its XML path correctly independent of cwd.
+
+---
+
+### Task 1: Add failing tests for task registration, samplers, and asset wiring
+
+**Files:**
+- Create: `tests/test_air_insert_env.py`
+- Modify: `tests/test_eval_vla_headless.py`
+- Modify: `tests/test_robot_asset_paths.py`
+- Modify: `roboimi/utils/act_ex_utils.py` (later in implementation)
+- Modify: `roboimi/utils/constants.py` (later in implementation)
+- Modify: `roboimi/assets/robots/diana_med.py` (later in implementation)
+- Modify: `roboimi/envs/double_pos_ctrl_env.py` (later in implementation)
+- Create: `roboimi/envs/double_air_insert_env.py` (minimal stub in this task)
+
+- [ ] **Step 1: Write failing tests for task config and sampling helpers**
+
+Add tests in `tests/test_air_insert_env.py` covering:
+- `SIM_TASK_CONFIGS['sim_air_insert_ring_bar']` exists
+- `sample_air_insert_ring_bar_pose()` (or equivalent helper) returns ring/bar positions with fixed z and correct left/right planar ranges
+- output structure is explicit and easy for reset/eval code to consume
+
+- [ ] **Step 2: Write failing tests for environment factory dispatch and robot asset resolution**
+
+Add tests covering:
+- `make_sim_env('sim_air_insert_ring_bar', headless=True)` dispatches to the new environment with rendering disabled
+- a new robot asset class resolves the new XML path independent of cwd, similar to the existing `BiDianaMed` test pattern
+
+- [ ] **Step 3: Write failing tests for eval reset helper dispatch**
+
+Extend `tests/test_eval_vla_headless.py` so headless eval can reset the new task using the new sampler instead of hard-coding `sample_transfer_pose()`.
+
+- [ ] **Step 4: Run the targeted tests to verify they fail for the expected missing-feature reasons**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v`
+
+Expected:
+- FAIL because the new task config/helper/class/dispatch branch does not exist yet
+
+- [ ] **Step 5: Implement the minimal production code to satisfy the new task registration and helper tests**
+
+Implement only enough to make the new tests pass:
+- add new task config entry
+- add the new placement sampler
+- add the new robot asset class
+- create a minimal importable `double_air_insert_env.py` stub and class/function surface needed for factory dispatch tests
+- add the factory dispatch branch / headless wiring
+- update eval reset dispatch for the new task
+
+- [ ] **Step 6: Re-run the targeted tests to verify they pass**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v`
+
+Expected:
+- PASS for the new registration/sampler/dispatch/asset tests
+
+- [ ] **Step 7: Commit Task 1**
+
+Run:
+`git add tests/test_air_insert_env.py tests/test_eval_vla_headless.py tests/test_robot_asset_paths.py roboimi/utils/act_ex_utils.py roboimi/utils/constants.py roboimi/assets/robots/diana_med.py roboimi/envs/double_pos_ctrl_env.py roboimi/envs/double_air_insert_env.py roboimi/demos/vla_scripts/eval_vla.py && git commit -m "feat(env): register sim air insert ring bar task"`
+
+---
+
+### Task 2: Add the MuJoCo ring+bar scene assets and reset helpers
+
+**Files:**
+- Create: `roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml`
+- Create: `roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml`
+- Create or Modify: `roboimi/envs/double_air_insert_env.py`
+- Modify: `tests/test_air_insert_env.py`
+
+- [ ] **Step 1: Write failing tests for object reset helpers and scene-specific joint naming assumptions**
+
+In `tests/test_air_insert_env.py`, add unit tests for helper functions that:
+- write ring pose to `ring_block_joint` from the named task-state mapping
+- write bar pose to `bar_block_joint` from the named task-state mapping
+- read back `env_state` as a stable 14D vector `[ring_pos, ring_quat, bar_pos, bar_quat]`
+
+Use fake `mj_data` objects so tests stay fast and deterministic.
+
+- [ ] **Step 2: Run the focused test slice and verify it fails**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- FAIL because reset/state helper functions and joint conventions are not implemented yet
+
+- [ ] **Step 3: Implement the scene XML files and reset/state helper code**
+
+Implement:
+- the object XML with one rigid ring body and one rigid bar body
+- the task scene XML entrypoint using the shared world/table/robot includes
+- reset helper(s) in `double_air_insert_env.py` that set qpos for both free joints with fixed quaternions
+- task-state accessor(s) returning both object poses in a stable structure
+
+- [ ] **Step 4: Re-run the focused test slice and verify it passes**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- PASS for reset/state helper tests
+
+- [ ] **Step 5: Commit Task 2**
+
+Run:
+`git add roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml roboimi/envs/double_air_insert_env.py tests/test_air_insert_env.py && git commit -m "feat(scene): add ring and bar insertion scene assets"`
+
+---
+
+### Task 3: Implement strict reward and finite-geometry success detection
+
+**Files:**
+- Modify: `roboimi/envs/double_air_insert_env.py`
+- Modify: `tests/test_air_insert_env.py`
+
+- [ ] **Step 1: Write failing tests for reward stages and strict success detection**
+
+Add tests in `tests/test_air_insert_env.py` for:
+- left contact stage reward
+- right contact stage reward
+- ring lifted off table stage
+- bar lifted off table stage
+- positive success case where a finite bar truly passes through the aperture
+- negative case where the centerline would pass but the finite square body would clip
+- negative case where the bar has not crossed the ring thickness direction enough
+- negative case where one/both objects are still on the table
+
+Structure the tests around pure helper functions and light fake contact/state objects so the geometry logic is directly regression tested.
+
+- [ ] **Step 2: Run the focused tests and verify they fail for missing reward/success logic**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- FAIL because the staged reward and finite-geometry insertion logic are not implemented yet
+
+- [ ] **Step 3: Implement minimal strict success helpers and reward logic**
+
+Implement in `roboimi/envs/double_air_insert_env.py`:
+- pure helper(s) for transforming bar geometry into ring-local coordinates
+- finite-geometry insertion predicate (not centerline-only)
+- table-contact / airborne checks
+- staged reward function returning the highest achieved stage with `max_reward = 5`
+
+- [ ] **Step 4: Re-run the focused tests to verify the logic passes**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- PASS for reward and success-detection regression tests
+
+- [ ] **Step 5: Commit Task 3**
+
+Run:
+`git add roboimi/envs/double_air_insert_env.py tests/test_air_insert_env.py && git commit -m "feat(env): add strict air insertion reward and success logic"`
+
+---
+
+### Task 4: Add the scripted policy and integration smoke coverage
+
+**Files:**
+- Create: `roboimi/demos/diana_air_insert_policy.py`
+- Modify: `roboimi/demos/diana_record_sim_episodes.py`
+- Modify: `tests/test_air_insert_env.py`
+- Optionally Modify: `roboimi/demos/vla_scripts/eval_vla.py` (only if integration gaps remain after Task 1)
+
+- [ ] **Step 1: Write failing tests for scripted-policy action shape and basic generation**
+
+Add tests covering:
+- the new policy produces a 16D action
+- trajectory generation accepts sampled named task state without error
+- the first action is a valid open-gripper safe pose command
+- a deterministic nominal smoke path (with canonical sampled state or fake env shim) reaches the intended terminal interface contract without shape/reward mismatches
+
+Keep the tests unit-level; do not require a full MuJoCo rollout for every assertion.
+
+- [ ] **Step 2: Write failing tests for the scripted rollout entrypoint and a real headless smoke path**
+
+Add coverage for both:
+- the standard scripted rollout entrypoint (`roboimi/demos/diana_record_sim_episodes.py`) can select the new task sampler/policy instead of remaining sim_transfer-only
+- a deterministic integration/smoke test that instantiates `make_sim_env('sim_air_insert_ring_bar', headless=True)`, resets with sampled named task state, and steps a few actions or scripted-policy outputs using the real task XML and task-specific wiring
+
+- [ ] **Step 3: Run the scripted-policy tests and verify they fail**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- FAIL because the new scripted policy does not exist yet
+
+- [ ] **Step 4: Implement the waypoint-based scripted policy**
+
+Implement a conservative open-loop policy with phases:
+- safe wait pose
+- above-target approach
+- descend + grasp
+- dual lift
+- airborne meeting alignment
+- bar push-through insertion
+
+Use fixed orientations for version 1 and follow the existing repository style from `diana_policy.py`.
+
+- [ ] **Step 5: Re-run the scripted-policy tests to verify they pass**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- PASS for scripted-policy tests
+
+- [ ] **Step 6: Run the combined verification suite for this feature**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_eval_vla_rollout_artifacts tests.test_train_vla_rollout_validation tests.test_robot_asset_paths -v`
+
+Expected:
+- PASS with 0 failures
+
+- [ ] **Step 6b: Run the mandatory real headless smoke check**
+
+Run a focused smoke command that instantiates the real task, resets with sampled state, and steps a few actions using the new scripted policy or a deterministic action sequence.
+
+Example command (adjust module/test helper if needed):
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env.AirInsertEnvSmokeTest -v`
+
+Expected:
+- PASS, proving the real XML/assets/env wiring instantiate and step correctly in headless mode
+
+- [ ] **Step 7: Commit Task 4**
+
+Run:
+`git add roboimi/demos/diana_air_insert_policy.py tests/test_air_insert_env.py tests/test_eval_vla_headless.py tests/test_robot_asset_paths.py roboimi/demos/vla_scripts/eval_vla.py && git commit -m "feat(policy): add scripted air insertion policy"`
+
+---
+
+### Task 5: Final verification and implementation review
+
+**Files:**
+- Review all files touched above
+
+- [ ] **Step 1: Run fresh end-to-end verification before claiming completion**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v`
+
+Expected:
+- PASS with 0 failures
+
+- [ ] **Step 2: Inspect git status and recent commits**
+
+Run:
+`git status --short && git log --oneline --decorate -n 8`
+
+Expected:
+- only intended feature files modified / committed
+
+- [ ] **Step 3: Request final code review for the completed feature**
+
+Use the requesting-code-review skill against the full diff from the feature branch starting point to current HEAD.
+
+- [ ] **Step 4: Address any review findings and re-run verification if code changes**
+
+If fixes are made, repeat the unittest command from Step 1.
+
+- [ ] **Step 5: Hand off using finishing-a-development-branch**
+
+After verification and review, use the finishing-a-development-branch skill to decide merge / PR / cleanup.

docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md

@@ -0,0 +1,184 @@
+# Native SmolVLA Model Migration Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Migrate only the SmolVLA model core into RoboIMI and expose it as a native VLA agent without importing the full LeRobot package.
+
+**Architecture:** Add a focused `roboimi.vla.models.smolvla` model package and a `SmolVLANativeAgent` wrapper that speaks the existing RoboIMI train/eval interface. Keep normalization, queues, camera ordering, and task fallback in the agent; keep model math in the migrated model package.
+
+**Tech Stack:** PyTorch, Transformers, Hydra/OmegaConf, unittest/pytest-style tests, existing RoboIMI normalization and VLA scripts.
+
+---
+
+## File Structure
+
+- Create `roboimi/vla/models/smolvla/__init__.py`: package exports.
+- Create `roboimi/vla/models/smolvla/configuration.py`: lightweight config dataclass.
+- Create `roboimi/vla/models/smolvla/modeling.py`: model utilities and `VLAFlowMatching`.
+- Create `roboimi/vla/models/smolvla/smolvlm_with_expert.py`: adapted VLM/expert core.
+- Create `roboimi/vla/agent_smolvla_native.py`: RoboIMI agent wrapper.
+- Create `roboimi/vla/conf/agent/smolvla_native.yaml`: Hydra config.
+- Create `tests/test_smolvla_native_agent.py`: TDD tests for wrapper behavior.
+- Create `tests/test_smolvla_native_modeling.py`: TDD tests for utilities/config where useful.
+
+## Task 1: Wrapper behavior tests and minimal agent skeleton
+
+**Files:**
+- Create: `tests/test_smolvla_native_agent.py`
+- Create: `roboimi/vla/agent_smolvla_native.py`
+
+- [ ] **Step 1: Write failing tests for task fallback, camera order, queue, and fake model calls**
+
+Implement tests using fake tokenizer and fake model. Test names:
+
+- `test_compute_loss_orders_cameras_tokenizes_task_and_calls_native_model`
+- `test_unknown_task_uses_configured_task_description`
+- `test_missing_camera_raises_clear_error`
+- `test_predict_action_chunk_denormalizes_fake_model_output`
+- `test_select_action_uses_action_queue_before_recomputing`
+
+- [ ] **Step 2: Run tests and verify import failure**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: FAIL because `roboimi.vla.agent_smolvla_native` does not exist.
+
+- [ ] **Step 3: Implement minimal `SmolVLANativeAgent` skeleton**
+
+Implement constructor injection for `model` and `tokenizer`, plus `_resolve_task`, `_order_images`, `_tokenize_tasks`, `reset`, `_prepare_observation_batch`, `compute_loss`, `predict_action_chunk`, and `select_action`. Use fake model in tests; real model construction can raise a clear ImportError until Task 3.
+
+- [ ] **Step 4: Run tests and verify pass**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: PASS.
+
+## Task 2: Config and modeling utility tests
+
+**Files:**
+- Create: `tests/test_smolvla_native_modeling.py`
+- Create: `roboimi/vla/models/smolvla/configuration.py`
+- Create: `roboimi/vla/models/smolvla/modeling.py`
+- Create: `roboimi/vla/models/smolvla/__init__.py`
+
+- [ ] **Step 1: Write failing tests**
+
+Cover:
+
+- `NativeSmolVLAConfig` validates `n_action_steps <= chunk_size`.
+- `pad_vector` pads and rejects truncation.
+- `resize_with_pad` preserves batch/channel shape and output size.
+- `make_att_2d_masks` matches prefix-LM mask semantics.
+
+- [ ] **Step 2: Run tests and verify failure**
+
+Run: `python -m unittest tests.test_smolvla_native_modeling -v`
+
+Expected: FAIL because package/modeling code is missing.
+
+- [ ] **Step 3: Implement config and utility functions**
+
+Port minimal functions from external SmolVLA while removing LeRobot imports.
+
+- [ ] **Step 4: Run tests and verify pass**
+
+Run: `python -m unittest tests.test_smolvla_native_modeling -v`
+
+Expected: PASS.
+
+## Task 3: Migrate model core
+
+**Files:**
+- Modify: `roboimi/vla/models/smolvla/modeling.py`
+- Create: `roboimi/vla/models/smolvla/smolvlm_with_expert.py`
+- Modify: `roboimi/vla/agent_smolvla_native.py`
+
+- [ ] **Step 1: Write fake-backed integration tests for lazy real model construction**
+
+Patch `VLAFlowMatching` and tokenizer loader so `SmolVLANativeAgent(model=None, tokenizer=None)` constructs a fake model without real downloads. Verify config fields are passed.
+
+- [ ] **Step 2: Run test and verify failure**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: FAIL because real construction path is not implemented.
+
+- [ ] **Step 3: Port `SmolVLMWithExpertModel` and `VLAFlowMatching`**
+
+Copy only model-core logic from external files, replacing LeRobot dependencies with local config and local helpers. Preserve `forward`, `sample_actions`, `denoise_step`, image resize/pad, state/action padding, and language token inputs.
+
+- [ ] **Step 4: Wire agent lazy construction**
+
+If `model` is not supplied, create `NativeSmolVLAConfig`, tokenizer, and `VLAFlowMatching`. Keep `model` injection path for tests.
+
+- [ ] **Step 5: Run unit tests**
+
+Run:
+
+```bash
+python -m unittest tests.test_smolvla_native_agent tests.test_smolvla_native_modeling -v
+```
+
+Expected: PASS.
+
+## Task 4: Hydra config
+
+**Files:**
+- Create: `roboimi/vla/conf/agent/smolvla_native.yaml`
+- Modify: `tests/test_smolvla_native_agent.py`
+
+- [ ] **Step 1: Add failing Hydra compose test**
+
+Test composing `agent=smolvla_native` exposes `_target_`, `action_dim=16`, `obs_dim=16`, `camera_names=${data.camera_names}`, and model config fields.
+
+- [ ] **Step 2: Run and verify failure**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: FAIL because yaml is missing.
+
+- [ ] **Step 3: Add yaml config**
+
+Add `smolvla_native.yaml` with `_target_: roboimi.vla.agent_smolvla_native.SmolVLANativeAgent` and sane defaults.
+
+- [ ] **Step 4: Run test and verify pass**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: PASS.
+
+## Task 5: Regression checks
+
+**Files:**
+- No production changes expected unless tests reveal integration gaps.
+
+- [ ] **Step 1: Run focused existing tests**
+
+Run:
+
+```bash
+python -m unittest tests.test_smolvla_prefix_encoder tests.test_smolvla_imf_agent tests.test_eval_vla_execution -v
+```
+
+Expected: PASS.
+
+- [ ] **Step 2: Search for forbidden import**
+
+Run:
+
+```bash
+rg -n "import lerobot|from lerobot" roboimi/vla/models/smolvla roboimi/vla/agent_smolvla_native.py
+```
+
+Expected: no matches.
+
+- [ ] **Step 3: Commit**
+
+Run:
+
+```bash
+git add roboimi/vla/models/smolvla roboimi/vla/agent_smolvla_native.py roboimi/vla/conf/agent/smolvla_native.yaml tests/test_smolvla_native_agent.py tests/test_smolvla_native_modeling.py docs/superpowers/specs/2026-05-25-native-smolvla-model-migration-design.md docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md
+git commit -m "feat(vla): add native SmolVLA model agent"
+```
+
+Expected: commit succeeds.

docs/superpowers/specs/2026-04-23-sim-air-insert-ring-bar-design.md

@@ -0,0 +1,316 @@
+# sim_air_insert_ring_bar Design
+
+## Summary
+
+Add a new independent MuJoCo simulation task named `sim_air_insert_ring_bar` that keeps the existing dual-Diana tabletop setup but replaces the single transfer box with two randomized objects:
+
+- a square ring block grasped by the left arm
+- a square bar block grasped by the right arm
+
+The task is to pick both objects off the table and complete an in-air insertion where the bar truly passes through the ring aperture. The existing `sim_transfer` task must remain unchanged.
+
+## Goals
+
+- Reuse the current dual-Diana EE-control simulation stack
+- Keep the same table/base robot arrangement as the existing transfer task
+- Add an independent task entrypoint and scene definition
+- Randomize planar placement of both objects within left/right task-specific regions
+- Implement reward staging for contact, lift, and successful in-air insertion
+- Add a scripted policy that performs pick, lift, align, and in-air insertion
+- Preserve compatibility with existing environment creation, evaluation, and rollout patterns
+
+## Non-Goals
+
+- No random yaw in the first version
+- No visual servoing or closed-loop insertion controller
+- No general multi-task environment framework refactor
+- No guarantee that the VLA training stack is immediately tuned for this new task
+- No replacement or behavior change for `sim_transfer`
+
+## Task Name
+
+Use a new task name:
+
+- `sim_air_insert_ring_bar`
+
+This task should be exposed alongside `sim_transfer`, not as a replacement.
+
+## Scene Geometry
+
+### Shared Base Scene
+
+Keep the dual Diana robot, the table, and the existing camera layout conceptually unchanged.
+
+### Ring Block
+
+Represent the square ring as a rigid free body composed from simple MuJoCo box geoms rather than an external mesh.
+
+Dimensions:
+
+- outer side length: 68 mm
+- inner aperture side length: 32 mm
+- thickness: 18 mm
+- ring wall width: 18 mm
+
+The ring should behave as a single object body with a single free joint.
+
+### Bar Block
+
+Represent the bar as a rigid free body with a single box geom.
+
+Dimensions:
+
+- length: 90 mm
+- cross-section: 18 mm x 18 mm
+
+The bar should also be a single free-joint body.
+
+## Initial Placement / Reset
+
+The first version uses position-only randomization with fixed orientation. Reset sampling stays **caller-driven**, matching the existing `sim_transfer` usage pattern in rollout/eval code: a helper samples task state, then callers pass that state into `env.reset(...)`.
+
+Use an explicit sampled task-state structure with named fields:
+
+- `ring_pos`: 3D position
+- `ring_quat`: fixed 4D quaternion for version 1
+- `bar_pos`: 3D position
+- `bar_quat`: fixed 4D quaternion for version 1
+
+Behavior:
+
+- ring block: randomized only in a left-side planar sampling region
+- bar block: randomized only in a right-side planar sampling region
+- both objects start flat on the table
+- both objects use fixed orientation at reset
+- no random yaw, tilt, or flip in this version
+
+The sampling regions should be chosen conservatively so that:
+
+- the left arm can comfortably reach and grasp the ring
+- the right arm can comfortably reach and grasp the bar
+- scripted open-loop pick trajectories remain feasible
+
+## Control / Action Interface
+
+Reuse the current 16D EE-space action convention already used by the dual-Diana position-control environment:
+
+- left arm EE pose: 7D (`xyz + quat`)
+- right arm EE pose: 7D (`xyz + quat`)
+- left gripper command: 1D
+- right gripper command: 1D
+
+The new task should continue using EE targets transformed through the existing IK-based control path.
+
+## Environment Structure
+
+Implement this as a new task-specific environment path while reusing the existing dual-Diana simulation base where possible.
+
+Expected responsibilities:
+
+- scene instantiation for the ring+bar setup
+- task reset for randomized object placement
+- environment-state accessors for both objects
+- reward computation
+- in-air insertion success detection
+
+The environment factory must dispatch by task name and leave the `sim_transfer` branch unchanged.
+
+## Observation / Environment State
+
+The task should retain the current observation structure style used by the dual-Diana environment:
+
+- `qpos`
+- multi-camera images
+
+For task state access, the environment should expose a stable `env_state` vector with this exact order:
+
+- `ring_pos[0:3]`
+- `ring_quat[3:7]`
+- `bar_pos[7:10]`
+- `bar_quat[10:14]`
+
+This 14D state should be sufficient for scripted-policy debugging and future rollout analysis, while reset itself remains caller-driven via the named task-state helper structure above.
+
+## Reward Design
+
+Use staged rewards in the same spirit as the current task, returning the highest achieved stage rather than accumulating one-time sparse bonuses per event.
+
+Maximum reward:
+
+- `max_reward = 5`
+
+Reward stages:
+
+1. left gripper touches the ring block
+2. right gripper touches the bar block
+3. ring block is lifted off the table
+4. bar block is lifted off the table
+5. while both objects are off the table, the bar truly passes through the ring aperture
+
+Notes:
+
+- contact rewards are intended as grasp-progress stages
+- lift rewards require the object to be off the table, not merely touched
+- final success reward only applies when both objects are airborne
+
+## Success Detection
+
+Success must **not** be based on a centerline-only check.
+
+A centerline-only test is insufficient because:
+
+- the bar has thickness, so a centerline can pass through while the body cannot
+- a square bar with imperfect orientation can have its centerline inside the aperture while its corners still collide with the ring
+
+### Required Success Semantics
+
+A successful insertion requires all of the following:
+
+1. the ring is off the table
+2. the bar is off the table
+3. the bar has actually crossed through the ring thickness direction
+4. the bar’s finite square cross-section fits through the square aperture during that crossing
+
+### Recommended Detection Approach
+
+Use a task-level geometric check in Python rather than relying on contact alone.
+
+Implementation intent:
+
+- transform the bar geometry into the ring’s local frame
+- reason about the bar as a finite oriented box (not a line)
+- verify that the bar has crossed the ring thickness direction
+- verify that the portion of the bar passing the aperture fits within the inner square opening, accounting for the bar’s cross-section and orientation
+
+This geometric check is the primary success test.
+
+### Role of Contacts
+
+Contacts may still be used for:
+
+- grasp-stage rewards
+- debugging / diagnostics
+
+But contact alone should **not** be the sole criterion for insertion success, since:
+
+- a true clean insertion may have limited aperture-wall contact
+- persistent contact can also happen while the bar is jammed and not actually inserted
+
+## Scripted Policy
+
+Add a new task-specific scripted policy for `sim_air_insert_ring_bar`.
+
+### Policy Intent
+
+The first version prioritizes a conservative, reliable open-loop demonstration rather than an optimized trajectory.
+
+### Action Phases
+
+The scripted policy should follow these phases:
+
+1. move both arms to safe initial / waiting poses with grippers open
+2. move left arm above the ring and right arm above the bar
+3. descend and grasp the assigned objects
+4. lift both objects clear of the table
+5. move both objects to an airborne meeting region above the table
+6. hold the ring stably while aligning the bar with the aperture
+7. push the bar along the intended insertion direction until the geometric success condition is met
+
+### Grasp Assignment
+
+- left arm: ring only
+- right arm: bar only
+
+### Motion Style
+
+Keep the current repository style:
+
+- waypoint-based trajectory definition
+- open-loop interpolation between waypoints
+- fixed grasp orientation in the first version
+
+No adaptive replanning is required for the first version.
+
+## Files / Integration Scope
+
+The implementation is expected to add task-specific files rather than broadly refactoring the codebase.
+
+Likely additions / changes:
+
+- a new MuJoCo scene XML for the ring+bar task
+- one or more XML fragments defining the two new objects
+- a new task-specific dual-Diana environment file
+- robot asset wiring for the new scene XML
+- reset sampling helpers for the new task
+- task registration in constants / environment factory paths
+- a new scripted policy file
+- focused tests for task creation, reset, rewards, success detection, and scripted policy shape/smoke behavior
+
+## Testing Requirements
+
+At minimum, add regression coverage for:
+
+### Environment Creation
+
+- the new task can be created via the task factory
+- the existing `sim_transfer` task remains unchanged
+
+### Reset / Sampling
+
+- ring reset positions are inside the left sampling region
+- bar reset positions are inside the right sampling region
+- reset orientation is fixed as intended
+
+### Environment State
+
+- environment-state access returns both object poses in the expected structure
+
+### Success Detection
+
+Must include both positive and negative cases.
+
+Positive case:
+
+- a configuration where the finite bar truly passes through the ring aperture is detected as success
+
+Negative cases:
+
+- centerline-inside but finite body would clip the aperture
+- not enough depth / not actually crossing the ring thickness direction
+- one or both objects still on the table
+
+### Reward Logic
+
+- left contact stage
+- right contact stage
+- ring lift stage
+- bar lift stage
+- final success stage with `max_reward = 5`
+
+### Scripted Policy
+
+At minimum:
+
+- policy emits valid 16D actions
+- trajectory generation does not error
+- rollout smoke path can step through the new environment
+
+## Risks / Constraints
+
+- MuJoCo contact naming must remain stable enough for stage rewards
+- geometric insertion checks must be strict enough to avoid false positives but not so brittle that numerically valid insertions are missed
+- scripted open-loop insertion may require conservative alignment and lift heights to keep the first version reliable
+
+## Acceptance Criteria
+
+The feature is complete when all of the following are true:
+
+- `sim_air_insert_ring_bar` is creatable as an independent task
+- the scene contains the dual Diana, table, ring block, and bar block
+- reset randomizes ring and bar positions in left/right planar regions with fixed orientation
+- the environment exposes task state for both objects
+- staged rewards progress to `max_reward = 5`
+- final success is based on finite-geometry insertion semantics, not a centerline-only shortcut
+- a new scripted policy can execute the intended pick-lift-align-insert behavior in the new environment
+- a canonical nominal smoke path (unit-level or deterministic integration-level) exists for the new scripted-policy interface so success is not judged purely by interpretation
+- existing `sim_transfer` behavior is preserved

docs/superpowers/specs/2026-05-25-native-smolvla-model-migration-design.md

@@ -0,0 +1,86 @@
+# Native SmolVLA Model Migration Design
+
+## Goal
+
+Migrate only the SmolVLA model core from `/data/lerobot-imf-attnres-exp/lerobot-imf-attnres` into this `roboimi` project so Diana simulation can train/evaluate through the existing `roboimi.vla` agent interface without importing or installing the full LeRobot tree.
+
+## Non-goals
+
+- Do not vendor the full `lerobot` package.
+- Do not change the current Python environment to LeRobot 0.5.x requirements.
+- Do not alter Diana environment action semantics.
+- Do not change `train_vla.py` or `eval_vla.py` main control flow unless a minimal compatibility hook is unavoidable.
+- Do not implement RTC in the first pass.
+
+## Architecture
+
+Add a native RoboIMI SmolVLA package under `roboimi/vla/models/smolvla/`. It will contain a lightweight config dataclass, a minimally adapted `SmolVLMWithExpertModel`, and a `VLAFlowMatching` implementation. A new `SmolVLANativeAgent` wraps the model with the existing RoboIMI agent contract: `compute_loss`, `predict_action_chunk`, `select_action`, `reset`, and `get_normalization_stats`.
+
+The new code will preserve the original model math where practical, but replace LeRobot framework dependencies with local constants, tokenizer handling, queue management, and `roboimi.vla.models.normalization.NormalizationModule`.
+
+## Data flow
+
+Training batch input remains the current RoboIMI format:
+
+```python
+{
+  "images": {cam: Tensor[B, T, C, H, W]},
+  "qpos": Tensor[B, T, 16],
+  "action": Tensor[B, H, 16],
+  "action_is_pad": optional BoolTensor[B, H],
+  "task": optional list[str],
+}
+```
+
+`SmolVLANativeAgent` normalizes `qpos` and `action` using RoboIMI dataset stats, tokenizes task strings, and passes images/state/language/action into the native SmolVLA model. In inference, the model returns normalized action chunks; the agent denormalizes them to 16-dim Diana EE actions.
+
+## Components
+
+### `roboimi/vla/models/smolvla/configuration.py`
+
+Defines `NativeSmolVLAConfig`, a small dataclass with fields needed by the model: state/action dimensions, padding dimensions, image resize target, tokenizer settings, VLM model name, VLM loading flags, expert layer configuration, sampling steps, dtype/device behavior, and compile flags.
+
+### `roboimi/vla/models/smolvla/smolvlm_with_expert.py`
+
+Migrates the model-core helper from the old repo. It should depend only on PyTorch and Transformers. It will expose `SmolVLMWithExpertModel` and attention helpers.
+
+### `roboimi/vla/models/smolvla/modeling.py`
+
+Defines model-core utilities (`resize_with_pad`, `pad_vector`, `make_att_2d_masks`, sinusoidal time embedding) plus `VLAFlowMatching`, with `forward` and `sample_actions`.
+
+### `roboimi/vla/agent_smolvla_native.py`
+
+RoboIMI-native agent wrapper. It owns tokenizer, normalization, task fallback, camera ordering, observation/action queues, loss masking, and denormalized rollout actions.
+
+### `roboimi/vla/conf/agent/smolvla_native.yaml`
+
+Hydra config for the native model, defaulting to Diana 16-dim state/action and configurable camera names.
+
+## Error handling
+
+- Missing configured camera raises `ValueError` with expected/missing names.
+- Task list length not matching batch size raises `ValueError`.
+- State/action dimensions exceeding configured `max_state_dim` / `max_action_dim` raises `ValueError`.
+- Missing Transformers SmolVLM classes raises `ImportError` explaining the required package.
+- Invalid action chunk shape raises `RuntimeError` explaining expected `(B,H,A)`.
+
+## Testing
+
+Use TDD with fake VLM/tokenizer/model components first, so tests do not download or instantiate the real SmolVLM. Cover:
+
+1. Config and Hydra instantiation with fake injected components.
+2. Camera ordering and missing-camera errors.
+3. Task fallback and newline/tokenization behavior.
+4. Loss path shape/mask behavior using a fake native model.
+5. `predict_action_chunk` normalization/denormalization and shape.
+6. `select_action` queue behavior.
+
+A later smoke test may instantiate the real model in an environment that already has the required Transformers/weights, but the first implementation must pass without external downloads.
+
+## Acceptance criteria
+
+- `agent=smolvla_native` can be composed by Hydra.
+- Unit tests pass without importing external `/data/.../src/lerobot`.
+- The new production code has no `import lerobot`.
+- The agent accepts the same batch/observation structure used by current train/eval scripts.
+- The agent emits 16-dim denormalized Diana EE actions for rollout.

render.py

@@ -0,0 +1,200 @@
+import mujoco
+from mujoco import viewer
+import sys
+import numpy as np
+import time
+import threading
+
+
+class MjBasicRenderer:
+    def __new__(cls, *args, **kwargs):
+        return super().__new__(cls)
+
+    def __init__(self, mj_model=None, mj_data=None):
+        # keyboard flag
+        self.render_paused = True
+        self.exit_flag = False
+        # init param
+        self.mj_model = mj_model
+        self.mj_data = mj_data
+        self.renderer = "viewer"  # default
+        self.viewer = None
+        self._image = None
+
+        # Set up mujoco viewer
+        self.image_renderer = mujoco.Renderer(self.mj_model)
+
+    def __del__(self):
+        pass
+
+    def _init_renderer(self):
+        """Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3),
+        another renderer with "mujoco_viewer"
+        """
+
+        def key_callback(keycode):
+            if keycode == 32:  # space
+                self.render_paused = not self.render_paused
+            elif keycode == 256:  # escape
+                self.exit_flag = not self.exit_flag
+
+        if self.renderer == "viewer":
+            # This function does not block, allowing user code to continue execution.
+            self.viewer = viewer.launch_passive(
+                self.mj_model,
+                self.mj_data,
+                key_callback=key_callback,
+                show_left_ui=False,
+                show_right_ui=False,
+            )
+            self.set_renderer_config()
+        else:
+            raise ValueError("Invalid renderer for some reason.")
+
+    def render(self):
+        """mujoco render"""
+        if self.viewer is not None and self.render_paused is True:
+            if self.viewer.is_running() and self.exit_flag is False:
+                self.viewer: viewer.Handle
+                self.viewer.sync()
+            else:
+                self.viewer.close()
+
+    def set_renderer_config(self):
+        """Setup mujoco global config while using viewer as renderer.
+        It should be noted that the render thread need locked.
+        """
+        self.viewer.cam.lookat = np.array([0.4, 0, 0.5])
+        self.viewer.cam.azimuth -= 0.005
+        with self.viewer.lock():
+            self.viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(
+                self.mj_data.time % 2
+            )
+
+
+try:
+    import cv2
+except ImportError:
+    print("Could not import cv2, please install it to enable camera viewer.")
+
+
+class MjMultiRenderer(MjBasicRenderer):
+
+    # __slots__=('mj_model','mj_data','renderer','enable_camera_viewer')
+    def __new__(cls, *args, **kwargs):
+        return super().__new__(cls)
+
+    def __init__(
+        self,
+        mj_model=None,
+        mj_data=None,
+        renderer=None,
+        enable_camera_viewer=False,
+        enable_depth=False,
+    ):
+        super().__init__(mj_model, mj_data)
+        self._depth = None
+        self.renderer = renderer
+        self._init_renderer()
+        self.enable_camera_viewer = enable_camera_viewer
+        if self.enable_camera_viewer:
+            self.enable_depth = enable_depth
+            self._init_window()
+        else:
+            self.enable_depth = False
+            print("No Camera View")
+
+    def __del__(self):
+        self.close()
+
+    def _init_renderer(self):
+        """
+        Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3)
+        """
+        if self.renderer == "unity":
+            # TODO: Support unity renderer.
+            raise ValueError("Unity renderer init failed for no supporting reason")
+        elif self.renderer == "viewer":
+            super()._init_renderer()
+            print("mujoco viewer init !")
+        else:
+            raise ValueError("renderer init failed for some reason.")
+
+    def _init_window(self, name="Camera view"):
+        if not self.enable_depth:
+            cv2.namedWindow(name, cv2.WINDOW_NORMAL)
+        else:
+            cv2.namedWindow(name, cv2.WINDOW_NORMAL)
+            cv2.namedWindow("Camera depth view", cv2.WINDOW_NORMAL)
+
+    def render(self):
+        """render mujoco"""
+        if self.renderer == "viewer":
+            super().render()
+        elif self.renderer == "unity":
+            # TODO: Support unity renderer.
+            raise ValueError("Unity renderer not supported now.")
+        else:
+            raise ValueError("Invalid renderer for some reason.")
+
+    def render(self):
+        """mujoco render"""
+        if self.viewer is not None and self.render_paused is True:
+            if self.viewer.is_running() and self.exit_flag is False:
+                self.viewer: viewer.Handle
+                self.viewer.sync()
+            else:
+                self.viewer.close()
+
+    def camera_render(self, cam=None):
+        if self.enable_camera_viewer:
+            if not self.enable_depth:
+                rgb, depth = self.render_from_camera(cam)
+                rgb = cv2.resize(rgb, (1920, 1600))
+                cv2.imshow("Camera view", rgb)
+                cv2.waitKey(1)
+
+            else:
+                rgb, depth = self.render_from_camera(cam)
+                cv2.imshow("Camera view", rgb)
+                cv2.imshow("Camera depth view", depth)
+                cv2.waitKey(1)
+
+        else:
+            print("camera info disable")
+            return
+
+    def render_from_camera(self, cam=None):
+        self.image_renderer.update_scene(self.mj_data, camera=cam)
+        if self.enable_depth is True:
+            self.image_renderer.enable_depth_rendering()
+            org = self.image_renderer.render()
+            depth = org[:, :]
+            self.image_renderer.disable_depth_rendering()
+            org = self.image_renderer.render()
+            image = org[:, :, ::-1]
+        else:
+            org = self.image_renderer.render()
+            image = org[:, :, ::-1]
+            depth = np.zeros([240, 320])
+        return image, depth
+
+    def close(self):
+        """close the environment."""
+        if self.enable_camera_viewer and self.viewer.is_running() == False:
+            cv2.destroyAllWindows()
+        self.viewer.close()
+        # sys.exit(0)
+
+    # def get_cam_intrinsic(self, fovy=45.0, width=320, height=240):
+    #     aspect = width * 1.0 / height
+    #     fovx = np.degrees(2 * np.arctan(aspect * np.tan(np.radians(fovy / 2))))
+
+    #     cx = 0.5 * width
+    #     cy = 0.5 * height
+    #     fx = cx / np.tan(fovx * np.pi / 180 * 0.5)
+    #     fy = cy / np.tan(fovy * np.pi / 180 * 0.5)
+
+    #     K = np.array([[fx, 0, cx],
+    #                   [0, fy, cy],
+    #                   [0, 0, 1]], dtype=np.float32)

roboimi/assets/models/manipulators/DianaMed/BiDianaMed_rethink.xml

@@ -76,7 +76,7 @@
                                                 <body name="ee_cam_left" pos="0.00 0.046 -0.075" euler="0.0 0.0 -0.0">
                                                         <inertial pos="0 0 0" quat="1 0 0 0" mass="0" diaginertia="0 0 0" />
                                                         <geom type="mesh" contype="1" conaffinity="1" group="1" rgba="0.69804 0.69804 0.69804 1" mesh="realsense_cam" />
-                                                        <camera name="rs_cam_left" mode="fixed" pos="0.0 0.0 -0.25" euler="0.0 9.4 0.0 " fovy="15" resolution="1920 1200"/>
+                                                        <camera name="rs_cam_left" mode="fixed" pos="0.0 0.0 0.01" euler="0.0 9.4 0.0 " fovy="50" resolution="1920 1200"/>
                                                 </body>
                                             </body>
                                             <body name="l_finger_left" pos="0 0.01 0.0444">

roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml

@@ -0,0 +1,6 @@
+<mujoco model="bi_diana_socket_peg">
+    <include file="./empty_world.xml" />
+    <include file="./table_square.xml" />
+    <include file="./socket_peg_objects.xml" />
+    <include file="./BiDianaMed_rethink.xml" />
+</mujoco>

roboimi/assets/models/manipulators/DianaMed/socket_peg_objects.xml

@@ -0,0 +1,19 @@
+<mujoco model="socket_peg_objects">
+  <worldbody>
+     <body name="peg" pos="0.12 0.90 0.46">
+          <joint name="red_peg_joint" type="free" frictionloss="0.01" />
+          <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg" rgba="1 0 0 1" />
+     </body>
+
+     <body name="socket" pos="-0.12 0.90 0.472">
+          <joint name="blue_socket_joint" type="free" frictionloss="0.01" />
+          <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 -0.02" size="0.06 0.018 0.002" type="box" name="socket-1" rgba="0 0 1 1" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 0.02" size="0.06 0.018 0.002" type="box" name="socket-2" rgba="0 0 1 1" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0.02 0" size="0.06 0.002 0.018" type="box" name="socket-3" rgba="0 0 1 1" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 -0.02 0" size="0.06 0.002 0.018" type="box" name="socket-4" rgba="0 0 1 1" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.04 0.01 0.01" type="box" name="pin" rgba="1 0 0 1" />
+     </body>
+  </worldbody>
+</mujoco>

roboimi/assets/models/manipulators/DianaMed/table_square.xml

@@ -7,7 +7,6 @@
           <geom name="table" condim="4" contype="1" conaffinity="1"  type="box" rgba="0.4 0.4 0.4 1" size="0.62 0.62 0.01" density="1500" friction="0.9 0.9 0.9"/>
       </body>
       <camera name="top" pos="0.0 1.0 2.0" fovy="44" mode="targetbody" target="table"/>
-      <camera name="angle" pos="0.0 0.0 2.0" fovy="37" mode="targetbody" target="table"/>
       <camera name="front" pos="0 0 0.8" fovy="65" mode="fixed" quat="0.7071 0.7071 0 0"/>
   </worldbody>
 </mujoco>

roboimi/assets/robots/diana_med.py

@@ -90,4 +90,40 @@ class BiDianaMed(ArmBase):
     def init_qpos(self):
         """ Robot's init joint position. """
         return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
+
+
+class BiDianaMedSocketPeg(ArmBase):
+    def __init__(self):
+        super().__init__(
+            name="Bidiana_socket_peg",
+            urdf_path="roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf",
+            xml_path="roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml",
+            gripper=None
+        )
+        self.left_arm = self.Arm(self, 'single', self.urdf_path)
+        self.left_arm.set_Arm_base_link('left_base_link')
+        self.left_arm.set_Arm_ee_link('left_link7')
+        self.left_arm.InitKDL
+        self.left_arm.joint_index = ['l_j1','l_j2','l_j3','l_j4','l_j5','l_j6','l_j7']
+        self.left_arm.gripper_index = ['l_finger_joint_left','r_finger_joint_left']
+        self.left_arm.actuator_index = ['a1_l','a2_l','a3_l','a4_l','a5_l','a6_l','a7_l','gripper_left']
+        self.left_arm.setArmInitPose(self.init_qpos)
+        self.arms.append(self.left_arm)
+        self.right_arm = self.Arm(self,'single', self.urdf_path)
+        self.right_arm.set_Arm_base_link('right_base_link')
+        self.right_arm.set_Arm_ee_link('right_link7')
+        self.right_arm.InitKDL
+        self.right_arm.joint_index = ['r_j1','r_j2','r_j3','r_j4','r_j5','r_j6','r_j7']
+        self.right_arm.gripper_index = ['l_finger_joint_right','r_finger_joint_right']
+        self.right_arm.actuator_index = ['a1_r','a2_r','a3_r','a4_r','a5_r','a6_r','a7_r','gripper_right']
+        self.right_arm.setArmInitPose(self.init_qpos)
+        self.arms.append(self.right_arm)
+        self.jnt_num = self.left_arm.jnt_num + self.right_arm.jnt_num
+        self.kp = 500 * np.ones(self.jnt_num)
+        self.kd = 44.57 * np.ones(self.jnt_num)
+
+    @property
+    def init_qpos(self):
+        """ Robot's init joint position. """
+        return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
     

roboimi/demos/diana_air_insert_policy.py

@@ -0,0 +1,184 @@
+import numpy as np
+from pyquaternion import Quaternion
+
+from roboimi.demos.diana_policy import PolicyBase
+
+
+class TestAirInsertPolicy(PolicyBase):
+    ACTION_OBJECT_Z_OFFSET = 0.078
+    SOCKET_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64)
+    PEG_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64)
+    SOCKET_OUTER_GRASP_STRATEGY = "socket_outer"
+    LEGACY_GRASP_STRATEGY = "legacy"
+    SOCKET_HOLD_Z = 0.85
+    PEG_INSERT_START_OFFSET = np.array([0.105, 0.0, 0.0], dtype=np.float64)
+    INSERT_END_T = 580
+    LEFT_SOCKET_GRIPPER_CLOSED = -100
+    RIGHT_PEG_GRIPPER_CLOSED = -100
+    SOCKET_APPROACH_Z = 1.05
+    EPISODE_END_T = 600
+
+    def __init__(self, inject_noise=False, grasp_strategy=SOCKET_OUTER_GRASP_STRATEGY):
+        super().__init__(inject_noise=inject_noise)
+        valid_strategies = {
+            self.SOCKET_OUTER_GRASP_STRATEGY,
+            self.LEGACY_GRASP_STRATEGY,
+        }
+        if grasp_strategy not in valid_strategies:
+            raise ValueError(
+                f"Unsupported air insert grasp_strategy={grasp_strategy!r}; "
+                f"expected one of {sorted(valid_strategies)}"
+            )
+        self.grasp_strategy = grasp_strategy
+
+    def generate_trajectory(self, task_state):
+        return self._generate_socket_peg_trajectory(task_state)
+
+    def _generate_socket_peg_trajectory(self, task_state):
+        socket_xyz = np.asarray(task_state["socket_pos"], dtype=np.float64)
+        peg_xyz = np.asarray(task_state["peg_pos"], dtype=np.float64)
+
+        init_mocap_pose_left = np.array(
+            [
+                -0.17297014,
+                1.00485877,
+                1.32773627,
+                7.06825181e-01,
+                8.20281078e-06,
+                -7.07388269e-01,
+                -5.20399313e-06,
+            ],
+            dtype=np.float64,
+        )
+        init_mocap_pose_right = np.array(
+            [
+                0.17297014,
+                0.9951369,
+                1.32773623,
+                2.59463975e-06,
+                7.07388269e-01,
+                5.59551158e-06,
+                7.06825181e-01,
+            ],
+            dtype=np.float64,
+        )
+
+        left_init_quat = Quaternion(init_mocap_pose_left[3:])
+        right_init_quat = Quaternion(init_mocap_pose_right[3:])
+
+        left_pick_quat = (
+            left_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45)
+        ).elements
+        right_pick_quat = (
+            right_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45)
+        ).elements
+
+        socket_hold_action = np.array(
+            [socket_xyz[0] - 0.078, socket_xyz[1], self.SOCKET_HOLD_Z], dtype=np.float64
+        )
+
+        peg_init_xyz = peg_xyz + np.array(
+            [0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET + 0.01]
+        )
+        peg_lift_center = np.array(
+            [peg_xyz[0] + 0.078, socket_hold_action[1], self.SOCKET_HOLD_Z - 0.01],
+            dtype=np.float64,
+        )
+        # The front camera looks along +Y, so visual right-to-left insertion is
+        # world +X -> -X.  With the socket XML in identity orientation, its
+        # tunnel axis is local/world X, so the peg approaches from +X and stops
+        # when its leading face reaches the socket's internal pin.
+        peg_insert_end_center = np.array(
+            [
+                socket_hold_action[0] + 0.078 * 2 + 0.04 + 0.06 - 0.01,
+                socket_hold_action[1],
+                self.SOCKET_HOLD_Z - 0.01,
+            ],
+            dtype=np.float64,
+        )
+
+        self.left_trajectory = [
+            {
+                "t": 1,
+                "xyz": init_mocap_pose_left[:3],
+                "quat": init_mocap_pose_left[3:],
+                "gripper": 100,
+            },
+            {
+                "t": 130,
+                "xyz": socket_xyz
+                + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]),
+                "quat": left_pick_quat,
+                "gripper": 100,
+            },
+            {
+                "t": 180,
+                "xyz": socket_xyz
+                + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]),
+                "quat": left_pick_quat,
+                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
+            },
+            {
+                "t": 350,
+                "xyz": socket_hold_action,
+                "quat": left_pick_quat,
+                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
+            },
+            {
+                "t": self.EPISODE_END_T,
+                "xyz": socket_hold_action,
+                "quat": left_pick_quat,
+                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
+            },
+        ]
+
+        self.right_trajectory = [
+            {
+                "t": 1,
+                "xyz": init_mocap_pose_right[:3],
+                "quat": init_mocap_pose_right[3:],
+                "gripper": 100,
+            },
+            {
+                "t": 80,
+                "xyz": peg_init_xyz,
+                "quat": right_pick_quat,
+                "gripper": 100,
+            },
+            {
+                "t": 150,
+                "xyz": peg_init_xyz,
+                "quat": right_pick_quat,
+                "gripper": 100,
+            },
+            {
+                "t": 180,
+                "xyz": peg_init_xyz,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": 350,
+                "xyz": peg_init_xyz,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": 450,
+                "xyz": peg_lift_center,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": self.INSERT_END_T,
+                "xyz": peg_insert_end_center,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": self.EPISODE_END_T,
+                "xyz": peg_insert_end_center,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+        ]

roboimi/demos/diana_record_sim_episodes.py

@@ -2,9 +2,11 @@ import time
 import os
 import numpy as np
 from roboimi.envs.double_pos_ctrl_env import make_sim_env
-from diana_policy import TestPickAndTransferPolicy
+from roboimi.demos.diana_air_insert_policy import TestAirInsertPolicy
+from roboimi.demos.diana_policy import TestPickAndTransferPolicy
 import cv2
-from roboimi.utils.act_ex_utils import sample_transfer_pose
+from roboimi.utils.act_ex_utils import sample_air_insert_socket_peg_state, sample_transfer_pose
+from roboimi.utils.constants import SIM_TASK_CONFIGS
 from roboimi.utils.streaming_episode_writer import StreamingEpisodeWriter
 
 import pathlib
@@ -12,16 +14,34 @@ HOME_PATH = str(pathlib.Path(__file__).parent.resolve())
 DATASET_DIR = HOME_PATH + '/dataset'
 
 
-def main():
-    task_name = 'sim_transfer'
-    dataset_dir = DATASET_DIR + '/sim_transfer'   #SIM_TASK_CONFIGS[task_name]['dataset_dir']
-    num_episodes = 100       #SIM_TASK_CONFIGS[task_name]['num_episodes']
+def sample_task_state(task_name):
+    if task_name == 'sim_transfer':
+        return sample_transfer_pose()
+    if task_name == 'sim_air_insert_socket_peg':
+        return sample_air_insert_socket_peg_state()
+    raise NotImplementedError(f'Unsupported scripted rollout task: {task_name}')
+
+
+def make_policy(task_name, inject_noise=False, grasp_strategy=None):
+    if task_name == 'sim_transfer':
+        return TestPickAndTransferPolicy(inject_noise)
+    if task_name == 'sim_air_insert_socket_peg':
+        if grasp_strategy is None:
+            return TestAirInsertPolicy(inject_noise)
+        return TestAirInsertPolicy(inject_noise, grasp_strategy=grasp_strategy)
+    raise NotImplementedError(f'Unsupported scripted rollout task: {task_name}')
+
+
+def main(task_name='sim_transfer'):
+    task_cfg = SIM_TASK_CONFIGS[task_name]
+    dataset_dir = task_cfg['dataset_dir']
+    num_episodes = 100
     inject_noise = False
 
-    episode_len  = 700      #SIM_TASK_CONFIGS[task_name]['episode_len']
-    camera_names = ['angle','r_vis', 'top', 'front']   #SIM_TASK_CONFIGS[task_name]['camera_names']
+    episode_len = task_cfg['episode_len']
+    camera_names = task_cfg['camera_names']
     image_size = (256, 256)
-    if task_name == 'sim_transfer':
+    if task_name in {'sim_transfer', 'sim_air_insert_socket_peg'}:
         print(task_name)
     else:
         raise NotImplementedError
@@ -29,7 +49,7 @@ def main():
     success = []
     
     env = make_sim_env(task_name)
-    policy = TestPickAndTransferPolicy(inject_noise)
+    policy = make_policy(task_name, inject_noise=inject_noise)
 
     # 等待osmesa完全启动后再开始收集数据
     print("等待osmesa线程启动...")
@@ -41,8 +61,8 @@ def main():
         max_reward = float('-inf')
         print(f'\n{episode_idx=}')
         print('Rollout out EE space scripted policy')
-        box_pos = sample_transfer_pose()
-        env.reset(box_pos)
+        task_state = sample_task_state(task_name)
+        env.reset(task_state)
         episode_writer = StreamingEpisodeWriter(
             dataset_path=os.path.join(dataset_dir, f'episode_{episode_idx}.hdf5'),
             max_timesteps=episode_len,
@@ -50,7 +70,7 @@ def main():
             image_size=image_size,
         )
         for step in range(episode_len):
-            raw_action = policy.predict(box_pos,step)
+            raw_action = policy.predict(task_state, step)
             env.step(raw_action)
             env.render()
             sum_reward += env.rew

roboimi/demos/vla_scripts/train_vla.py

@@ -71,6 +71,19 @@ from hydra.utils import instantiate
 
 log = logging.getLogger(__name__)
 
+
+SMOLVLA_NATIVE_TRAINING_PRESET = {
+    'lr': 1e-4,
+    'betas': (0.9, 0.95),
+    'eps': 1e-8,
+    'weight_decay': 1e-10,
+    'grad_clip': 10.0,
+    'warmup_steps': 1000,
+    'scheduler_type': 'cosine',
+    'scheduler_decay_steps': 30000,
+    'scheduler_decay_lr': 2.5e-6,
+}
+
 # 注册列表长度解析器（用于配置中如 ${len:${data.camera_names}}）
 if not OmegaConf.has_resolver("len"):
     OmegaConf.register_new_resolver("len", lambda x: len(x))
@@ -154,7 +167,7 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
     Args:
         optimizer: PyTorch 优化器
         warmup_steps: 预热步数
-        max_steps: 总训练步数
+        max_steps: 余弦衰减步数
         scheduler_type: 预热后的调度器类型 ('cosine' 或 'constant')
         min_lr: 最小学习率（用于余弦衰减）
 
@@ -167,16 +180,24 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
     min_lr_ratio = min_lr / base_lr if base_lr > 0 else 0.0
 
     def lr_lambda(step):
-        # 预热阶段：从 0 线性增加到 1
+        # LeRobot CosineDecayWithWarmupSchedulerConfig 的线性预热：
+        # 从一个很小的非零 LR 开始，避免首步完全为 0。
         if step < warmup_steps:
-            return float(step) / float(max(1, warmup_steps))
+            if step <= 0:
+                return 1.0 / float(max(1, warmup_steps + 1))
+            frac = 1.0 - float(step) / float(max(1, warmup_steps))
+            return (1.0 / float(max(1, warmup_steps + 1)) - 1.0) * frac + 1.0
 
         # 预热后阶段
         if scheduler_type == 'cosine':
-            # 从 1 到 min_lr_ratio 的余弦退火
-            progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps))
-            cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
-            return max(min_lr_ratio, cosine_decay)
+            # 与 LeRobot SmolVLA/PI0 的 CosineDecayWithWarmupSchedulerConfig 对齐：
+            # 1) 余弦衰减步数是固定的 num_decay_steps（SmolVLA 默认 30k），
+            # 2) 超过 decay_steps 后 clamp 到 decay_lr，不能继续 cos() 进入下一周期；
+            #    否则 150k 训练会显示成“正弦波”式反复升降。
+            decay_steps = max(1, int(max_steps))
+            clamped_step = min(max(int(step), 0), decay_steps)
+            cosine_decay = 0.5 * (1.0 + math.cos(math.pi * clamped_step / decay_steps))
+            return (1.0 - min_lr_ratio) * cosine_decay + min_lr_ratio
         else:
             # 恒定学习率
             return 1.0
@@ -184,7 +205,119 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
     return LambdaLR(optimizer, lr_lambda)
 
 
-def build_training_optimizer(agent, lr, weight_decay):
+def _is_smolvla_native_agent_config(agent_cfg) -> bool:
+    target = str(agent_cfg.get('_target_', '')) if hasattr(agent_cfg, 'get') else ''
+    return target == 'roboimi.vla.agent_smolvla_native.SmolVLANativeAgent'
+
+
+def resolve_training_recipe(cfg):
+    """Return optimizer/scheduler knobs, using LeRobot SmolVLA defaults for smolvla_native."""
+    if _is_smolvla_native_agent_config(cfg.agent):
+        preset = SMOLVLA_NATIVE_TRAINING_PRESET
+        scheduler_steps = int(cfg.train.get(
+            'scheduler_decay_steps',
+            cfg.train.get('lr_scheduler_steps', cfg.train.max_steps),
+        ))
+        return {
+            'lr': float(preset['lr']),
+            'betas': tuple(preset['betas']),
+            'eps': float(preset['eps']),
+            'weight_decay': float(preset['weight_decay']),
+            'grad_clip': float(preset['grad_clip']),
+            'warmup_steps': int(preset['warmup_steps']),
+            'scheduler_type': str(preset['scheduler_type']),
+            'scheduler_steps': scheduler_steps,
+            'min_lr': float(preset['scheduler_decay_lr']),
+            'preset_name': 'lerobot_smolvla',
+        }
+
+    return {
+        'lr': float(cfg.train.lr),
+        'betas': (0.9, 0.999),
+        'eps': 1e-8,
+        'weight_decay': float(cfg.train.get('weight_decay', 1e-5)),
+        'grad_clip': float(cfg.train.get('grad_clip', 1.0)),
+        'warmup_steps': int(cfg.train.get('warmup_steps', 500)),
+        'scheduler_type': str(cfg.train.get('scheduler_type', 'cosine')),
+        'scheduler_steps': int(cfg.train.max_steps),
+        'min_lr': float(cfg.train.get('min_lr', 1e-6)),
+        'preset_name': 'config',
+    }
+
+
+def _instantiate_dataset(cfg, dataset_image_resize_shape, episode_indices=None):
+    kwargs = {'image_resize_shape': dataset_image_resize_shape}
+    if episode_indices is not None:
+        kwargs['episode_indices'] = episode_indices
+    return instantiate(cfg.data, **kwargs)
+
+
+def _resolve_dataset_image_resize_shape(cfg):
+    dataset_image_resize_shape = cfg.data.get('image_resize_shape', (224, 224))
+    agent_cfg = cfg.agent
+    if 'dataset_image_resize_shape' in agent_cfg:
+        return agent_cfg.get('dataset_image_resize_shape')
+    for backbone_key in ('vision_backbone', 'condition_encoder'):
+        backbone_cfg = agent_cfg.get(backbone_key, None)
+        if backbone_cfg is not None and 'dataset_image_resize_shape' in backbone_cfg:
+            dataset_image_resize_shape = backbone_cfg.get('dataset_image_resize_shape')
+            break
+    return dataset_image_resize_shape
+
+
+def build_train_val_datasets(cfg, dataset_image_resize_shape):
+    val_episode_indices = cfg.train.get('val_episode_indices', None)
+    if val_episode_indices:
+        dataset = _instantiate_dataset(cfg, dataset_image_resize_shape)
+        available_episode_indices = list(getattr(dataset, 'available_episode_indices', []))
+        if not available_episode_indices:
+            raise ValueError('显式 val_episode_indices 需要数据集暴露 available_episode_indices')
+
+        requested_val_episode_indices = sorted(int(idx) for idx in val_episode_indices)
+        available_set = set(available_episode_indices)
+        missing = sorted(set(requested_val_episode_indices) - available_set)
+        if missing:
+            raise ValueError(
+                f'val_episode_indices {missing} 不存在于数据集可用 episodes {available_episode_indices}'
+            )
+
+        val_set = set(requested_val_episode_indices)
+        train_episode_indices = [
+            idx for idx in available_episode_indices
+            if idx not in val_set
+        ]
+        if not train_episode_indices:
+            raise ValueError('显式 val_episode_indices 不能覆盖全部 episodes，训练集将为空')
+
+        train_dataset = _instantiate_dataset(
+            cfg,
+            dataset_image_resize_shape,
+            episode_indices=train_episode_indices,
+        )
+        val_dataset = _instantiate_dataset(
+            cfg,
+            dataset_image_resize_shape,
+            episode_indices=requested_val_episode_indices,
+        )
+        return dataset, train_dataset, val_dataset, requested_val_episode_indices
+
+    dataset = _instantiate_dataset(cfg, dataset_image_resize_shape)
+    val_split = float(cfg.train.get('val_split', 0.1))
+    seed = int(cfg.train.get('seed', 42))
+    val_size = int(len(dataset) * val_split)
+    train_size = len(dataset) - val_size
+    if val_size > 0:
+        train_dataset, val_dataset = random_split(
+            dataset,
+            [train_size, val_size],
+            generator=torch.Generator().manual_seed(seed)
+        )
+    else:
+        train_dataset, val_dataset = dataset, None
+    return dataset, train_dataset, val_dataset, None
+
+
+def build_training_optimizer(agent, lr, weight_decay, betas=(0.9, 0.999), eps=1e-8):
     """为训练脚本构建优化器，优先复用任意 head 自带的参数分组。"""
     trainable_params = [param for param in agent.parameters() if param.requires_grad]
     noise_pred_net = getattr(agent, 'noise_pred_net', None)
@@ -192,7 +325,7 @@ def build_training_optimizer(agent, lr, weight_decay):
     use_head_groups = callable(get_optim_groups)
 
     if not use_head_groups:
-        return AdamW(trainable_params, lr=lr, weight_decay=weight_decay)
+        return AdamW(trainable_params, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
 
     head_groups = []
     grouped_param_ids = set()
@@ -234,7 +367,7 @@ def build_training_optimizer(agent, lr, weight_decay):
     if grouped_param_ids != all_trainable_param_ids:
         raise ValueError('Optimizer parameter groups must include each trainable parameter exactly once')
 
-    return AdamW(optim_groups, lr=lr, weight_decay=weight_decay)
+    return AdamW(optim_groups, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
 
 
 def _init_swanlab(cfg):
@@ -281,7 +414,12 @@ def _init_swanlab(cfg):
         init_kwargs['experiment_name'] = run_name
 
     try:
-        swanlab.init(**init_kwargs)
+        run = swanlab.init(**init_kwargs)
+        if run is not None:
+            try:
+                setattr(swanlab, "_roboimi_swanlab_run", run)
+            except Exception:
+                pass
     except Exception as exc:
         raise RuntimeError(
             f"SwanLab logging is enabled, but SwanLab init/login failed: {exc}"
@@ -290,6 +428,24 @@ def _init_swanlab(cfg):
     return swanlab
 
 
+def _log_swanlab_init_details(swanlab_module):
+    if swanlab_module is None:
+        return
+    run = getattr(swanlab_module, "_roboimi_swanlab_run", None)
+    if run is None:
+        run = getattr(swanlab_module, "run", None)
+    if run is None:
+        return
+
+    url = getattr(run, "url", None)
+    swanlog_dir = getattr(run, "swanlog_dir", None)
+    log.info(
+        "🦢 SwanLab initialized%s%s",
+        f" | url={url}" if url else "",
+        f" | swanlog_dir={swanlog_dir}" if swanlog_dir else "",
+    )
+
+
 def _log_to_swanlab(swanlab_module, payload, step=None):
     if swanlab_module is None:
         return
@@ -368,7 +524,13 @@ def _run_training(cfg: DictConfig):
     log.info(f"🚀 开始 VLA 训练 (设备: {cfg.train.device})")
     _configure_cuda_runtime(cfg)
     swanlab_module = _init_swanlab(cfg)
+    _log_swanlab_init_details(swanlab_module)
     try:
+        action_mse_val_freq_epochs = int(cfg.train.get('action_mse_val_freq_epochs', 0) or 0)
+        explicit_val_episode_indices = cfg.train.get('val_episode_indices', None)
+        if action_mse_val_freq_epochs > 0 and not explicit_val_episode_indices:
+            raise ValueError('action_mse_val_freq_epochs > 0 requires train.val_episode_indices')
+
         # 创建检查点目录
         run_output_dir = _resolve_run_output_dir()
         checkpoint_dir = run_output_dir / "checkpoints"
@@ -380,34 +542,31 @@ def _run_training(cfg: DictConfig):
         # =========================================================================
         log.info("📦 加载数据集...")
         try:
-            dataset_image_resize_shape = cfg.data.get('image_resize_shape', (224, 224))
-            vision_backbone_cfg = cfg.agent.get('vision_backbone', None)
-            if vision_backbone_cfg is not None and 'dataset_image_resize_shape' in vision_backbone_cfg:
-                dataset_image_resize_shape = vision_backbone_cfg.get('dataset_image_resize_shape')
-            dataset = instantiate(
-                cfg.data,
-                image_resize_shape=dataset_image_resize_shape,
+            dataset_image_resize_shape = _resolve_dataset_image_resize_shape(cfg)
+            dataset, train_dataset, val_dataset, explicit_val_episode_indices = (
+                build_train_val_datasets(cfg, dataset_image_resize_shape)
             )
             log.info(f"✅ 数据集加载成功。总样本数: {len(dataset)}")
         except Exception as e:
             log.error(f"❌ 数据集加载失败: {e}")
             raise
 
-        # 训练/验证集划分
-        val_split = float(cfg.train.get('val_split', 0.1))
-        seed = int(cfg.train.get('seed', 42))
-        val_size = int(len(dataset) * val_split)
-        train_size = len(dataset) - val_size
-        if val_size > 0:
-            train_dataset, val_dataset = random_split(
-                dataset,
-                [train_size, val_size],
-                generator=torch.Generator().manual_seed(seed)
+        if explicit_val_episode_indices is not None:
+            log.info(
+                "✅ 数据集划分: 训练集=%s, 验证集=%s (显式 held-out episodes=%s)",
+                len(train_dataset),
+                len(val_dataset),
+                explicit_val_episode_indices,
             )
-            log.info(f"✅ 数据集划分: 训练集={train_size}, 验证集={val_size} (验证比例={val_split})")
         else:
-            train_dataset, val_dataset = dataset, None
-            log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)")
+            val_split = float(cfg.train.get('val_split', 0.1))
+            val_size = len(val_dataset) if val_dataset is not None else 0
+            if val_size > 0:
+                log.info(
+                    f"✅ 数据集划分: 训练集={len(train_dataset)}, 验证集={val_size} (验证比例={val_split})"
+                )
+            else:
+                log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)")
 
         train_batch_size = int(cfg.train.batch_size)
         train_drop_last = len(train_dataset) >= train_batch_size
@@ -535,21 +694,35 @@ def _run_training(cfg: DictConfig):
         # =========================================================================
         # 4. 设置优化器与学习率调度器
         # =========================================================================
-        weight_decay = float(cfg.train.get('weight_decay', 1e-5))
-        grad_clip = float(cfg.train.get('grad_clip', 1.0))
-
-        optimizer = build_training_optimizer(agent, lr=cfg.train.lr, weight_decay=weight_decay)
-        log.info(f"🔧 优化器: AdamW (学习率={cfg.train.lr}, weight_decay={weight_decay})")
+        training_recipe = resolve_training_recipe(cfg)
+        weight_decay = training_recipe['weight_decay']
+        grad_clip = training_recipe['grad_clip']
+        train_lr = training_recipe['lr']
+        optimizer = build_training_optimizer(
+            agent,
+            lr=train_lr,
+            weight_decay=weight_decay,
+            betas=training_recipe['betas'],
+            eps=training_recipe['eps'],
+        )
+        log.info(
+            "🔧 优化器: AdamW (preset=%s, 学习率=%s, betas=%s, eps=%s, weight_decay=%s)",
+            training_recipe['preset_name'],
+            train_lr,
+            training_recipe['betas'],
+            training_recipe['eps'],
+            weight_decay,
+        )
 
         # 设置带预热的学習率调度器
-        warmup_steps = int(cfg.train.get('warmup_steps', 500))
-        scheduler_type = cfg.train.get('scheduler_type', 'cosine')
-        min_lr = float(cfg.train.get('min_lr', 1e-6))
+        warmup_steps = training_recipe['warmup_steps']
+        scheduler_type = training_recipe['scheduler_type']
+        min_lr = training_recipe['min_lr']
 
         scheduler = get_lr_schedule_with_warmup(
             optimizer,
             warmup_steps=warmup_steps,
-            max_steps=cfg.train.max_steps,
+            max_steps=training_recipe['scheduler_steps'],
             scheduler_type=scheduler_type,
             min_lr=min_lr
         )
@@ -652,12 +825,16 @@ def _run_training(cfg: DictConfig):
                 if key in batch_data:
                     images[cam_name] = batch_data[key]
 
-            return {
+            agent_input = {
                 'images': images,
                 'qpos': batch_data['observation.state'],  # SimpleRobotDataset 使用 observation.state
                 'action': batch_data['action'],
                 'action_is_pad': batch_data.get('action_is_pad', None)  # 传递padding mask
             }
+            if 'task' in batch_data:
+                agent_input['task'] = batch_data['task']
+
+            return agent_input
 
         def save_checkpoint(checkpoint_path: Path, step: int, loss_value, val_loss=None, rollout_avg_reward=None):
             agent_stats = agent.get_normalization_stats()
@@ -702,10 +879,28 @@ def _run_training(cfg: DictConfig):
             from roboimi.demos.vla_scripts import eval_vla
 
             rollout_cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=False))
+            rollout_num_episodes = int(cfg.train.get('rollout_num_episodes', 1))
+            rollout_device = str(cfg.train.get('rollout_device', cfg.train.device))
+            configured_rollout_workers = cfg.train.get('rollout_num_workers', None)
+            if configured_rollout_workers is None:
+                if rollout_device.startswith('cuda'):
+                    rollout_num_workers = min(max(rollout_num_episodes, 1), 8)
+                else:
+                    rollout_num_workers = 1
+            else:
+                rollout_num_workers = int(configured_rollout_workers)
             rollout_cfg.eval.ckpt_path = str(checkpoint_path)
-            rollout_cfg.eval.num_episodes = int(cfg.train.get('rollout_num_episodes', 1))
+            rollout_cfg.eval.num_episodes = rollout_num_episodes
+            rollout_cfg.eval.num_workers = rollout_num_workers
             rollout_cfg.eval.headless = True
-            rollout_cfg.eval.device = 'cpu'
+            rollout_cfg.eval.device = rollout_device
+            rollout_cfg.eval.cuda_devices = cfg.train.get('rollout_cuda_devices', None)
+            rollout_cfg.eval.response_timeout_s = float(
+                cfg.train.get('rollout_response_timeout_s', 300.0)
+            )
+            rollout_cfg.eval.server_startup_timeout_s = float(
+                cfg.train.get('rollout_server_startup_timeout_s', 300.0)
+            )
             rollout_cfg.eval.verbose_action = False
             rollout_cfg.eval.record_video = False
             rollout_cfg.eval.save_trajectory_image = True
@@ -716,9 +911,11 @@ def _run_training(cfg: DictConfig):
             )
 
             log.info(
-                "🎯 开始 checkpoint rollout 验证: %s (episodes=%s, headless=True)",
+                "🎯 开始 checkpoint rollout 验证: %s (episodes=%s, device=%s, workers=%s, headless=True)",
                 checkpoint_path,
                 rollout_cfg.eval.num_episodes,
+                rollout_cfg.eval.device,
+                rollout_cfg.eval.num_workers,
             )
             return eval_vla._run_eval(rollout_cfg)
 
@@ -884,6 +1081,31 @@ def _run_training(cfg: DictConfig):
                 completed_steps // steps_per_epoch
                 if steps_per_epoch > 0 else 0
             )
+            should_run_action_mse_val = (
+                val_loader is not None
+                and explicit_val_episode_indices is not None
+                and action_mse_val_freq_epochs > 0
+                and steps_per_epoch > 0
+                and completed_steps % steps_per_epoch == 0
+                and completed_epoch > 0
+                and completed_epoch % action_mse_val_freq_epochs == 0
+            )
+            if should_run_action_mse_val:
+                val_loss = run_validation()
+                if val_loss is not None:
+                    log.info(
+                        f"步骤 {step}/{cfg.train.max_steps} | Epoch {completed_epoch} "
+                        f"held-out action MSE: {val_loss:.6f}"
+                    )
+                    _log_to_swanlab(
+                        swanlab_module,
+                        {
+                            'val/action_mse': val_loss,
+                            'val/epoch': completed_epoch,
+                        },
+                        step=step,
+                    )
+
             should_run_epoch_rollout = (
                 rollout_validation_enabled
                 and steps_per_epoch > 0

roboimi/envs/double_air_insert_env.py

@@ -0,0 +1,157 @@
+import copy as cp
+import time
+
+import numpy as np
+
+from roboimi.envs.double_base import DualDianaMed
+from roboimi.envs.double_pos_ctrl_env import DualDianaMed_Pos_Ctrl
+
+
+SOCKET_JOINT_NAME = "blue_socket_joint"
+PEG_JOINT_NAME = "red_peg_joint"
+REQUIRED_TASK_STATE_KEYS = ("socket_pos", "socket_quat", "peg_pos", "peg_quat")
+SOCKET_GEOM_NAMES = ("socket-1", "socket-2", "socket-3", "socket-4")
+SOCKET_SUCCESS_GEOM_NAMES = ("pin",)
+SOCKET_BODY_GEOM_NAMES = SOCKET_GEOM_NAMES + SOCKET_SUCCESS_GEOM_NAMES
+PEG_GEOM_NAMES = ("red_peg",)
+LEFT_GRIPPER_GEOM_NAMES = (
+    "l_finger_left",
+    "r_finger_left",
+    "l_fingertip_g0_left",
+    "r_fingertip_g0_left",
+    "l_fingerpad_g0_left",
+    "r_fingerpad_g0_left",
+    "l_fingertip_g0_vis_left",
+    "r_fingertip_g0_vis_left",
+)
+RIGHT_GRIPPER_GEOM_NAMES = (
+    "l_finger_right",
+    "r_finger_right",
+    "l_fingertip_g0_right",
+    "r_fingertip_g0_right",
+    "l_fingerpad_g0_right",
+    "r_fingerpad_g0_right",
+    "l_fingertip_g0_vis_right",
+    "r_fingertip_g0_vis_right",
+)
+TABLE_GEOM_NAME = "table"
+
+
+def _set_free_joint_pose(joint, position, quat):
+    joint.qpos[:3] = np.asarray(position, dtype=np.float64)
+    joint.qpos[3:7] = np.asarray(quat, dtype=np.float64)
+
+
+def set_socket_peg_task_state(mj_data, task_state):
+    if not isinstance(task_state, dict) or tuple(task_state.keys()) != REQUIRED_TASK_STATE_KEYS:
+        raise ValueError(
+            "task_state must be an ordered dict-like mapping with keys "
+            "socket_pos, socket_quat, peg_pos, peg_quat"
+        )
+
+    _set_free_joint_pose(
+        mj_data.joint(SOCKET_JOINT_NAME),
+        task_state["socket_pos"],
+        task_state["socket_quat"],
+    )
+    _set_free_joint_pose(
+        mj_data.joint(PEG_JOINT_NAME),
+        task_state["peg_pos"],
+        task_state["peg_quat"],
+    )
+
+
+def get_socket_peg_env_state(mj_data):
+    socket_qpos = cp.deepcopy(np.asarray(mj_data.joint(SOCKET_JOINT_NAME).qpos[:7], dtype=np.float64))
+    peg_qpos = cp.deepcopy(np.asarray(mj_data.joint(PEG_JOINT_NAME).qpos[:7], dtype=np.float64))
+    return np.concatenate([socket_qpos, peg_qpos], dtype=np.float64)
+
+
+def _normalize_contact_pairs(contact_pairs):
+    return {frozenset(pair) for pair in contact_pairs}
+
+
+def _has_any_object_contact(contact_set, object_geom_names, other_geom_names):
+    return any(
+        frozenset((object_geom_name, other_geom_name)) in contact_set
+        for object_geom_name in object_geom_names
+        for other_geom_name in other_geom_names
+    )
+
+
+def _object_is_airborne(contact_set, object_geom_names):
+    return not _has_any_object_contact(contact_set, object_geom_names, (TABLE_GEOM_NAME,))
+
+
+def peg_inserted_into_socket(contact_pairs):
+    contact_set = _normalize_contact_pairs(contact_pairs)
+    return frozenset((PEG_GEOM_NAMES[0], SOCKET_SUCCESS_GEOM_NAMES[0])) in contact_set
+
+
+def compute_air_insert_reward(contact_pairs, env_state=None):
+    del env_state  # kept for API compatibility with rollout/eval code paths
+    contact_set = _normalize_contact_pairs(contact_pairs)
+    reward = 0
+
+    if _has_any_object_contact(contact_set, SOCKET_GEOM_NAMES, LEFT_GRIPPER_GEOM_NAMES):
+        reward += 1
+    if _has_any_object_contact(contact_set, PEG_GEOM_NAMES, RIGHT_GRIPPER_GEOM_NAMES):
+        reward += 1
+
+    socket_airborne = _object_is_airborne(contact_set, SOCKET_BODY_GEOM_NAMES)
+    peg_airborne = _object_is_airborne(contact_set, PEG_GEOM_NAMES)
+    if socket_airborne:
+        reward += 1
+    if peg_airborne:
+        reward += 1
+
+    if socket_airborne and peg_airborne and peg_inserted_into_socket(contact_pairs):
+        reward += 1
+
+    return reward
+
+
+class DualDianaMed_Air_Insert(DualDianaMed_Pos_Ctrl):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_reward = 5
+
+    def reset(self, task_state):
+        set_socket_peg_task_state(self.mj_data, task_state)
+        DualDianaMed.reset(self)
+        self.top = None
+        self.r_vis = None
+        self.l_vis = None
+        self.front = None
+        if not self.is_render:
+            self._update_camera_images_sync()
+            return
+        self.cam_flage = True
+        while self.cam_flage:
+            if (
+                type(self.top) == type(None)
+                or type(self.r_vis) == type(None)
+                or type(self.l_vis) == type(None)
+                or type(self.front) == type(None)
+            ):
+                time.sleep(0.001)
+            else:
+                self.cam_flage = False
+
+    def step(self, action=np.zeros(16)):
+        super().step(action)
+        self.rew = self._get_reward()
+        self.obs = self._get_obs()
+
+    def get_env_state(self):
+        return get_socket_peg_env_state(self.mj_data)
+
+    def _get_reward(self):
+        contact_pairs = []
+        for collision_num in range(self.mj_data.ncon):
+            geom1 = self.mj_data.contact[collision_num].geom1
+            geom2 = self.mj_data.contact[collision_num].geom2
+            contact_pairs.append(
+                (self.getID2Name("geom", geom1), self.getID2Name("geom", geom2))
+            )
+        return compute_air_insert_reward(contact_pairs, self.get_env_state())

roboimi/envs/double_base.py

@@ -52,7 +52,6 @@ class DualDianaMed(MujocoEnv):
         self.r_vis = None
         self.l_vis = None
         self.top = None
-        self.angle = None
         self.front = None
         self.obs = None
         
@@ -92,7 +91,6 @@ class DualDianaMed(MujocoEnv):
 
     def step(self,action):
         self.compute_qpos = action  #for observation !
-        self.obs = self._get_obs()
         if self.interpolator_left is not None and self.interpolator_right is not None:
             self.interpolator_left.updateInput(action[:7], control_cycle=self.base_time)
             self.interpolator_right.updateInput(action[7:-2], control_cycle=self.base_time)
@@ -105,6 +103,7 @@ class DualDianaMed(MujocoEnv):
 
             super().step(action)
         self.base_time = time.time() - ctrl_cur_time
+        self.obs = self._get_obs()
             
 
     def preStep(self, action):
@@ -168,10 +167,9 @@ class DualDianaMed(MujocoEnv):
         obs['qpos'] = self.get_obs_qpos
         obs['action'] = self.compute_qpos
         obs['images'] = dict()
-        obs['images']['top']   = self.top
-        obs['images']['angle'] = self.angle
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
+        obs['images']['top'] = self.top
         obs['images']['front'] = self.front
         return obs
 
@@ -180,10 +178,9 @@ class DualDianaMed(MujocoEnv):
             self._update_camera_images_sync()
         obs = collections.OrderedDict()
         obs['images'] = dict()
-        obs['images']['top']   = self.top
-        obs['images']['angle'] = self.angle
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
+        obs['images']['top'] = self.top
         obs['images']['front'] = self.front
         return obs
     
@@ -202,14 +199,12 @@ class DualDianaMed(MujocoEnv):
 
     @property
     def cam_view(self):
-        if self.cam == 'top':
-            return self.top
-        elif self.cam == 'angle':
-            return self.angle
-        elif self.cam == 'r_vis':
+        if self.cam == 'r_vis':
             return self.r_vis
         elif self.cam == 'l_vis':
             return self.l_vis
+        elif self.cam == 'top':
+            return self.top
         elif self.cam == 'front':
             return self.front
         else:
@@ -230,8 +225,6 @@ class DualDianaMed(MujocoEnv):
         self.l_vis = img_renderer.render()[:, :, ::-1]
         img_renderer.update_scene(self.mj_data, camera="top")
         self.top = img_renderer.render()[:, :, ::-1]
-        img_renderer.update_scene(self.mj_data, camera="angle")
-        self.angle = img_renderer.render()[:, :, ::-1]
         img_renderer.update_scene(self.mj_data, camera="front")
         self.front = img_renderer.render()[:, :, ::-1]
 

roboimi/envs/double_pos_ctrl_env.py

@@ -34,19 +34,19 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
             is_interpolate=is_interpolate,
             cam_view=cam_view
         )
-        
+
         self.max_reward = 4
-        
+
         self.cam_start()
 
-        
+
     def step(self,action=np.zeros(16)):
         action_left  = self.ik_solve(action[:3],action[3:7],self.arm_left)
         action_right = self.ik_solve(action[7:10],action[10:14],self.arm_right)
         action = np.hstack((action_left,action_right,action[14:]))
         super().step(action)
         self.rew = self._get_reward()
-        
+
     def step_jnt(self,action):
         super().step(action)
 
@@ -63,8 +63,8 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         return Arm.kdl_solver.ikSolver(p_goal, mat_goal, Arm.arm_qpos)
 
     def reset(self,box_pos):
-        
-        self.mj_data.joint('red_box_joint').qpos[0] = box_pos[0] 
+
+        self.mj_data.joint('red_box_joint').qpos[0] = box_pos[0]
         self.mj_data.joint('red_box_joint').qpos[1] = box_pos[1]
         self.mj_data.joint('red_box_joint').qpos[2] = box_pos[2]
         self.mj_data.joint('red_box_joint').qpos[3] = 1.0
@@ -73,8 +73,8 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         self.mj_data.joint('red_box_joint').qpos[6] = 0.0
         super().reset()
         self.top = None
-        self.angle = None
         self.r_vis = None
+        self.l_vis = None
         self.front = None
         if not self.is_render:
             self._update_camera_images_sync()
@@ -82,16 +82,16 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         self.cam_flage = True
         t=0
         while self.cam_flage:
-            if(type(self.top)==type(None) 
-               or type(self.angle)==type(None) 
+            if(type(self.top)==type(None)
                or type(self.r_vis)==type(None)
+               or type(self.l_vis)==type(None)
                or type(self.front)==type(None)):
                 time.sleep(0.001)
                 t+=1
             else:
                 self.cam_flage=False
-        
-        
+
+
 
     def preStep(self, action):
         if isinstance(action,np.ndarray) and len(action)==16:
@@ -104,7 +104,7 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         for i in range(3):
             box_pose[i] = cp.deepcopy(self.mj_data.joint('red_box_joint').qpos[i])
         return box_pose
-    
+
 
     def _get_reward(self):
         all_contact_pairs = []
@@ -127,16 +127,28 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         reward = 0
         if touch_right_gripper and not touch_table:
             reward = 1
-        if touch_right_gripper and not box_touch_table: 
+        if touch_right_gripper and not box_touch_table:
             reward = 2
         if touch_left_gripper: # attempted transfer
             reward = 3
         if touch_left_gripper and not box_touch_table: # successful transfer
             reward = 4
         return reward
-        
+
 
 def make_sim_env(task_name, headless=False):
+    if task_name == 'sim_air_insert_socket_peg':
+        from roboimi.assets.robots.diana_med import BiDianaMedSocketPeg
+        from roboimi.envs.double_air_insert_env import DualDianaMed_Air_Insert
+
+        env = DualDianaMed_Air_Insert(
+            robot=BiDianaMedSocketPeg(),
+            is_render=not headless,
+            control_freq=30,
+            is_interpolate=True,
+            cam_view='front'
+        )
+        return env
     if 'sim_transfer' in task_name:
         from roboimi.assets.robots.diana_med import BiDianaMed
         env = DualDianaMed_Pos_Ctrl(
@@ -144,7 +156,7 @@ def make_sim_env(task_name, headless=False):
             is_render=not headless,
             control_freq=30,
             is_interpolate=True,
-            cam_view='angle'
+            cam_view='top'
         )
         return env
     else:
@@ -170,4 +182,3 @@ if __name__ == "__main__":
         env.step(action)
         if env.is_render:
             env.render()
-        

roboimi/utils/act_ex_utils.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+
 def sample_insertion_pose():
         # Peg
         x_range = [0.1, 0.2]
@@ -35,4 +36,23 @@ def sample_transfer_pose():
         box_position = np.random.uniform(ranges[:, 0], ranges[:, 1])
 
 
-        return box_position
+        return box_position
+
+
+def sample_air_insert_socket_peg_state():
+        socket_position = np.random.uniform(
+            low=np.array([-0.20, 0.80, 0.472], dtype=np.float32),
+            high=np.array([-0.10, 1.00, 0.472], dtype=np.float32),
+        )
+        peg_position = np.random.uniform(
+            low=np.array([0.10, 0.80, 0.46], dtype=np.float32),
+            high=np.array([0.20, 1.00, 0.46], dtype=np.float32),
+        )
+        socket_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
+        peg_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
+        return {
+            "socket_pos": socket_position.astype(np.float32, copy=False),
+            "socket_quat": socket_quat,
+            "peg_pos": peg_position.astype(np.float32, copy=False),
+            "peg_quat": peg_quat,
+        }

roboimi/utils/constants.py

@@ -20,7 +20,14 @@ SIM_TASK_CONFIGS = {
         'dataset_dir': DATASET_DIR + '/sim_transfer',
         'num_episodes': 20,
         'episode_len': 700,
-        'camera_names': ['top','r_vis','front'],
+        'camera_names': ['r_vis', 'top', 'front'],
+        'xml_dir': HOME_PATH + '/assets'
+    },
+    'sim_air_insert_socket_peg': {
+        'dataset_dir': DATASET_DIR + '/sim_air_insert_socket_peg',
+        'num_episodes': 20,
+        'episode_len': 750,
+        'camera_names': ['l_vis', 'r_vis', 'front'],
         'xml_dir': HOME_PATH + '/assets'
     },
 
@@ -52,13 +59,3 @@ PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) /
 MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
 PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
 MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x))
-
-MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
-PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
-
-MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
-MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE))
-PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
-PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE))
-
-MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2

roboimi/vla/agent_smolvla_conditioned.py

@@ -0,0 +1,277 @@
+from __future__ import annotations
+
+from collections import deque
+from typing import Dict, Optional, Sequence
+
+import torch
+import torch.nn as nn
+
+from roboimi.vla.agent_imf import IMFVLAAgent
+from roboimi.vla.models.normalization import NormalizationModule
+
+
+class SmolVLAIMFAttnResAgent(IMFVLAAgent):
+    """IMF-AttnRes action expert conditioned by SmolVLA-style VLM prefix tokens.
+
+    Unlike ``VLAAgent`` this agent does not concatenate ResNet features and raw
+    state at every observation step.  A ``condition_encoder`` receives images,
+    normalized state, and variable task language, and returns a condition token
+    sequence `(B, S, D)` that is passed directly to the IMF head.
+    """
+
+    def __init__(
+        self,
+        condition_encoder: nn.Module,
+        action_encoder,
+        head,
+        action_dim: int,
+        obs_dim: int,
+        pred_horizon: int = 16,
+        obs_horizon: int = 2,
+        diffusion_steps: int = 100,
+        inference_steps: int = 1,
+        num_cams: int = 3,
+        camera_names: Optional[Sequence[str]] = None,
+        dataset_stats=None,
+        normalization_type: str = 'min_max',
+        num_action_steps: int = 8,
+        head_type: str = 'transformer',
+        condition_dim: int | None = None,
+        condition_sequence_length: int | None = None,
+        task_description: str | None = None,
+        **kwargs,
+    ) -> None:
+        # Intentionally bypass VLAAgent.__init__; its condition dimensions are
+        # tied to ResNet-style visual+state concatenation.
+        nn.Module.__init__(self)
+        if inference_steps != 1:
+            raise ValueError(
+                'SmolVLAIMFAttnResAgent only supports one-step IMF inference; '
+                f'inference_steps must be 1, got {inference_steps}.'
+            )
+        if head_type != 'transformer':
+            raise ValueError(f'SmolVLAIMFAttnResAgent requires head_type="transformer", got {head_type!r}')
+        del diffusion_steps, kwargs
+
+        self.action_dim = int(action_dim)
+        self.obs_dim = int(obs_dim)
+        self.pred_horizon = int(pred_horizon)
+        self.obs_horizon = int(obs_horizon)
+        self.num_cams = int(num_cams)
+        self.num_action_steps = int(num_action_steps)
+        self.inference_steps = 1
+        self.head_type = head_type
+        self.camera_names = tuple(camera_names) if camera_names is not None else None
+        if self.camera_names is not None and len(self.camera_names) != self.num_cams:
+            raise ValueError(f'camera_names length({len(self.camera_names)}) does not match num_cams({self.num_cams})')
+
+        self.normalization = NormalizationModule(stats=dataset_stats, normalization_type=normalization_type)
+        self.condition_encoder = condition_encoder
+        # Compatibility aliases used by some utilities/tests.
+        self.vision_encoder = condition_encoder
+        self.action_encoder = action_encoder
+        self.state_encoder = None
+        self.task_description = task_description
+
+        encoder_dim = getattr(condition_encoder, 'output_dim', None)
+        if encoder_dim is None:
+            encoder_dim = getattr(condition_encoder, 'joint_output_dim', None)
+        if condition_dim is None:
+            if encoder_dim is None:
+                raise ValueError('condition_dim must be provided when condition_encoder has no output_dim')
+            condition_dim = int(encoder_dim)
+        self.per_step_cond_dim = int(condition_dim)
+        self.raw_per_step_cond_dim = self.per_step_cond_dim
+
+        encoder_seq_len = getattr(condition_encoder, 'condition_sequence_length', None)
+        if condition_sequence_length is None:
+            if encoder_seq_len is None:
+                raise ValueError(
+                    'condition_sequence_length must be provided when condition_encoder has no condition_sequence_length'
+                )
+            condition_sequence_length = int(encoder_seq_len)
+        self.condition_sequence_length = int(condition_sequence_length)
+        self.condition_tokens_per_step = self.condition_sequence_length
+        self.global_cond_dim = self.per_step_cond_dim * self.condition_sequence_length
+
+        if isinstance(head, nn.Module):
+            self.noise_pred_net = head
+        else:
+            self.noise_pred_net = head(
+                input_dim=self.action_dim,
+                output_dim=self.action_dim,
+                horizon=self.pred_horizon,
+                n_obs_steps=self.condition_sequence_length,
+                cond_dim=self.per_step_cond_dim,
+            )
+        self.reset()
+
+    def _get_model_device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def _move_to_device(self, data, device: torch.device):
+        if torch.is_tensor(data):
+            return data.to(device)
+        if isinstance(data, dict):
+            return {k: self._move_to_device(v, device) for k, v in data.items()}
+        if isinstance(data, list):
+            return [self._move_to_device(v, device) for v in data]
+        if isinstance(data, tuple):
+            return tuple(self._move_to_device(v, device) for v in data)
+        return data
+
+    def _order_images(self, images: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        if self.camera_names is None:
+            names = tuple(sorted(images.keys()))
+            if len(names) != self.num_cams:
+                raise ValueError(f'image camera count({len(names)}) does not match num_cams({self.num_cams})')
+            return {name: images[name] for name in names}
+        missing = [name for name in self.camera_names if name not in images]
+        if missing:
+            raise ValueError(f'image condition missing required cameras. missing={missing}, expected={list(self.camera_names)}')
+        return {name: images[name] for name in self.camera_names}
+
+    def _resolve_task(self, task, batch_size: int):
+        def is_missing_task(item) -> bool:
+            return item is None or (isinstance(item, str) and item.strip().lower() in {'', 'unknown'})
+
+        def fallback_or(item):
+            if self.task_description is not None and is_missing_task(item):
+                return self.task_description
+            if item is None:
+                return ''
+            return item
+
+        if task is None:
+            task = self.task_description
+        if task is None:
+            return [''] * batch_size
+        if isinstance(task, str):
+            return [fallback_or(task)] * batch_size
+        task = list(task)
+        if len(task) != batch_size:
+            raise ValueError(f'task batch size ({len(task)}) must match batch size ({batch_size})')
+        return [fallback_or(item) for item in task]
+
+    def _build_cond(self, images: Dict[str, torch.Tensor], states: torch.Tensor, task=None) -> torch.Tensor:
+        ordered_images = self._order_images(images)
+        batch_size = states.shape[0]
+        tasks = self._resolve_task(task, batch_size=batch_size)
+        cond = self.condition_encoder(ordered_images, state=states, task=tasks)
+        if cond.ndim != 3:
+            raise RuntimeError(f'condition_encoder must return (B,S,D), got {tuple(cond.shape)}')
+        if cond.shape[0] != batch_size:
+            raise RuntimeError(f'condition batch mismatch: got {cond.shape[0]}, expected {batch_size}')
+        if cond.shape[1] != self.condition_sequence_length:
+            raise RuntimeError(
+                f'condition sequence length mismatch: got {cond.shape[1]}, expected {self.condition_sequence_length}'
+            )
+        if cond.shape[2] != self.per_step_cond_dim:
+            raise RuntimeError(f'condition dim mismatch: got {cond.shape[2]}, expected {self.per_step_cond_dim}')
+        head_dtype = next(self.noise_pred_net.parameters()).dtype
+        return cond.to(dtype=head_dtype)
+
+    def compute_loss(self, batch):
+        actions, states, images = batch['action'], batch['qpos'], batch['images']
+        action_is_pad = batch.get('action_is_pad', None)
+        batch_size = actions.shape[0]
+
+        states = self.normalization.normalize_qpos(states)
+        actions = self.normalization.normalize_action(actions)
+        cond = self._build_cond(images, states, task=batch.get('task', None))
+
+        x = actions
+        e = torch.randn_like(x)
+        t = torch.rand(batch_size, device=x.device, dtype=x.dtype)
+        r = torch.rand(batch_size, device=x.device, dtype=x.dtype)
+        t, r = torch.maximum(t, r), torch.minimum(t, r)
+
+        t_broadcast = self._broadcast_batch_time(t, x)
+        z_t = (1 - t_broadcast) * x + t_broadcast * e
+
+        v = self.fn(z_t, t, t, cond=cond)
+        u, du_dt = self._compute_u_and_du_dt(z_t, r, t, cond=cond, v=v)
+        V = self._compound_velocity(u, du_dt, r, t)
+        target = e - x
+
+        loss = nn.functional.mse_loss(V, target, reduction='none')
+        if action_is_pad is not None:
+            mask = (~action_is_pad).unsqueeze(-1).to(loss.dtype)
+            valid_count = mask.sum() * loss.shape[-1]
+            loss = (loss * mask).sum() / valid_count.clamp_min(1.0)
+        else:
+            loss = loss.mean()
+        return loss
+
+    @torch.no_grad()
+    def predict_action(self, images, proprioception, task=None):
+        batch_size = proprioception.shape[0]
+        proprioception = self.normalization.normalize_qpos(proprioception)
+        cond = self._build_cond(images, proprioception, task=task)
+        z_t = torch.randn((batch_size, self.pred_horizon, self.action_dim), device=cond.device, dtype=cond.dtype)
+        action = self._sample_one_step(z_t, cond=cond)
+        return self.normalization.denormalize_action(action)
+
+    @torch.no_grad()
+    def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
+        return self.predict_action(batch['images'], batch['qpos'], task=batch.get('task', None))
+
+    def reset(self):
+        self._queues = {
+            'qpos': deque(maxlen=self.obs_horizon),
+            'images': deque(maxlen=self.obs_horizon),
+            'task': deque(maxlen=self.obs_horizon),
+            'action': deque(maxlen=self.pred_horizon - self.obs_horizon + 1),
+        }
+
+    def _populate_queues(self, observation: Dict[str, torch.Tensor]) -> None:
+        if 'qpos' in observation:
+            self._queues['qpos'].append(observation['qpos'].clone())
+        if 'images' in observation:
+            ordered_images = self._order_images(observation['images'])
+            self._queues['images'].append({k: v.clone() for k, v in ordered_images.items()})
+        if 'task' in observation:
+            self._queues['task'].append(observation['task'])
+
+    def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]:
+        qpos_list = list(self._queues['qpos'])
+        if not qpos_list:
+            raise ValueError('observation queue is empty.')
+        while len(qpos_list) < self.obs_horizon:
+            qpos_list.append(qpos_list[-1])
+        batch_qpos = torch.stack(qpos_list, dim=0).unsqueeze(0)
+
+        images_list = list(self._queues['images'])
+        if not images_list:
+            raise ValueError('image queue is empty.')
+        while len(images_list) < self.obs_horizon:
+            images_list.append(images_list[-1])
+        names = self.camera_names if self.camera_names is not None else tuple(sorted(images_list[0].keys()))
+        batch_images = {
+            name: torch.stack([item[name] for item in images_list], dim=0).unsqueeze(0)
+            for name in names
+        }
+        batch = {'qpos': batch_qpos, 'images': batch_images}
+        if self._queues['task']:
+            batch['task'] = [list(self._queues['task'])[-1]]
+        elif self.task_description is not None:
+            batch['task'] = [self.task_description]
+        return batch
+
+    @torch.no_grad()
+    def select_action(self, observation: Dict[str, torch.Tensor]) -> torch.Tensor:
+        device = self._get_model_device()
+        observation = self._move_to_device(observation, device)
+        self._populate_queues(observation)
+        if len(self._queues['action']) == 0:
+            batch = self._prepare_observation_batch()
+            actions = self.predict_action_chunk(batch)
+            start = self.obs_horizon - 1
+            end = start + self.num_action_steps
+            executable_actions = actions[:, start:end]
+            for i in range(executable_actions.shape[1]):
+                self._queues['action'].append(executable_actions[:, i].squeeze(0))
+        return self._queues['action'].popleft()
+
+    def get_normalization_stats(self):
+        return self.normalization.get_stats()

roboimi/vla/agent_smolvla_native.py

@@ -0,0 +1,398 @@
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import fields
+from typing import Dict, Optional, Sequence
+
+import torch
+import torch.nn as nn
+
+from roboimi.vla.models.normalization import NormalizationModule
+
+
+class SmolVLANativeAgent(nn.Module):
+    """RoboIMI wrapper for the native SmolVLA flow-matching model.
+
+    The wrapper owns RoboIMI-facing concerns only: dataset normalization,
+    camera ordering, language fallback/tokenization, rollout queues, and action
+    denormalization.  The native SmolVLA model itself is imported lazily so unit
+    tests can inject fakes without downloading or loading a real VLM.
+    """
+
+    def __init__(
+        self,
+        model: Optional[nn.Module] = None,
+        tokenizer=None,
+        action_dim: int = 16,
+        obs_dim: int = 16,
+        chunk_size: int = 16,
+        n_action_steps: int = 8,
+        obs_horizon: int = 1,
+        action_chunk_start: int = 0,
+        num_cams: int = 3,
+        camera_names: Optional[Sequence[str]] = None,
+        dataset_stats=None,
+        normalization_type: str = 'min_max',
+        task_description: Optional[str] = None,
+        model_config: Optional[dict] = None,
+        tokenizer_name: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        del kwargs
+        self.action_dim = int(action_dim)
+        self.obs_dim = int(obs_dim)
+        self.chunk_size = int(chunk_size)
+        self.pred_horizon = self.chunk_size
+        self.n_action_steps = int(n_action_steps)
+        self.num_action_steps = self.n_action_steps
+        self.obs_horizon = int(obs_horizon)
+        self.action_chunk_start = int(action_chunk_start)
+        self.num_cams = int(num_cams)
+        self.camera_names = tuple(camera_names) if camera_names is not None else None
+        if self.camera_names is not None and len(self.camera_names) != self.num_cams:
+            raise ValueError(f'camera_names length({len(self.camera_names)}) does not match num_cams({self.num_cams})')
+        if self.n_action_steps < 1:
+            raise ValueError('n_action_steps must be >= 1')
+        if self.action_chunk_start < 0:
+            raise ValueError('action_chunk_start must be >= 0')
+        if self.action_chunk_start + self.n_action_steps > self.chunk_size:
+            raise ValueError('action_chunk_start + n_action_steps must be <= chunk_size')
+
+        self.normalization = NormalizationModule(stats=dataset_stats, normalization_type=normalization_type)
+        self.task_description = task_description
+        self.model_config = dict(model_config or {})
+        self.max_state_dim = int(self.model_config.get('max_state_dim', self.obs_dim))
+        self.max_action_dim = int(self.model_config.get('max_action_dim', self.action_dim))
+        self.resize_imgs_with_padding = self._normalize_resize_shape(
+            self.model_config.get('resize_imgs_with_padding', self.model_config.get('image_resize_shape', (512, 512)))
+        )
+        self.tokenizer_max_length = int(self.model_config.get('tokenizer_max_length', 48))
+        self.pad_language_to = str(self.model_config.get('pad_language_to', 'longest'))
+        self.tokenizer_name = tokenizer_name
+        self.tokenizer = tokenizer if tokenizer is not None else self._build_tokenizer(tokenizer_name)
+        self.model = model if model is not None else self._build_model()
+        self.reset()
+
+    @staticmethod
+    def _normalize_resize_shape(shape):
+        if shape is None:
+            return None
+        normalized = tuple(int(v) for v in shape)
+        if len(normalized) != 2:
+            raise ValueError(f'resize_imgs_with_padding must contain exactly two values, got {normalized}')
+        return normalized
+
+    def _build_tokenizer(self, tokenizer_name):
+        if tokenizer_name is None:
+            tokenizer_name = self.model_config.get('tokenizer_name') or self.model_config.get('vlm_model_name')
+        if tokenizer_name is None:
+            raise ImportError('A tokenizer or tokenizer_name/model_config.vlm_model_name is required for SmolVLANativeAgent')
+        try:
+            from transformers import AutoTokenizer
+        except ImportError as exc:
+            raise ImportError('SmolVLANativeAgent requires transformers to construct the real tokenizer') from exc
+        return AutoTokenizer.from_pretrained(tokenizer_name)
+
+    def _native_config_kwargs(self) -> dict:
+        from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
+
+        cfg_kwargs = dict(self.model_config)
+
+        # Backwards-compatible aliases from earlier RoboIMI-facing drafts.  The
+        # native model core intentionally keeps only SmolVLA model fields.
+        if 'image_resize_shape' in cfg_kwargs and 'resize_imgs_with_padding' not in cfg_kwargs:
+            cfg_kwargs['resize_imgs_with_padding'] = cfg_kwargs['image_resize_shape']
+        if 'freeze_vlm' in cfg_kwargs and 'train_expert_only' not in cfg_kwargs:
+            cfg_kwargs['train_expert_only'] = bool(cfg_kwargs['freeze_vlm'])
+
+        for wrapper_only_key in (
+            'state_dim',
+            'action_dim',
+            'tokenizer_name',
+            'freeze_vlm',
+            'image_resize_shape',
+            'num_cameras',
+        ):
+            cfg_kwargs.pop(wrapper_only_key, None)
+
+        cfg_kwargs.setdefault('max_state_dim', self.max_state_dim)
+        cfg_kwargs.setdefault('max_action_dim', self.max_action_dim)
+        cfg_kwargs.setdefault('chunk_size', self.chunk_size)
+        cfg_kwargs.setdefault('n_action_steps', self.n_action_steps)
+        cfg_kwargs['resize_imgs_with_padding'] = self._normalize_resize_shape(
+            cfg_kwargs.get('resize_imgs_with_padding', self.resize_imgs_with_padding)
+        )
+
+        allowed_keys = {field.name for field in fields(NativeSmolVLAConfig)}
+        return {key: value for key, value in cfg_kwargs.items() if key in allowed_keys}
+
+    def _build_model(self):
+        try:
+            from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
+            from roboimi.vla.models.smolvla.modeling import VLAFlowMatching
+        except ImportError as exc:
+            raise ImportError(
+                'Native SmolVLA model modules are unavailable. Pass injected model/tokenizer fakes in tests, '
+                'or add roboimi.vla.models.smolvla.configuration/modeling for real construction.'
+            ) from exc
+        config = NativeSmolVLAConfig(**self._native_config_kwargs())
+        return VLAFlowMatching(config=config)
+
+    def _get_model_device(self) -> torch.device:
+        try:
+            return next(self.model.parameters()).device
+        except (StopIteration, AttributeError):
+            return torch.device('cpu')
+
+    def _move_to_device(self, data, device: torch.device):
+        if torch.is_tensor(data):
+            return data.to(device)
+        if isinstance(data, dict):
+            return {k: self._move_to_device(v, device) for k, v in data.items()}
+        if isinstance(data, list):
+            return [self._move_to_device(v, device) for v in data]
+        if isinstance(data, tuple):
+            return tuple(self._move_to_device(v, device) for v in data)
+        return data
+
+    def _order_images(self, images: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        if self.camera_names is None:
+            names = tuple(sorted(images.keys()))
+            if len(names) != self.num_cams:
+                raise ValueError(f'image camera count({len(names)}) does not match num_cams({self.num_cams})')
+            return {name: images[name] for name in names}
+        missing = [name for name in self.camera_names if name not in images]
+        if missing:
+            raise ValueError(f'image batch missing required cameras. missing={missing}, expected={list(self.camera_names)}')
+        return {name: images[name] for name in self.camera_names}
+
+    @staticmethod
+    def _is_missing_task(item) -> bool:
+        return item is None or (isinstance(item, str) and item.strip().lower() in {'', 'unknown'})
+
+    def _resolve_task(self, task, batch_size: int):
+        def fallback_or(item):
+            if self._is_missing_task(item):
+                if self.task_description is not None:
+                    return self.task_description
+                if item is None:
+                    return ''
+            return item
+
+        if task is None:
+            task = self.task_description
+        if task is None:
+            return [''] * batch_size
+        if isinstance(task, str):
+            return [fallback_or(task)] * batch_size
+        task = list(task)
+        if len(task) != batch_size:
+            raise ValueError(f'task batch size ({len(task)}) must match batch size ({batch_size})')
+        return [fallback_or(item) for item in task]
+
+    def _tokenize_tasks(self, task, batch_size: int, device: torch.device) -> dict:
+        tasks = []
+        for text in self._resolve_task(task, batch_size):
+            text = str(text)
+            tasks.append(text if text.endswith('\n') else f'{text}\n')
+        old_padding_side = getattr(self.tokenizer, 'padding_side', None)
+        if old_padding_side is not None:
+            self.tokenizer.padding_side = 'right'
+        try:
+            tokenized = self.tokenizer(
+                tasks,
+                padding=self.pad_language_to,
+                max_length=self.tokenizer_max_length,
+                truncation=True,
+                return_tensors='pt',
+            )
+        finally:
+            if old_padding_side is not None:
+                self.tokenizer.padding_side = old_padding_side
+        return {
+            'lang_tokens': tokenized['input_ids'].to(device=device),
+            'lang_masks': tokenized['attention_mask'].to(device=device, dtype=torch.bool),
+        }
+
+    @staticmethod
+    def _pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor:
+        current_dim = vector.shape[-1]
+        if current_dim == new_dim:
+            return vector
+        if current_dim > new_dim:
+            raise ValueError(f'cannot pad vector with dim {current_dim} to smaller dim {new_dim}')
+        padded_shape = list(vector.shape)
+        padded_shape[-1] = int(new_dim)
+        padded = torch.zeros(*padded_shape, dtype=vector.dtype, device=vector.device)
+        padded[..., :current_dim] = vector
+        return padded
+
+    @staticmethod
+    def _resize_with_pad(img: torch.Tensor, width: int, height: int, pad_value: float = 0.0) -> torch.Tensor:
+        if img.ndim != 4:
+            raise ValueError(f'expected image tensor shaped (B,C,H,W), got {tuple(img.shape)}')
+        import torch.nn.functional as F
+
+        cur_height, cur_width = img.shape[2:]
+        ratio = max(cur_width / width, cur_height / height)
+        resized_height = int(cur_height / ratio)
+        resized_width = int(cur_width / ratio)
+        resized = F.interpolate(img, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
+        pad_height = max(0, int(height - resized_height))
+        pad_width = max(0, int(width - resized_width))
+        return F.pad(resized, (pad_width, 0, pad_height, 0), value=pad_value)
+
+    def _prepare_native_images(self, images: Dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+        ordered_images = self._order_images(images)
+        prepared_images: list[torch.Tensor] = []
+        img_masks: list[torch.Tensor] = []
+        reference_batch_size: int | None = None
+        for camera_name, image in ordered_images.items():
+            if image.ndim == 5:
+                image = image[:, -1]
+            elif image.ndim != 4:
+                raise ValueError(
+                    f'image for camera {camera_name!r} must be shaped (B,T,C,H,W) or (B,C,H,W), got {tuple(image.shape)}'
+                )
+            if reference_batch_size is None:
+                reference_batch_size = int(image.shape[0])
+            elif int(image.shape[0]) != reference_batch_size:
+                raise ValueError(f'image batch size mismatch for camera {camera_name!r}')
+
+            image = image.contiguous().float().clamp(0.0, 1.0)
+            if self.resize_imgs_with_padding is not None:
+                image = self._resize_with_pad(image, *self.resize_imgs_with_padding, pad_value=0.0)
+            image = image * 2.0 - 1.0
+            prepared_images.append(image)
+            img_masks.append(torch.ones(image.shape[0], dtype=torch.bool, device=image.device))
+        return prepared_images, img_masks
+
+    def _prepare_native_state(self, states: torch.Tensor) -> torch.Tensor:
+        states = self.normalization.normalize_qpos(states)
+        if states.ndim > 2:
+            states = states[:, -1, :]
+        return self._pad_vector(states.float(), self.max_state_dim)
+
+    def _prepare_native_actions(self, actions: torch.Tensor) -> torch.Tensor:
+        actions = self.normalization.normalize_action(actions)
+        return self._pad_vector(actions.float(), self.max_action_dim)
+
+    def _prepare_model_inputs(self, batch: Dict[str, torch.Tensor], include_actions: bool = False) -> dict:
+        states = batch['qpos']
+        batch_size = states.shape[0]
+        device = states.device
+        images, img_masks = self._prepare_native_images(batch['images'])
+        inputs = {
+            'images': images,
+            'img_masks': img_masks,
+            'state': self._prepare_native_state(states),
+        }
+        inputs.update(self._tokenize_tasks(batch.get('task', None), batch_size, device))
+        if include_actions:
+            inputs['actions'] = self._prepare_native_actions(batch['action'])
+        return inputs
+
+    def _reduce_native_losses(self, losses: torch.Tensor, action_is_pad: torch.Tensor | None = None) -> torch.Tensor:
+        if losses.ndim == 0:
+            return losses
+        if losses.shape[-1] < self.action_dim:
+            raise RuntimeError(f'loss action dim mismatch: got {losses.shape[-1]}, expected at least {self.action_dim}')
+        losses = losses[..., : self.action_dim]
+        if action_is_pad is None:
+            return losses.mean()
+        if losses.ndim != 3:
+            raise RuntimeError(f'action padding mask requires per-element losses shaped (B,H,A), got {tuple(losses.shape)}')
+        if tuple(action_is_pad.shape) != tuple(losses.shape[:2]):
+            raise RuntimeError(
+                f'action_is_pad shape {tuple(action_is_pad.shape)} does not match loss prefix {tuple(losses.shape[:2])}'
+            )
+        mask = (~action_is_pad).to(device=losses.device, dtype=losses.dtype).unsqueeze(-1)
+        denom = (mask.sum() * losses.shape[-1]).clamp_min(1.0)
+        return (losses * mask).sum() / denom
+
+    def compute_loss(self, batch):
+        inputs = self._prepare_model_inputs(batch, include_actions=True)
+        action_is_pad = batch.get('action_is_pad', None)
+        output = self.model(**inputs)
+        if isinstance(output, dict):
+            if 'loss' not in output:
+                raise RuntimeError('SmolVLA model forward returned a dict without loss')
+            return output['loss']
+        if torch.is_tensor(output):
+            return self._reduce_native_losses(output, action_is_pad=action_is_pad)
+        if hasattr(output, 'loss'):
+            return output.loss
+        raise RuntimeError(f'Unsupported SmolVLA forward output type: {type(output)!r}')
+
+    @torch.no_grad()
+    def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
+        inputs = self._prepare_model_inputs(batch, include_actions=False)
+        if not hasattr(self.model, 'sample_actions'):
+            raise RuntimeError('SmolVLA model must implement sample_actions for inference')
+        actions = self.model.sample_actions(**inputs)
+        if not torch.is_tensor(actions) or actions.ndim != 3:
+            raise RuntimeError(f'sample_actions must return (B,H,A), got {type(actions)!r} {getattr(actions, "shape", None)}')
+        if actions.shape[-1] < self.action_dim:
+            raise RuntimeError(f'action dim mismatch: got {actions.shape[-1]}, expected at least {self.action_dim}')
+        actions = actions[..., : self.action_dim]
+        return self.normalization.denormalize_action(actions)
+
+    def reset(self):
+        self._queues = {
+            'qpos': deque(maxlen=self.obs_horizon),
+            'images': deque(maxlen=self.obs_horizon),
+            'task': deque(maxlen=self.obs_horizon),
+            'action': deque(maxlen=self.n_action_steps),
+        }
+
+    def _populate_queues(self, observation: Dict[str, torch.Tensor]) -> None:
+        if 'qpos' in observation:
+            self._queues['qpos'].append(observation['qpos'].clone())
+        if 'images' in observation:
+            ordered_images = self._order_images(observation['images'])
+            self._queues['images'].append({k: v.clone() for k, v in ordered_images.items()})
+        if 'task' in observation:
+            self._queues['task'].append(observation['task'])
+
+    def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]:
+        qpos_list = list(self._queues['qpos'])
+        if not qpos_list:
+            raise ValueError('observation queue is empty.')
+        while len(qpos_list) < self.obs_horizon:
+            qpos_list.append(qpos_list[-1])
+        batch_qpos = torch.stack(qpos_list, dim=0).unsqueeze(0)
+
+        images_list = list(self._queues['images'])
+        if not images_list:
+            raise ValueError('image queue is empty.')
+        while len(images_list) < self.obs_horizon:
+            images_list.append(images_list[-1])
+        names = self.camera_names if self.camera_names is not None else tuple(sorted(images_list[0].keys()))
+        batch_images = {name: torch.stack([item[name] for item in images_list], dim=0).unsqueeze(0) for name in names}
+        batch = {'qpos': batch_qpos, 'images': batch_images}
+        if self._queues['task']:
+            batch['task'] = [list(self._queues['task'])[-1]]
+        elif self.task_description is not None:
+            batch['task'] = [self.task_description]
+        return batch
+
+    @torch.no_grad()
+    def select_action(self, observation: Dict[str, torch.Tensor]) -> torch.Tensor:
+        device = self._get_model_device()
+        observation = self._move_to_device(observation, device)
+        self._populate_queues(observation)
+        if len(self._queues['action']) == 0:
+            batch = self._prepare_observation_batch()
+            actions = self.predict_action_chunk(batch)
+            start = self.action_chunk_start
+            end = start + self.n_action_steps
+            if actions.shape[1] < end:
+                raise RuntimeError(f'action chunk too short: got {actions.shape[1]}, need at least {end}')
+            executable_actions = actions[:, start:end]
+            for i in range(executable_actions.shape[1]):
+                self._queues['action'].append(executable_actions[:, i].squeeze(0))
+        return self._queues['action'].popleft()
+
+    def get_normalization_stats(self):
+        return self.normalization.get_stats()

roboimi/vla/conf/agent/smolvla_imf_attnres.yaml

@@ -0,0 +1,46 @@
+# @package agent
+defaults:
+  - /backbone@condition_encoder: smolvla_prefix_encoder
+  - /modules@action_encoder: identity_action_encoder
+  - /head: imf_transformer1d
+  - _self_
+
+_target_: roboimi.vla.agent_smolvla_conditioned.SmolVLAIMFAttnResAgent
+
+action_dim: 16
+obs_dim: 16
+normalization_type: "min_max"
+pred_horizon: 16
+obs_horizon: 2
+num_action_steps: 8
+camera_names: ${data.camera_names}
+num_cams: ${len:${agent.camera_names}}
+
+# Optional fallback language instruction used only when a batch/observation
+# does not provide `task`. Keep null by default so this architecture remains
+# generic and dataset/eval can supply variable language.
+task_description: null
+condition_dim: 960
+condition_sequence_length: 241
+
+condition_encoder:
+  num_cameras: ${agent.num_cams}
+  camera_names: ${agent.camera_names}
+
+diffusion_steps: 100
+inference_steps: 1
+head_type: "transformer"
+
+head:
+  input_dim: ${agent.action_dim}
+  output_dim: ${agent.action_dim}
+  horizon: ${agent.pred_horizon}
+  n_obs_steps: ${agent.condition_sequence_length}
+  cond_dim: ${agent.condition_dim}
+  causal_attn: false
+  time_as_cond: true
+  obs_as_cond: true
+  n_cond_layers: 0
+  backbone_type: attnres_full
+  n_head: 1
+  n_kv_head: 1

roboimi/vla/conf/agent/smolvla_native.yaml

@@ -0,0 +1,35 @@
+# @package agent
+_target_: roboimi.vla.agent_smolvla_native.SmolVLANativeAgent
+
+model: null
+tokenizer: null
+tokenizer_name: ${agent.model_config.vlm_model_name}
+
+action_dim: 16
+obs_dim: 16
+normalization_type: "gaussian"
+chunk_size: 32
+pred_horizon: ${agent.chunk_size}
+obs_horizon: 2
+n_action_steps: 16
+num_action_steps: ${agent.n_action_steps}
+action_chunk_start: 0
+camera_names: ${data.camera_names}
+num_cams: ${len:${agent.camera_names}}
+task_description: null
+# SmolVLA performs its own SigLIP-style resize+pad inside the wrapper/core.
+# Keeping dataset/eval resize disabled avoids lossy double resizing.
+dataset_image_resize_shape: null
+eval_image_resize_shape: null
+
+model_config:
+  max_state_dim: 32
+  max_action_dim: 32
+  chunk_size: ${agent.chunk_size}
+  n_action_steps: ${agent.n_action_steps}
+  vlm_model_name: HuggingFaceTB/SmolVLM2-500M-Video-Instruct
+  load_vlm_weights: true
+  train_expert_only: true
+  freeze_vision_encoder: true
+  num_vlm_layers: 16
+  resize_imgs_with_padding: [512, 512]

roboimi/vla/conf/backbone/smolvla_prefix_encoder.yaml

@@ -0,0 +1,19 @@
+_target_: roboimi.vla.models.backbones.smolvla_prefix_encoder.SmolVLAPrefixEncoder
+
+model_name: HuggingFaceTB/SmolVLM2-500M-Video-Instruct
+load_vlm_weights: true
+local_files_only: false
+num_vlm_layers: 16
+freeze_vlm: true
+freeze_vision_encoder: true
+train_state_proj: true
+max_state_dim: 32
+resize_imgs_with_padding: [512, 512]
+tokenizer_max_length: 48
+pad_language_to: max_length
+run_text_model: true
+camera_names: [r_vis, top, front]
+num_cameras: 3
+
+dataset_image_resize_shape: null
+eval_image_resize_shape: null

roboimi/vla/conf/config.yaml

@@ -29,6 +29,11 @@ train:
   rollout_val_freq_epochs: 50  # 每隔多少个 epoch 执行一次 rollout 验证
   rollout_validate_on_checkpoint: false  # 是否在保存 checkpoint 后立即运行 rollout 验证
   rollout_num_episodes: 3  # rollout 验证的回合数
+  rollout_device: ${train.device}  # rollout 使用的设备；默认跟随训练设备
+  rollout_num_workers: null  # rollout 并行 worker 数；null 时 CUDA 自动推断，CPU 保持 1
+  rollout_cuda_devices: null  # rollout CUDA 并行使用的逻辑 device 列表；null 时默认 [0]
+  rollout_response_timeout_s: 300.0  # rollout worker 等待 inference server 响应的超时时间
+  rollout_server_startup_timeout_s: 300.0  # rollout 等待 inference server 就绪的超时时间
 
   # 学习率调度器（带预热）
   warmup_steps: 2000  # 预热步数（Transformer建议更长）

roboimi/vla/conf/eval/eval.yaml

@@ -2,9 +2,14 @@
 # 评估配置
 ckpt_path: "checkpoints/vla_model_best.pt"  # 模型检查点路径
 num_episodes: 3           # 评估回合数
+num_workers: 1            # 并行 worker 数；1 表示保持单进程评估
+cuda_devices: null        # CUDA 并行评估时使用的逻辑设备列表；null 表示默认 [0]
+response_timeout_s: 300.0 # worker 等待 inference server 响应的超时时间（秒）
+server_startup_timeout_s: 300.0 # parent 等待 inference server 就绪的超时时间（秒）
 max_timesteps: 700        # 每回合最大时间步
 device: ${train.device}   # 与训练保持一致
 task_name: "sim_transfer" # 环境任务名称
+task_description: null    # 可选语言指令；环境 obs 没有 task 时注入给语言条件策略
 
 # ====================
 # 策略执行参数

roboimi/vla/data/simpe_robot_dataset.py

@@ -4,6 +4,7 @@ from torch.utils.data import Dataset
 from typing import List, Dict, Union, Optional, Sequence
 from pathlib import Path
 from collections import OrderedDict
+import re
 
 
 class SimpleRobotDataset(Dataset):
@@ -24,6 +25,7 @@ class SimpleRobotDataset(Dataset):
         camera_names: List[str] = None,
         image_resize_shape: Optional[Sequence[int]] = (224, 224),
         max_open_files: int = 64,
+        episode_indices: Optional[Sequence[int]] = None,
     ):
         """
         Args:
@@ -33,6 +35,7 @@ class SimpleRobotDataset(Dataset):
             camera_names: 相机名称列表，如 ["r_vis", "top", "front"]
             image_resize_shape: 图像缩放尺寸 (W, H)；为 None 时保留原始分辨率
             max_open_files: 每个 worker 最多缓存的 HDF5 文件句柄数
+            episode_indices: 可选的原始 episode 编号子集
 
         HDF5 文件格式:
             - action: [T, action_dim]
@@ -48,6 +51,9 @@ class SimpleRobotDataset(Dataset):
         )
         self.max_open_files = max(1, int(max_open_files))
         self._file_cache: "OrderedDict[str, h5py.File]" = OrderedDict()
+        self.requested_episode_indices = (
+            None if episode_indices is None else tuple(sorted(int(idx) for idx in episode_indices))
+        )
 
         self.dataset_dir = Path(dataset_dir)
         if not self.dataset_dir.exists():
@@ -59,6 +65,18 @@ class SimpleRobotDataset(Dataset):
             self.hdf5_files = sorted(self.dataset_dir.glob("episode_*.hdf5"))
         if not self.hdf5_files:
             raise FileNotFoundError(f"在 {dataset_dir} 中未找到 HDF5 文件")
+        if self.requested_episode_indices is not None:
+            requested = set(self.requested_episode_indices)
+            filtered = []
+            for hdf5_path in self.hdf5_files:
+                match = re.search(r'episode_(\d+)$', hdf5_path.stem)
+                if match and int(match.group(1)) in requested:
+                    filtered.append(hdf5_path)
+            self.hdf5_files = filtered
+            if not self.hdf5_files:
+                raise FileNotFoundError(
+                    f"在 {dataset_dir} 中未找到 episode_indices={sorted(requested)} 对应的 HDF5 文件"
+                )
 
         # 构建 episode 索引（只存储元数据，不加载数据）
         self.episodes = {}
@@ -66,14 +84,18 @@ class SimpleRobotDataset(Dataset):
         for ep_idx, hdf5_path in enumerate(self.hdf5_files):
             with h5py.File(hdf5_path, 'r') as f:
                 T = f['action'].shape[0]
+                dataset_episode_idx = ep_idx
+                match = re.search(r'episode_(\d+)$', hdf5_path.stem)
+                if match:
+                    dataset_episode_idx = int(match.group(1))
                 start_idx = len(self.frame_meta)
                 for t in range(T):
                     self.frame_meta.append({
-                        "ep_idx": ep_idx,
+                        "ep_idx": dataset_episode_idx,
                         "frame_idx": t,
                         "hdf5_path": hdf5_path,
                     })
-                self.episodes[ep_idx] = list(range(start_idx, len(self.frame_meta)))
+                self.episodes[dataset_episode_idx] = list(range(start_idx, len(self.frame_meta)))
 
         print(f"懒加载模式: {len(self.hdf5_files)} 个 episodes, 共 {len(self.frame_meta)} 帧")
 
@@ -227,6 +249,10 @@ class SimpleRobotDataset(Dataset):
         """获取所有相机键名 (LeRobotDataset 格式)"""
         return [f"observation.{cam_name}" for cam_name in self.camera_names]
 
+    @property
+    def available_episode_indices(self) -> List[int]:
+        return sorted(self.episodes.keys())
+
     @property
     def camera_info(self) -> dict:
         """获取相机信息"""

roboimi/vla/models/backbones/smolvla_prefix_encoder.py

@@ -0,0 +1,410 @@
+from __future__ import annotations
+
+import warnings
+from typing import Dict, Sequence
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from roboimi.vla.core.interfaces import VLABackbone
+
+try:  # pragma: no cover - exercised by tests via monkeypatch/fakes
+    from transformers import AutoModelForImageTextToText, AutoTokenizer
+except Exception:  # pragma: no cover
+    AutoModelForImageTextToText = None
+    AutoTokenizer = None
+
+
+def _resize_with_pad(img: torch.Tensor, width: int, height: int, pad_value: float = 0.0) -> torch.Tensor:
+    if img.ndim != 4:
+        raise ValueError(f'expected image tensor shaped (B,C,H,W), got {tuple(img.shape)}')
+    cur_height, cur_width = img.shape[2:]
+    ratio = max(cur_width / width, cur_height / height)
+    resized_height = int(cur_height / ratio)
+    resized_width = int(cur_width / ratio)
+    resized = F.interpolate(img, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
+    pad_height = max(0, int(height - resized_height))
+    pad_width = max(0, int(width - resized_width))
+    return F.pad(resized, (pad_width, 0, pad_height, 0), value=pad_value)
+
+
+def _pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor:
+    if vector.shape[-1] == new_dim:
+        return vector
+    if vector.shape[-1] > new_dim:
+        raise ValueError(f'cannot pad vector with dim {vector.shape[-1]} to smaller dim {new_dim}')
+    shape = list(vector.shape)
+    shape[-1] = int(new_dim)
+    padded = torch.zeros(*shape, dtype=vector.dtype, device=vector.device)
+    padded[..., : vector.shape[-1]] = vector
+    return padded
+
+
+class SmolVLAPrefixEncoder(VLABackbone):
+    """SmolVLA-compatible VLM prefix encoder for RoboIMI action experts.
+
+    This module intentionally extracts only the conditioning path from SmolVLA:
+    multiview image tokens, variable language task tokens, and one projected
+    state token.  It freezes the pretrained VLM by default and returns a fixed
+    condition token sequence `(B, S, D)` that can be consumed by existing IMF /
+    transformer-style action heads.
+    """
+
+    def __init__(
+        self,
+        model_name: str = 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
+        *,
+        model_name_or_path: str | None = None,
+        vlm: nn.Module | None = None,
+        tokenizer=None,
+        load_vlm_weights: bool = True,
+        local_files_only: bool = False,
+        num_vlm_layers: int = 16,
+        freeze_vlm: bool = True,
+        freeze_vision_encoder: bool = True,
+        train_state_proj: bool = True,
+        max_state_dim: int = 32,
+        resize_imgs_with_padding: Sequence[int] | None = (512, 512),
+        dataset_image_resize_shape: Sequence[int] | None = None,
+        eval_image_resize_shape: Sequence[int] | None = None,
+        tokenizer_max_length: int = 48,
+        pad_language_to: str = 'max_length',
+        run_text_model: bool = True,
+        camera_names: Sequence[str] = ('r_vis', 'top', 'front'),
+        num_cameras: int | None = None,
+    ) -> None:
+        super().__init__()
+        if model_name_or_path is not None:
+            model_name = model_name_or_path
+        self.model_name = str(model_name)
+        self.camera_names = tuple(camera_names)
+        self.num_cameras = int(num_cameras) if num_cameras is not None else len(self.camera_names)
+        if len(self.camera_names) != self.num_cameras:
+            raise ValueError(
+                f'camera_names length ({len(self.camera_names)}) must match num_cameras ({self.num_cameras})'
+            )
+        self.load_vlm_weights = bool(load_vlm_weights)
+        self.local_files_only = bool(local_files_only)
+        self.num_vlm_layers_requested = int(num_vlm_layers)
+        self.freeze_vlm = bool(freeze_vlm)
+        self.freeze_vision_encoder = bool(freeze_vision_encoder)
+        self.max_state_dim = int(max_state_dim)
+        self.resize_imgs_with_padding = self._normalize_resize_shape(resize_imgs_with_padding)
+        self.dataset_image_resize_shape = self._normalize_resize_shape(dataset_image_resize_shape)
+        self.eval_image_resize_shape = self._normalize_resize_shape(eval_image_resize_shape)
+        self.tokenizer_max_length = int(tokenizer_max_length)
+        self.pad_language_to = str(pad_language_to)
+        self.run_text_model = bool(run_text_model)
+        self._warned_trainable_frozen_text_model = False
+
+        if vlm is None:
+            if AutoModelForImageTextToText is None:
+                raise ImportError('transformers AutoModelForImageTextToText is required for SmolVLAPrefixEncoder')
+            if not self.load_vlm_weights:
+                raise ValueError('SmolVLAPrefixEncoder currently requires load_vlm_weights=True')
+            vlm = AutoModelForImageTextToText.from_pretrained(
+                self.model_name,
+                torch_dtype='bfloat16',
+                low_cpu_mem_usage=True,
+                local_files_only=self.local_files_only,
+            )
+        self.vlm = vlm
+
+        if tokenizer is None:
+            if AutoTokenizer is None:
+                raise ImportError('transformers AutoTokenizer is required for SmolVLAPrefixEncoder')
+            tokenizer = AutoTokenizer.from_pretrained(self.model_name, local_files_only=self.local_files_only)
+        self.tokenizer = tokenizer
+
+        if self.num_vlm_layers_requested > 0:
+            text_layers = self.vlm.model.text_model.layers
+            if self.num_vlm_layers_requested > len(text_layers):
+                raise ValueError(
+                    f'num_vlm_layers ({self.num_vlm_layers_requested}) exceeds available text layers '
+                    f'({len(text_layers)})'
+                )
+            self.vlm.model.text_model.layers = text_layers[: self.num_vlm_layers_requested]
+        self.num_vlm_layers = len(self.vlm.model.text_model.layers)
+        text_config = getattr(self.vlm.config, 'text_config', None)
+        if text_config is not None and hasattr(text_config, 'num_hidden_layers'):
+            text_config.num_hidden_layers = self.num_vlm_layers
+
+        hidden_size = int(self.vlm.config.text_config.hidden_size)
+        self._output_dim = hidden_size
+        self.state_proj = nn.Linear(self.max_state_dim, hidden_size)
+        for param in self.state_proj.parameters():
+            param.requires_grad = bool(train_state_proj)
+
+        self.last_prefix_pad_mask: torch.Tensor | None = None
+        self.last_prefix_att_mask: torch.Tensor | None = None
+        self.last_attention_2d_mask: torch.Tensor | None = None
+        self.last_position_ids: torch.Tensor | None = None
+        self._configured_condition_sequence_length = self._infer_configured_condition_sequence_length()
+
+        self.set_requires_grad()
+
+    @staticmethod
+    def _normalize_resize_shape(shape: Sequence[int] | None) -> tuple[int, int] | None:
+        if shape is None:
+            return None
+        normalized = tuple(int(v) for v in shape)
+        if len(normalized) != 2:
+            raise ValueError(f'resize_imgs_with_padding must contain exactly two values, got {normalized}')
+        return normalized
+
+    @property
+    def output_dim(self) -> int:
+        return self._output_dim
+
+    @property
+    def joint_output_dim(self) -> int:
+        return self._output_dim
+
+    def _infer_configured_condition_sequence_length(self) -> int:
+        image_tokens_per_camera = 0
+        config = getattr(self.vlm, 'config', None)
+        vision_config = getattr(config, 'vision_config', None)
+        if vision_config is not None:
+            image_size = int(getattr(vision_config, 'image_size', 0) or 0)
+            patch_size = int(getattr(vision_config, 'patch_size', 0) or 0)
+            scale_factor = int(getattr(config, 'scale_factor', 1) or 1)
+            if image_size > 0 and patch_size > 0 and scale_factor > 0:
+                image_tokens_per_camera = int(((image_size // patch_size) ** 2) / (scale_factor**2))
+        return self.num_cameras * image_tokens_per_camera + self.tokenizer_max_length + 1
+
+    @property
+    def tokens_per_step(self) -> int:
+        if self.last_prefix_pad_mask is not None:
+            return int(self.last_prefix_pad_mask.shape[1])
+        return int(self._configured_condition_sequence_length)
+
+    @property
+    def condition_sequence_length(self) -> int:
+        return self.tokens_per_step
+
+    @staticmethod
+    def _freeze_module(module) -> None:
+        if module is None:
+            return
+        if hasattr(module, 'eval'):
+            module.eval()
+        parameters = getattr(module, 'parameters', None)
+        if callable(parameters):
+            for param in parameters():
+                param.requires_grad = False
+
+    def set_requires_grad(self) -> None:
+        if self.freeze_vision_encoder:
+            self._freeze_module(self.vlm.model.vision_model)
+        if self.freeze_vlm:
+            self._freeze_module(self.vlm)
+            self._freeze_module(getattr(self.vlm.model, 'vision_model', None))
+            self._freeze_module(getattr(self.vlm.model, 'connector', None))
+            self._freeze_module(getattr(self.vlm.model, 'text_model', None))
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        if self.freeze_vlm:
+            self.vlm.eval()
+        elif self.freeze_vision_encoder:
+            self.vlm.model.vision_model.eval()
+        return self
+
+    def _ordered_camera_names(self, images: Dict[str, torch.Tensor]) -> tuple[str, ...]:
+        missing = [name for name in self.camera_names if name not in images]
+        if missing:
+            raise ValueError(f'image input missing required cameras. missing={missing}, expected={list(self.camera_names)}')
+        return self.camera_names
+
+    def _batch_size_from_images(self, images: Dict[str, torch.Tensor]) -> int:
+        return int(next(iter(images.values())).shape[0])
+
+    def _normalize_tasks(self, task, batch_size: int) -> list[str]:
+        if task is None:
+            tasks = [''] * batch_size
+        elif isinstance(task, str):
+            tasks = [task] * batch_size
+        elif isinstance(task, tuple):
+            tasks = list(task)
+        elif isinstance(task, list):
+            tasks = task
+        else:
+            raise TypeError(f'task must be str/list/tuple/None, got {type(task)!r}')
+        if len(tasks) != batch_size:
+            raise ValueError(f'task batch size ({len(tasks)}) must match image batch size ({batch_size})')
+        return [item if item.endswith('\n') else f'{item}\n' for item in tasks]
+
+    def tokenize_task(self, task, batch_size: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+        tasks = self._normalize_tasks(task, batch_size)
+        old_padding_side = getattr(self.tokenizer, 'padding_side', None)
+        if old_padding_side is not None:
+            self.tokenizer.padding_side = 'right'
+        try:
+            tokenized = self.tokenizer(
+                tasks,
+                padding=self.pad_language_to,
+                max_length=self.tokenizer_max_length,
+                return_tensors='pt',
+                truncation=True,
+            )
+        finally:
+            if old_padding_side is not None:
+                self.tokenizer.padding_side = old_padding_side
+        tokens = tokenized['input_ids'].to(device=device)
+        masks = tokenized['attention_mask'].to(device=device, dtype=torch.bool)
+        return tokens, masks
+
+    def prepare_images(self, images: Dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+        camera_names = self._ordered_camera_names(images)
+        reference = images[camera_names[0]]
+        if reference.ndim != 5:
+            raise ValueError(f'expected image tensor shaped (B,T,C,H,W), got {tuple(reference.shape)}')
+        batch_size = reference.shape[0]
+        prepared: list[torch.Tensor] = []
+        masks: list[torch.Tensor] = []
+        for camera_name in camera_names:
+            image = images[camera_name]
+            if image.shape[:2] != reference.shape[:2] or image.shape[2] != reference.shape[2]:
+                raise ValueError(f'camera {camera_name!r} shape {tuple(image.shape)} does not match reference {tuple(reference.shape)}')
+            image = image[:, -1].contiguous().float().clamp(0.0, 1.0)
+            if self.resize_imgs_with_padding is not None:
+                image = _resize_with_pad(image, *self.resize_imgs_with_padding, pad_value=0.0)
+            image = image * 2.0 - 1.0
+            prepared.append(image)
+            masks.append(torch.ones(batch_size, dtype=torch.bool, device=image.device))
+        return prepared, masks
+
+    def prepare_state(self, state: torch.Tensor) -> torch.Tensor:
+        if state.ndim > 2:
+            state = state[:, -1, :]
+        return _pad_vector(state.float(), self.max_state_dim)
+
+    def embed_image(self, image: torch.Tensor) -> torch.Tensor:
+        if hasattr(self.vlm.model, 'get_image_features'):
+            pixel_values = image[:, None, ...]
+            pixel_attention_mask = torch.ones(
+                image.shape[0],
+                1,
+                image.shape[2],
+                image.shape[3],
+                dtype=torch.bool,
+                device=image.device,
+            )
+            with torch.set_grad_enabled(
+                torch.is_grad_enabled() and not self.freeze_vlm and not self.freeze_vision_encoder
+            ):
+                hidden = self.vlm.model.get_image_features(
+                    pixel_values=pixel_values,
+                    pixel_attention_mask=pixel_attention_mask,
+                    return_dict=True,
+                ).pooler_output
+        else:
+            vision_model = self.vlm.model.vision_model
+            with torch.set_grad_enabled(
+                torch.is_grad_enabled() and not self.freeze_vlm and not self.freeze_vision_encoder
+            ):
+                patch_attention_mask = torch.ones(
+                    image.shape[0],
+                    image.shape[2] // self.vlm.config.vision_config.patch_size,
+                    image.shape[3] // self.vlm.config.vision_config.patch_size,
+                    dtype=torch.bool,
+                    device=image.device,
+                ) if hasattr(self.vlm.config, 'vision_config') and hasattr(self.vlm.config.vision_config, 'patch_size') else None
+                hidden = vision_model(
+                    pixel_values=image.to(dtype=vision_model.dtype),
+                    patch_attention_mask=patch_attention_mask,
+                ).last_hidden_state
+                hidden = self.vlm.model.connector(hidden)
+        return hidden
+
+    def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
+        return self.vlm.model.text_model.get_input_embeddings()(tokens)
+
+    @staticmethod
+    def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor:
+        cumsum = torch.cumsum(att_masks, dim=1)
+        att_2d = cumsum[:, None, :] <= cumsum[:, :, None]
+        pad_2d = pad_masks[:, None, :] * pad_masks[:, :, None]
+        return att_2d & pad_2d
+
+    def embed_prefix(self, images: Dict[str, torch.Tensor], state: torch.Tensor, task=None) -> torch.Tensor:
+        prepared_images, image_masks = self.prepare_images(images)
+        batch_size = prepared_images[0].shape[0]
+        device = prepared_images[0].device
+
+        embs: list[torch.Tensor] = []
+        pad_masks: list[torch.Tensor] = []
+        att_masks: list[int] = []
+
+        for image, image_mask in zip(prepared_images, image_masks, strict=False):
+            img_emb = self.embed_image(image)
+            img_emb = img_emb * torch.tensor(
+                img_emb.shape[-1] ** 0.5,
+                dtype=img_emb.dtype,
+                device=img_emb.device,
+            )
+            num_img_tokens = img_emb.shape[1]
+            embs.append(img_emb)
+            pad_masks.append(image_mask[:, None].expand(batch_size, num_img_tokens))
+            att_masks += [0] * num_img_tokens
+
+        lang_tokens, lang_masks = self.tokenize_task(task, batch_size=batch_size, device=device)
+        lang_emb = self.embed_language_tokens(lang_tokens)
+        lang_emb = lang_emb * (lang_emb.shape[-1] ** 0.5)
+        embs.append(lang_emb)
+        pad_masks.append(lang_masks)
+        att_masks += [0] * lang_emb.shape[1]
+
+        state = self.prepare_state(state).to(device=device)
+        state_emb = self.state_proj(state).unsqueeze(1)
+        embs.append(state_emb)
+        pad_masks.append(torch.ones(batch_size, 1, dtype=torch.bool, device=device))
+        att_masks += [1]
+
+        prefix = torch.cat(embs, dim=1)
+        prefix_pad_mask = torch.cat(pad_masks, dim=1)
+        prefix_att_mask = torch.tensor(att_masks, dtype=torch.bool, device=device)[None, :].expand(batch_size, -1)
+
+        self.last_prefix_pad_mask = prefix_pad_mask
+        self.last_prefix_att_mask = prefix_att_mask
+        self.last_attention_2d_mask = self.make_att_2d_masks(prefix_pad_mask, prefix_att_mask)
+        self.last_position_ids = torch.cumsum(prefix_pad_mask, dim=1) - 1
+        return prefix
+
+    def encode_prefix(self, prefix: torch.Tensor) -> torch.Tensor:
+        if not self.run_text_model:
+            return prefix
+        if self.last_attention_2d_mask is None or self.last_position_ids is None:
+            raise RuntimeError('encode_prefix requires masks from embed_prefix')
+        text_model = self.vlm.model.text_model
+        text_dtype = getattr(text_model, 'dtype', prefix.dtype)
+        if (
+            self.freeze_vlm
+            and not self._warned_trainable_frozen_text_model
+            and any(param.requires_grad for param in text_model.parameters())
+        ):
+            warnings.warn(
+                'freeze_vlm=True but text_model has trainable parameters; temporarily disabling '
+                'their gradients while keeping gradients for trainable prefix inputs.',
+                RuntimeWarning,
+            )
+            self._warned_trainable_frozen_text_model = True
+        with torch.set_grad_enabled(torch.is_grad_enabled()):
+            outputs = text_model(
+                inputs_embeds=prefix.to(dtype=text_dtype),
+                attention_mask=self.last_attention_2d_mask[:, None, :, :],
+                position_ids=self.last_position_ids,
+                use_cache=False,
+                return_dict=True,
+            )
+        return outputs.last_hidden_state
+
+    def forward(self, images: Dict[str, torch.Tensor], state: torch.Tensor | None = None, task=None) -> torch.Tensor:
+        if state is None:
+            raise ValueError('SmolVLAPrefixEncoder.forward requires `state` for SmolVLA-compatible state token')
+        prefix = self.embed_prefix(images=images, state=state, task=task)
+        return self.encode_prefix(prefix)
+
+
+SmolVLMPrefixEncoder = SmolVLAPrefixEncoder

roboimi/vla/models/smolvla/__init__.py

@@ -0,0 +1,13 @@
+"""Native SmolVLA model core."""
+
+from .configuration import NativeSmolVLAConfig, SmolVLAConfig
+from .modeling import VLAFlowMatching, make_att_2d_masks, pad_vector, resize_with_pad
+
+__all__ = [
+    "NativeSmolVLAConfig",
+    "SmolVLAConfig",
+    "VLAFlowMatching",
+    "make_att_2d_masks",
+    "pad_vector",
+    "resize_with_pad",
+]

roboimi/vla/models/smolvla/configuration.py

@@ -0,0 +1,75 @@
+"""Native SmolVLA configuration, independent of LeRobot."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class NativeSmolVLAConfig:
+    """Lightweight configuration for the native SmolVLA model core.
+
+    This intentionally keeps only model-core fields needed by
+    :class:`VLAFlowMatching`; policy/dataset/optimizer concerns stay outside of
+    the native model package.
+    """
+
+    # Input / output structure.
+    n_obs_steps: int = 1
+    chunk_size: int = 50
+    n_action_steps: int = 50
+
+    # Shorter state and action vectors are padded before entering the model.
+    max_state_dim: int = 32
+    max_action_dim: int = 32
+
+    # Image preprocessing.
+    resize_imgs_with_padding: tuple[int, int] = (512, 512)
+
+    # Tokenizer / decoding.
+    tokenizer_max_length: int = 48
+    num_steps: int = 10
+
+    # Attention utils.
+    use_cache: bool = True
+
+    # Finetuning settings.
+    freeze_vision_encoder: bool = True
+    train_expert_only: bool = True
+    train_state_proj: bool = True
+
+    # VLM / expert construction settings.
+    vlm_model_name: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
+    load_vlm_weights: bool = False
+    add_image_special_tokens: bool = False
+    attention_mode: str = "cross_attn"
+    prefix_length: int = -1
+    pad_language_to: str = "longest"
+    num_expert_layers: int = -1
+    num_vlm_layers: int = 16
+    self_attn_every_n_layers: int = 2
+    expert_width_multiplier: float = 0.75
+
+    # Flow-matching timestep embedding.
+    min_period: float = 4e-3
+    max_period: float = 4.0
+
+    # Runtime settings.
+    device: str | None = None
+    compile_model: bool = False
+    compile_mode: str = "max-autotune"
+
+    # RTC hook placeholder; the native core does not implement RTC, but keeping
+    # the field lets migrated call sites pass configs through unchanged.
+    rtc_config: object | None = None
+
+    def __post_init__(self) -> None:
+        if self.n_action_steps > self.chunk_size:
+            raise ValueError(
+                "The chunk size is the upper bound for the number of action steps per model invocation. "
+                f"Got {self.n_action_steps} for `n_action_steps` and {self.chunk_size} for `chunk_size`."
+            )
+
+
+# Compatibility alias for migrated code that still imports SmolVLAConfig.
+SmolVLAConfig = NativeSmolVLAConfig

roboimi/vla/models/smolvla/modeling.py

@@ -0,0 +1,421 @@
+"""Native SmolVLA flow-matching core.
+
+This module ports the lightweight helpers and ``VLAFlowMatching`` from the
+LeRobot SmolVLA implementation while avoiding any LeRobot imports.  The heavy
+Transformers-backed VLM/expert is optional and can be injected for tests.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TypedDict
+
+try:
+    from typing import Unpack
+except ImportError:  # Python 3.10
+    from typing_extensions import Unpack
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from .configuration import NativeSmolVLAConfig
+
+
+class ActionSelectKwargs(TypedDict, total=False):
+    inference_delay: int | None
+    prev_chunk_left_over: Tensor | None
+    execution_horizon: int | None
+
+
+def create_sinusoidal_pos_embedding(
+    time: torch.Tensor,
+    dimension: int,
+    min_period: float,
+    max_period: float,
+    device: torch.device | str = "cpu",
+) -> Tensor:
+    """Compute sine/cosine positional embeddings for scalar positions."""
+    if dimension % 2 != 0:
+        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
+    if time.ndim != 1:
+        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
+
+    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=torch.float64, device=device)
+    period = min_period * (max_period / min_period) ** fraction
+    scaling_factor = 1.0 / period * 2 * math.pi
+    sin_input = scaling_factor[None, :] * time[:, None].to(dtype=torch.float64)
+    pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+    return pos_emb
+
+
+def make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) -> Tensor:
+    """Build Big Vision-style 2D prefix-LM attention masks.
+
+    ``pad_masks`` is ``bool[B, N]`` and marks valid tokens. ``att_masks`` is
+    ``bool/int[B, N]`` where cumulative increments begin new causal groups.  A
+    query token can attend to valid key tokens whose cumulative attention group
+    is less than or equal to the query's group.
+    """
+    if att_masks.ndim != 2:
+        raise ValueError(att_masks.ndim)
+    if pad_masks.ndim != 2:
+        raise ValueError(pad_masks.ndim)
+
+    cumsum = torch.cumsum(att_masks.to(dtype=torch.long), dim=1)
+    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
+    pad_2d_masks = pad_masks[:, None, :].bool() & pad_masks[:, :, None].bool()
+    return att_2d_masks & pad_2d_masks
+
+
+def resize_with_pad(img: Tensor, width: int, height: int, pad_value: float = -1) -> Tensor:
+    """Resize a BCHW image batch preserving aspect ratio, then top/left pad."""
+    if img.ndim != 4:
+        raise ValueError(f"(b,c,h,w) expected, but {img.shape}")
+
+    cur_height, cur_width = img.shape[2:]
+    ratio = max(cur_width / width, cur_height / height)
+    resized_height = int(cur_height / ratio)
+    resized_width = int(cur_width / ratio)
+    resized_img = F.interpolate(
+        img,
+        size=(resized_height, resized_width),
+        mode="bilinear",
+        align_corners=False,
+    )
+
+    pad_height = max(0, int(height - resized_height))
+    pad_width = max(0, int(width - resized_width))
+    return F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
+
+
+def pad_vector(vector: Tensor, new_dim: int) -> Tensor:
+    """Pad a vector-like tensor's last dimension with zeros to ``new_dim``."""
+    current_dim = vector.shape[-1]
+    if current_dim == new_dim:
+        return vector
+    if current_dim > new_dim:
+        raise ValueError(f"Cannot pad vector with current dimension {current_dim} to smaller target dimension {new_dim}.")
+
+    shape = list(vector.shape)
+    shape[-1] = new_dim
+    new_vector = torch.zeros(*shape, dtype=vector.dtype, device=vector.device)
+    new_vector[..., :current_dim] = vector
+    return new_vector
+
+
+def pad_tensor(tensor: Tensor, max_len: int, pad_value: int | float) -> Tensor:
+    """Pad a tensor along sequence dimension to ``max_len``."""
+    bsize, seq_len = tensor.shape[:2]
+    if seq_len >= max_len:
+        return tensor
+    padded_tensor = torch.full(
+        (bsize, max_len, *tensor.shape[2:]),
+        pad_value,
+        dtype=tensor.dtype,
+        device=tensor.device,
+    )
+    padded_tensor[:, :seq_len] = tensor
+    return padded_tensor
+
+
+class VLAFlowMatching(nn.Module):
+    """SmolVLA flow-matching action head around a VLM plus action expert."""
+
+    def __init__(
+        self,
+        config: NativeSmolVLAConfig,
+        rtc_processor: object | None = None,
+        vlm_with_expert: nn.Module | None = None,
+    ):
+        super().__init__()
+        self.config = config
+
+        if vlm_with_expert is None:
+            from .smolvlm_with_expert import SmolVLMWithExpertModel
+
+            vlm_with_expert = SmolVLMWithExpertModel(
+                model_id=self.config.vlm_model_name,
+                freeze_vision_encoder=self.config.freeze_vision_encoder,
+                train_expert_only=self.config.train_expert_only,
+                load_vlm_weights=self.config.load_vlm_weights,
+                attention_mode=self.config.attention_mode,
+                num_expert_layers=self.config.num_expert_layers,
+                num_vlm_layers=self.config.num_vlm_layers,
+                self_attn_every_n_layers=self.config.self_attn_every_n_layers,
+                expert_width_multiplier=self.config.expert_width_multiplier,
+                device=self.config.device if self.config.device is not None else "auto",
+            )
+        self.vlm_with_expert = vlm_with_expert
+
+        vlm_hidden_size = self.vlm_with_expert.config.text_config.hidden_size
+        expert_hidden_size = self.vlm_with_expert.expert_hidden_size
+        self.state_proj = nn.Linear(self.config.max_state_dim, vlm_hidden_size)
+        self.action_in_proj = nn.Linear(self.config.max_action_dim, expert_hidden_size)
+        self.action_out_proj = nn.Linear(expert_hidden_size, self.config.max_action_dim)
+        self.action_time_mlp_in = nn.Linear(expert_hidden_size * 2, expert_hidden_size)
+        self.action_time_mlp_out = nn.Linear(expert_hidden_size, expert_hidden_size)
+
+        self.set_requires_grad()
+        tokenizer = self.vlm_with_expert.processor.tokenizer
+        self.fake_image_token = tokenizer.fake_image_token_id
+        self.global_image_token = tokenizer.global_image_token_id
+        self.global_image_start_token = torch.tensor(
+            [self.fake_image_token, self.global_image_token], dtype=torch.long
+        )
+        self.add_image_special_tokens = self.config.add_image_special_tokens
+        self.image_end_token = torch.tensor([self.fake_image_token], dtype=torch.long)
+        self.prefix_length = self.config.prefix_length
+        self.rtc_processor = rtc_processor
+
+        if config.compile_model:
+            torch.set_float32_matmul_precision("high")
+            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
+            self.forward = torch.compile(self.forward, mode=config.compile_mode)
+
+    def _rtc_enabled(self) -> bool:
+        return bool(self.config.rtc_config is not None and getattr(self.config.rtc_config, "enabled", False))
+
+    def set_requires_grad(self) -> None:
+        for params in self.state_proj.parameters():
+            params.requires_grad = self.config.train_state_proj
+
+    def sample_noise(self, shape: tuple[int, ...] | torch.Size, device: torch.device | str) -> Tensor:
+        return torch.normal(mean=0.0, std=1.0, size=shape, dtype=torch.float32, device=device)
+
+    def sample_time(self, bsize: int, device: torch.device | str) -> Tensor:
+        beta_dist = torch.distributions.Beta(concentration1=1.5, concentration0=1.0)
+        time_beta = beta_dist.sample((bsize,)).to(device=device, dtype=torch.float32)
+        return time_beta * 0.999 + 0.001
+
+    def _vlm_device(self) -> torch.device:
+        vlm = getattr(self.vlm_with_expert, "vlm", None)
+        return getattr(vlm, "device", next(self.parameters()).device)
+
+    def embed_prefix(
+        self,
+        images: list[Tensor],
+        img_masks: list[Tensor],
+        lang_tokens: Tensor,
+        lang_masks: Tensor,
+        state: Tensor,
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        """Embed images, language tokens, and robot state as prefix tokens."""
+        embs: list[Tensor] = []
+        pad_masks: list[Tensor] = []
+        att_masks: list[int] = []
+
+        for img, img_mask in zip(images, img_masks, strict=False):
+            if self.add_image_special_tokens:
+                image_start_token = (
+                    self.vlm_with_expert.embed_language_tokens(self.global_image_start_token.to(device=self._vlm_device()))
+                    .unsqueeze(0)
+                    .expand(img.shape[0], -1, -1)
+                )
+                image_start_mask = torch.ones_like(image_start_token[:, :, 0], dtype=torch.bool)
+                embs.append(image_start_token)
+                pad_masks.append(image_start_mask)
+                att_masks += [0] * image_start_mask.shape[-1]
+
+            img_emb = self.vlm_with_expert.embed_image(img)
+            img_emb = img_emb * torch.tensor(img_emb.shape[-1] ** 0.5, dtype=img_emb.dtype, device=img_emb.device)
+            bsize, num_img_embs = img_emb.shape[:2]
+            img_mask = img_mask.to(device=img_emb.device, dtype=torch.bool)[:, None].expand(bsize, num_img_embs)
+            embs.append(img_emb)
+            pad_masks.append(img_mask)
+            att_masks += [0] * num_img_embs
+
+            if self.add_image_special_tokens:
+                image_end_token = (
+                    self.vlm_with_expert.embed_language_tokens(self.image_end_token.to(device=self._vlm_device()))
+                    .unsqueeze(0)
+                    .expand(img.shape[0], -1, -1)
+                )
+                image_end_mask = torch.ones_like(image_end_token[:, :, 0], dtype=torch.bool)
+                embs.append(image_end_token)
+                pad_masks.append(image_end_mask)
+                att_masks += [0] * image_end_mask.shape[1]
+
+        lang_emb = self.vlm_with_expert.embed_language_tokens(lang_tokens)
+        lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
+        embs.append(lang_emb)
+        pad_masks.append(lang_masks.to(device=lang_emb.device, dtype=torch.bool))
+        att_masks += [0] * lang_emb.shape[1]
+
+        state_emb = self.state_proj(state)
+        state_emb = state_emb[:, None, :] if state_emb.ndim == 2 else state_emb
+        embs.append(state_emb)
+        state_mask = torch.ones(state_emb.shape[:2], dtype=torch.bool, device=state_emb.device)
+        pad_masks.append(state_mask)
+        att_masks += [1] * state_emb.shape[1]
+
+        all_embs = torch.cat(embs, dim=1)
+        all_pad_masks = torch.cat(pad_masks, dim=1)
+        all_att_masks = torch.tensor(att_masks, dtype=torch.bool, device=all_pad_masks.device)[None, :]
+
+        if self.prefix_length > 0 and all_pad_masks.shape[1] < self.prefix_length:
+            all_embs = pad_tensor(all_embs, self.prefix_length, pad_value=0)
+            all_pad_masks = pad_tensor(all_pad_masks, self.prefix_length, pad_value=0)
+            all_att_masks = pad_tensor(all_att_masks, self.prefix_length, pad_value=0)
+
+        all_att_masks = all_att_masks.expand(all_pad_masks.shape[0], -1)
+        return all_embs, all_pad_masks, all_att_masks
+
+    def embed_suffix(self, noisy_actions: Tensor, timestep: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        """Embed noisy action tokens and timestep for the expert suffix."""
+        action_emb = self.action_in_proj(noisy_actions)
+        device = action_emb.device
+        bsize = action_emb.shape[0]
+        dtype = action_emb.dtype
+
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.vlm_with_expert.expert_hidden_size,
+            self.config.min_period,
+            self.config.max_period,
+            device=device,
+        ).to(dtype=dtype)
+        time_emb = time_emb[:, None, :].expand_as(action_emb)
+        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
+        action_time_emb = self.action_time_mlp_in(action_time_emb)
+        action_time_emb = F.silu(action_time_emb)
+        action_time_emb = self.action_time_mlp_out(action_time_emb)
+
+        action_time_mask = torch.ones(action_time_emb.shape[:2], dtype=torch.bool, device=device)
+        att_masks = torch.ones(bsize, self.config.chunk_size, dtype=torch.bool, device=device)
+        return action_time_emb, action_time_mask, att_masks
+
+    def forward(
+        self,
+        images: list[Tensor],
+        img_masks: list[Tensor],
+        lang_tokens: Tensor,
+        lang_masks: Tensor,
+        state: Tensor,
+        actions: Tensor,
+        noise: Tensor | None = None,
+        time: Tensor | None = None,
+    ) -> Tensor:
+        """Run a training forward pass and return per-element flow loss."""
+        if noise is None:
+            noise = self.sample_noise(actions.shape, actions.device)
+        if time is None:
+            time = self.sample_time(actions.shape[0], actions.device)
+
+        time_expanded = time[:, None, None]
+        x_t = time_expanded * noise + (1 - time_expanded) * actions
+        u_t = noise - actions
+
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
+            images, img_masks, lang_tokens, lang_masks, state=state
+        )
+        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(x_t, time)
+
+        pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)
+        att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)
+        att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
+        position_ids = torch.cumsum(pad_masks, dim=1) - 1
+
+        (_, suffix_out), _ = self.vlm_with_expert.forward(
+            attention_mask=att_2d_masks,
+            position_ids=position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, suffix_embs],
+            use_cache=False,
+            fill_kv_cache=False,
+        )
+        suffix_out = suffix_out[:, -self.config.chunk_size :].to(dtype=torch.float32)
+        v_t = self.action_out_proj(suffix_out)
+        return F.mse_loss(u_t, v_t, reduction="none")
+
+    def sample_actions(
+        self,
+        images: list[Tensor],
+        img_masks: list[Tensor],
+        lang_tokens: Tensor,
+        lang_masks: Tensor,
+        state: Tensor,
+        noise: Tensor | None = None,
+        **kwargs: Unpack[ActionSelectKwargs],
+    ) -> Tensor:
+        """Sample an action chunk with Euler integration over the flow field."""
+        bsize = state.shape[0]
+        device = state.device
+        if noise is None:
+            actions_shape = (bsize, self.config.chunk_size, self.config.max_action_dim)
+            noise = self.sample_noise(actions_shape, device)
+
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
+            images, img_masks, lang_tokens, lang_masks, state=state
+        )
+        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
+        prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
+        _, past_key_values = self.vlm_with_expert.forward(
+            attention_mask=prefix_att_2d_masks,
+            position_ids=prefix_position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, None],
+            use_cache=self.config.use_cache,
+            fill_kv_cache=True,
+        )
+
+        dt = -1.0 / self.config.num_steps
+        x_t = noise
+        for step in range(self.config.num_steps):
+            time = 1.0 + step * dt
+            time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
+
+            def denoise_step_partial_call(input_x_t: Tensor, current_timestep: Tensor = time_tensor) -> Tensor:
+                return self.denoise_step(
+                    x_t=input_x_t,
+                    prefix_pad_masks=prefix_pad_masks,
+                    past_key_values=past_key_values,
+                    timestep=current_timestep,
+                )
+
+            if self._rtc_enabled() and self.rtc_processor is not None:
+                v_t = self.rtc_processor.denoise_step(
+                    x_t=x_t,
+                    prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
+                    inference_delay=kwargs.get("inference_delay"),
+                    time=time,
+                    original_denoise_step_partial=denoise_step_partial_call,
+                    execution_horizon=kwargs.get("execution_horizon"),
+                )
+            else:
+                v_t = denoise_step_partial_call(x_t)
+            x_t = x_t + dt * v_t
+
+            if self.rtc_processor is not None and getattr(self.rtc_processor, "is_debug_enabled", lambda: False)():
+                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
+
+        return x_t
+
+    def denoise_step(
+        self,
+        prefix_pad_masks: Tensor,
+        past_key_values: object,
+        x_t: Tensor,
+        timestep: Tensor,
+    ) -> Tensor:
+        """Apply one denoising step at a given timestep."""
+        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(x_t, timestep)
+        suffix_len = suffix_pad_masks.shape[1]
+        batch_size = prefix_pad_masks.shape[0]
+        prefix_len = prefix_pad_masks.shape[1]
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
+        prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
+        position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
+
+        outputs_embeds, _ = self.vlm_with_expert.forward(
+            attention_mask=full_att_2d_masks,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, suffix_embs],
+            use_cache=self.config.use_cache,
+            fill_kv_cache=False,
+        )
+        suffix_out = outputs_embeds[1][:, -self.config.chunk_size :].to(dtype=torch.float32)
+        return self.action_out_proj(suffix_out)

roboimi/vla/models/smolvla/smolvlm_with_expert.py

@@ -0,0 +1,571 @@
+# Copyright 2025 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.
+
+import copy
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from transformers import (
+        AutoConfig,
+        AutoModel,
+        AutoModelForImageTextToText,
+        AutoProcessor,
+        SmolVLMForConditionalGeneration,
+    )
+
+
+def _require_transformers():
+    try:
+        from transformers import (
+            AutoConfig,
+            AutoModel,
+            AutoModelForImageTextToText,
+            AutoProcessor,
+            SmolVLMForConditionalGeneration,
+        )
+    except ImportError as exc:
+        raise ImportError(
+            "Native SmolVLA requires the optional `transformers` package to construct "
+            "SmolVLMWithExpertModel. Install transformers or inject `vlm_with_expert` "
+            "when constructing VLAFlowMatching."
+        ) from exc
+    return AutoConfig, AutoModel, AutoModelForImageTextToText, AutoProcessor, SmolVLMForConditionalGeneration
+
+
+def apply_rope(x, positions, max_wavelength=10_000):
+    """
+    Applies RoPE positions [B, L] to x [B, L, H, D].
+    """
+    d_half = x.shape[-1] // 2
+    device = x.device
+    dtype = x.dtype
+    x = x.to(torch.float32)
+
+    freq_exponents = (2.0 / x.shape[-1]) * torch.arange(d_half, dtype=torch.float32, device=device)
+    timescale = max_wavelength**freq_exponents
+    radians = positions[..., None].to(torch.float32) / timescale[None, None, :].to(torch.float32)
+
+    radians = radians[..., None, :]
+
+    sin = torch.sin(radians)  # .to(dtype=dtype)
+    cos = torch.cos(radians)  # .to(dtype=dtype)
+
+    x1, x2 = x.split(d_half, dim=-1)
+    res = torch.empty_like(x)
+    res[..., :d_half] = x1 * cos - x2 * sin
+    res[..., d_half:] = x2 * cos + x1 * sin
+
+    return res.to(dtype)
+
+
+def get_intermediate_size(hidden_dim, ffn_dim_multiplier=4, multiple_of=256):
+    hidden_dim = int(2 * hidden_dim / 3)
+    hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+    hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+    return hidden_dim
+
+
+class SmolVLMWithExpertModel(nn.Module):
+    def __init__(
+        self,
+        model_id: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct",
+        load_vlm_weights: bool = True,
+        train_expert_only: bool = True,
+        freeze_vision_encoder: bool = False,
+        attention_mode: str = "self_attn",
+        num_expert_layers: int = -1,
+        num_vlm_layers: int = -1,
+        self_attn_every_n_layers: int = -1,
+        expert_width_multiplier: float = 0.5,
+        device: str = "auto",
+    ):
+        super().__init__()
+        AutoConfig, AutoModel, AutoModelForImageTextToText, AutoProcessor, SmolVLMForConditionalGeneration = _require_transformers()
+        if load_vlm_weights:
+            print(f"Loading  {model_id} weights ...")
+            self.vlm = AutoModelForImageTextToText.from_pretrained(
+                model_id,
+                torch_dtype="bfloat16",
+                low_cpu_mem_usage=True,
+            )
+            config = self.vlm.config
+        else:
+            config = AutoConfig.from_pretrained(model_id)
+            self.vlm = SmolVLMForConditionalGeneration(config=config)
+        self.processor = AutoProcessor.from_pretrained(model_id)
+        if num_vlm_layers > 0:
+            print(f"Reducing the number of VLM layers to {num_vlm_layers} ...")
+            self.get_vlm_model().text_model.layers = self.get_vlm_model().text_model.layers[:num_vlm_layers]
+        self.num_vlm_layers = len(self.get_vlm_model().text_model.layers)
+        self.config = config
+        # Smaller lm expert
+        lm_expert_config = copy.deepcopy(config.text_config)
+        hidden_size = lm_expert_config.hidden_size
+        lm_expert_config.hidden_size = int(hidden_size * expert_width_multiplier)  # hidden_size // 2
+        lm_expert_config.intermediate_size = get_intermediate_size(int(hidden_size * expert_width_multiplier))
+        lm_expert_config.num_hidden_layers = self.num_vlm_layers
+        if num_expert_layers > 0:
+            assert len(self.get_vlm_model().text_model.layers) % num_expert_layers == 0, (
+                f"Number of layers in the VLM {len(self.get_vlm_model().text_model.layers)} are not multiple of num_expert_layers {num_expert_layers}"
+            )
+            lm_expert_config.num_hidden_layers = num_expert_layers
+        self.lm_expert = AutoModel.from_config(lm_expert_config)
+
+        self.num_expert_layers = len(self.lm_expert.layers)
+        self.self_attn_every_n_layers = self_attn_every_n_layers
+        if "cross" in attention_mode:
+            # Reshape qkv projections to have the same input dimension as the vlm
+            for layer_idx in range(len(self.lm_expert.layers)):
+                if self.self_attn_every_n_layers > 0 and layer_idx % self.self_attn_every_n_layers == 0:
+                    continue
+                self.lm_expert.layers[layer_idx].self_attn.k_proj = nn.Linear(
+                    config.text_config.num_key_value_heads * config.text_config.head_dim,
+                    lm_expert_config.num_key_value_heads * lm_expert_config.head_dim,
+                    bias=lm_expert_config.attention_bias,
+                )
+                self.lm_expert.layers[layer_idx].self_attn.v_proj = nn.Linear(
+                    config.text_config.num_key_value_heads * config.text_config.head_dim,
+                    lm_expert_config.num_key_value_heads * lm_expert_config.head_dim,
+                    bias=lm_expert_config.attention_bias,
+                )
+        # Remove unused embed_tokens
+        self.lm_expert.embed_tokens = None
+
+        self.num_attention_heads = self.config.text_config.num_attention_heads
+        self.num_key_value_heads = self.config.text_config.num_key_value_heads
+
+        self.freeze_vision_encoder = freeze_vision_encoder
+        self.train_expert_only = train_expert_only
+        self.attention_mode = attention_mode
+        self.expert_hidden_size = lm_expert_config.hidden_size
+        self.set_requires_grad()
+
+    def get_vlm_model(self):
+        return self.vlm.model
+
+    def set_requires_grad(self):
+        if self.freeze_vision_encoder:
+            self.get_vlm_model().vision_model.eval()
+            for params in self.get_vlm_model().vision_model.parameters():
+                params.requires_grad = False
+        if self.train_expert_only:
+            self.vlm.eval()
+            for params in self.vlm.parameters():
+                params.requires_grad = False
+        else:
+            # To avoid unused params issue with distributed training
+            last_layers = [self.num_vlm_layers - 1]
+            if (
+                self.num_vlm_layers != self.num_expert_layers
+                and self.num_vlm_layers % self.num_expert_layers == 0
+            ):
+                last_layers.append(self.num_vlm_layers - 2)
+            frozen_layers = [
+                "lm_head",
+                "text_model.model.norm.weight",
+            ]
+            for layer in last_layers:
+                frozen_layers.append(f"text_model.model.layers.{layer}.")
+
+            for name, params in self.vlm.named_parameters():
+                if any(k in name for k in frozen_layers):
+                    params.requires_grad = False
+        # To avoid unused params issue with distributed training
+        for name, params in self.lm_expert.named_parameters():
+            if "lm_head" in name:
+                params.requires_grad = False
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+
+        if self.freeze_vision_encoder:
+            self.get_vlm_model().vision_model.eval()
+
+        if self.train_expert_only:
+            self.vlm.eval()
+
+    def embed_image(self, image: torch.Tensor):
+        patch_attention_mask = None
+        # Get sequence from the vision encoder
+        image_hidden_states = (
+            self.get_vlm_model()
+            .vision_model(
+                pixel_values=image.to(dtype=self.get_vlm_model().vision_model.dtype),
+                patch_attention_mask=patch_attention_mask,
+            )
+            .last_hidden_state
+        )
+        # Modality projection & resampling
+        image_hidden_states = self.get_vlm_model().connector(image_hidden_states)
+        return image_hidden_states
+
+    def embed_language_tokens(self, tokens: torch.Tensor):
+        return self.get_vlm_model().text_model.get_input_embeddings()(tokens)
+
+    def forward_attn_layer(
+        self,
+        model_layers,
+        inputs_embeds,
+        layer_idx,
+        position_ids,
+        attention_mask,
+        batch_size,
+        head_dim,
+        use_cache: bool = True,
+        fill_kv_cache: bool = True,
+        past_key_values=None,
+    ) -> list[torch.Tensor]:
+        query_states = []
+        key_states = []
+        value_states = []
+        for i, hidden_states in enumerate(inputs_embeds):
+            layer = model_layers[i][layer_idx]
+            if hidden_states is None or layer is None:
+                continue
+            hidden_states = layer.input_layernorm(hidden_states)
+
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+
+            hidden_states = hidden_states.to(dtype=layer.self_attn.q_proj.weight.dtype)
+            query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
+            key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
+            value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
+
+            query_states.append(query_state)
+            key_states.append(key_state)
+            value_states.append(value_state)
+
+        # B,L,H,D with L sequence length, H number of heads, D head dim
+        # concatenate on the number of embeddings/tokens
+        query_states = torch.cat(query_states, dim=1)
+        key_states = torch.cat(key_states, dim=1)
+        value_states = torch.cat(value_states, dim=1)
+        seq_len = query_states.shape[1]
+        if seq_len < position_ids.shape[1]:
+            _position_ids = position_ids[:, :seq_len]
+            _attention_mask = attention_mask[:, :seq_len, :seq_len]
+        else:
+            _position_ids = position_ids
+            _attention_mask = attention_mask
+
+        attention_mask_ = _attention_mask
+        position_ids_ = _position_ids
+
+        query_states = apply_rope(query_states, position_ids_)
+        key_states = apply_rope(key_states, position_ids_)
+
+        if use_cache and past_key_values is None:
+            past_key_values = {}
+
+        if use_cache:
+            if fill_kv_cache:
+                past_key_values[layer_idx] = {
+                    "key_states": key_states,
+                    "value_states": value_states,
+                }
+            else:
+                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
+                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
+                # the max len, then we (for instance) double the cache size. This implementation already exists
+                # in `transformers`. (molbap)
+                key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1)
+                value_states = torch.cat([past_key_values[layer_idx]["value_states"], value_states], dim=1)
+
+        attention_interface = self.get_attention_interface()
+
+        att_output = attention_interface(
+            attention_mask_, batch_size, head_dim, query_states, key_states, value_states
+        )
+        return [att_output], past_key_values
+
+    def forward_cross_attn_layer(
+        self,
+        model_layers,
+        inputs_embeds,
+        layer_idx,
+        position_ids,
+        attention_mask,
+        batch_size,
+        head_dim,
+        use_cache: bool = True,
+        fill_kv_cache: bool = True,
+        past_key_values=None,
+    ) -> list[torch.Tensor]:
+        attention_interface = self.get_attention_interface()
+
+        att_outputs = []
+        assert len(inputs_embeds) == 2 or (use_cache and past_key_values is not None and not fill_kv_cache), (
+            f"Both len(inputs_embeds) == {len(inputs_embeds)} and past_key_values is {past_key_values}"
+        )
+
+        if len(inputs_embeds) == 2 and not past_key_values:
+            # Prefix attention
+            seq_len = inputs_embeds[0].shape[1]
+            position_id, expert_position_id = position_ids[:, :seq_len], position_ids[:, seq_len:]
+            prefix_attention_mask = attention_mask[:, :seq_len, :seq_len]
+
+            layer = model_layers[0][layer_idx]
+
+            hidden_states = layer.input_layernorm(inputs_embeds[0])
+
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+
+            hidden_states = hidden_states.to(dtype=layer.self_attn.q_proj.weight.dtype)
+            query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
+            key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
+            value_states = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
+
+            # B,L,H,D with L sequence length, H number of heads, D head dim
+            query_states = apply_rope(query_state, position_id)
+            key_states = apply_rope(key_state, position_id)
+
+            att_output = attention_interface(
+                prefix_attention_mask, batch_size, head_dim, query_states, key_states, value_states
+            )
+            att_outputs.append(att_output)
+        else:
+            expert_position_id = position_ids
+
+        if use_cache and past_key_values is None:
+            past_key_values = {}
+
+        if use_cache:
+            if fill_kv_cache:
+                past_key_values[layer_idx] = {
+                    "key_states": key_states,
+                    "value_states": value_states,
+                }
+            else:
+                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
+                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
+                # the max len, then we (for instance) double the cache size. This implementation already exists
+                # in `transformers`. (molbap)
+                key_states = past_key_values[layer_idx]["key_states"]
+                value_states = past_key_values[layer_idx]["value_states"]
+
+        # Expert
+        expert_layer = model_layers[1][layer_idx]
+        if expert_layer is not None:
+            expert_hidden_states = expert_layer.input_layernorm(inputs_embeds[1])
+
+            expert_input_shape = expert_hidden_states.shape[:-1]
+            expert_hidden_shape = (*expert_input_shape, -1, expert_layer.self_attn.head_dim)
+
+            expert_hidden_states = expert_hidden_states.to(dtype=expert_layer.self_attn.q_proj.weight.dtype)
+            expert_query_state = expert_layer.self_attn.q_proj(expert_hidden_states).view(expert_hidden_shape)
+
+            _key_states = key_states.to(dtype=expert_layer.self_attn.k_proj.weight.dtype).view(
+                *key_states.shape[:2], -1
+            )
+            expert_key_states = expert_layer.self_attn.k_proj(_key_states).view(
+                *_key_states.shape[:-1], -1, expert_layer.self_attn.head_dim
+            )  # k_proj should have same dim as kv
+
+            _value_states = value_states.to(dtype=expert_layer.self_attn.v_proj.weight.dtype).view(
+                *value_states.shape[:2], -1
+            )
+            expert_value_states = expert_layer.self_attn.v_proj(_value_states).view(
+                *_value_states.shape[:-1], -1, expert_layer.self_attn.head_dim
+            )
+
+            expert_position_id = (
+                expert_position_id - torch.min(expert_position_id, dim=1, keepdim=True).values
+            )  # start from 0
+            expert_attention_mask = attention_mask[
+                :, -inputs_embeds[1].shape[1] :, : expert_key_states.shape[1] :
+            ]  # take into account kv
+
+            expert_query_states = apply_rope(expert_query_state, expert_position_id)
+
+            att_output = attention_interface(
+                expert_attention_mask,
+                batch_size,
+                head_dim,
+                expert_query_states,
+                expert_key_states,
+                expert_value_states,
+            )
+            att_outputs.append(att_output)
+        else:
+            att_outputs.append(None)
+
+        # att_output = att_output.to(dtype=models[i].dtype)
+        return att_outputs, past_key_values
+
+    def get_model_layers(self, models: list) -> list:
+        vlm_layers = []
+        expert_layers = []
+        multiple_of = self.num_vlm_layers // self.num_expert_layers
+        for i in range(self.num_vlm_layers):
+            if multiple_of > 0 and i > 0 and i % multiple_of != 0:
+                expert_layer = None
+            else:
+                expert_layer_index = i // multiple_of if multiple_of > 0 else i
+                expert_layer = models[1].layers[expert_layer_index]
+            vlm_layers.append(models[0].layers[i])
+            expert_layers.append(expert_layer)
+        return [vlm_layers, expert_layers]
+
+    def forward(
+        self,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: list[torch.FloatTensor] | None = None,
+        inputs_embeds: list[torch.FloatTensor] = None,
+        use_cache: bool | None = None,
+        fill_kv_cache: bool | None = None,
+    ):
+        models = [self.get_vlm_model().text_model, self.lm_expert]
+        model_layers = self.get_model_layers(models)
+        for hidden_states in inputs_embeds:
+            # TODO this is very inefficient
+            # dtype is always the same, batch size too (if > 1 len)
+            # device could be trickier in multi gpu edge cases but that's it
+            if hidden_states is None:
+                continue
+            batch_size = hidden_states.shape[0]
+
+        # RMSNorm
+        num_layers = self.num_vlm_layers
+        head_dim = self.vlm.config.text_config.head_dim
+        for layer_idx in range(num_layers):
+            if (
+                fill_kv_cache
+                or "cross" not in self.attention_mode
+                or (self.self_attn_every_n_layers > 0 and layer_idx % self.self_attn_every_n_layers == 0)
+            ):
+                att_outputs, past_key_values = self.forward_attn_layer(
+                    model_layers,
+                    inputs_embeds,
+                    layer_idx,
+                    position_ids,
+                    attention_mask,
+                    batch_size,
+                    head_dim,
+                    use_cache=use_cache,
+                    fill_kv_cache=fill_kv_cache,
+                    past_key_values=past_key_values,
+                )
+            else:
+                att_outputs, past_key_values = self.forward_cross_attn_layer(
+                    model_layers,
+                    inputs_embeds,
+                    layer_idx,
+                    position_ids,
+                    attention_mask,
+                    batch_size,
+                    head_dim,
+                    use_cache=use_cache,
+                    fill_kv_cache=fill_kv_cache,
+                    past_key_values=past_key_values,
+                )
+            outputs_embeds = []
+            start = 0
+            for i, hidden_states in enumerate(inputs_embeds):
+                layer = model_layers[i][layer_idx]
+                att_output = (
+                    att_outputs[i] if i < len(att_outputs) else att_outputs[0]
+                )  # in case of self_attn
+                if hidden_states is not None:
+                    if layer is None:
+                        outputs_embeds.append(hidden_states)
+                        continue
+                    end = start + hidden_states.shape[1]
+
+                    if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
+                        att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
+                    att_out = att_output[:, start:end]
+                    out_emb = layer.self_attn.o_proj(att_out)
+
+                    out_emb += hidden_states
+                    after_first_residual = out_emb.clone()
+
+                    out_emb = layer.post_attention_layernorm(out_emb)
+                    out_emb = layer.mlp(out_emb)
+
+                    out_emb += after_first_residual
+
+                    outputs_embeds.append(out_emb)
+
+                    start = end if len(att_outputs) == 1 else 0
+                else:
+                    outputs_embeds.append(None)
+
+            inputs_embeds = outputs_embeds
+
+        # final norm
+        outputs_embeds = []
+        for i, hidden_states in enumerate(inputs_embeds):
+            if hidden_states is not None:
+                out_emb = models[i].norm(hidden_states)
+                outputs_embeds.append(out_emb)
+            else:
+                outputs_embeds.append(None)
+        return outputs_embeds, past_key_values
+
+    def get_attention_interface(self):
+        attention_interface = self.eager_attention_forward
+        return attention_interface
+
+    def eager_attention_forward(
+        self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
+    ):
+        num_att_heads = self.num_attention_heads
+        num_key_value_heads = self.num_key_value_heads
+        num_key_value_groups = num_att_heads // num_key_value_heads
+
+        sequence_length = key_states.shape[1]
+
+        key_states = key_states[:, :, :, None, :].expand(
+            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
+        )
+        key_states = key_states.reshape(
+            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
+        )
+
+        value_states = value_states[:, :, :, None, :].expand(
+            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
+        )
+        value_states = value_states.reshape(
+            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
+        )
+
+        # Attention here is upcasted to float32 to match the original eager implementation.
+        query_states = query_states.to(dtype=torch.float32)
+        key_states = key_states.to(dtype=torch.float32)
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+
+        att_weights = torch.matmul(query_states, key_states.transpose(2, 3))
+        att_weights *= head_dim**-0.5
+
+        att_weights = att_weights.to(dtype=torch.float32)
+        big_neg = torch.finfo(att_weights.dtype).min  # -2.3819763e38  # See gemma/modules.py
+        masked_att_weights = torch.where(attention_mask[:, None, :, :], att_weights, big_neg)
+        probs = nn.functional.softmax(masked_att_weights, dim=-1)
+        probs = probs.to(dtype=value_states.dtype)
+
+        att_output = torch.matmul(probs, value_states.permute(0, 2, 1, 3))
+
+        att_output = att_output.permute(0, 2, 1, 3)
+        # we use -1 because sequence length can change
+        att_output = att_output.reshape(batch_size, -1, num_key_value_heads * num_key_value_groups * head_dim)
+
+        return att_output

tests/test_air_insert_env.py

@@ -0,0 +1,510 @@
+import importlib
+import inspect
+import pathlib
+import unittest
+from unittest import mock
+import xml.etree.ElementTree as ET
+
+import numpy as np
+
+from roboimi.envs.double_pos_ctrl_env import make_sim_env
+from roboimi.utils import act_ex_utils
+from roboimi.utils.constants import SIM_TASK_CONFIGS
+
+
+TASK_NAME = "sim_air_insert_socket_peg"
+
+
+class AirInsertTaskRegistrationTest(unittest.TestCase):
+    def test_sim_task_configs_registers_air_insert_socket_peg(self):
+        self.assertIn(TASK_NAME, SIM_TASK_CONFIGS)
+        self.assertNotIn("sim_air_insert_ring_bar", SIM_TASK_CONFIGS)
+        self.assertEqual(SIM_TASK_CONFIGS[TASK_NAME]["episode_len"], 750)
+        self.assertEqual(SIM_TASK_CONFIGS[TASK_NAME]["camera_names"], ["l_vis", "r_vis", "front"])
+        self.assertTrue(SIM_TASK_CONFIGS[TASK_NAME]["dataset_dir"].endswith("/sim_air_insert_socket_peg"))
+
+    def test_sample_air_insert_socket_peg_state_returns_explicit_named_mapping(self):
+        sampler = getattr(act_ex_utils, "sample_air_insert_socket_peg_state", None)
+        self.assertIsNotNone(
+            sampler,
+            "Expected roboimi.utils.act_ex_utils.sample_air_insert_socket_peg_state()",
+        )
+        self.assertFalse(
+            hasattr(act_ex_utils, "sample_air_insert_ring_bar_state"),
+            "air insert sampler should use socket/peg naming after the task rename",
+        )
+
+        task_state = sampler()
+
+        self.assertEqual(
+            list(task_state.keys()),
+            ["socket_pos", "socket_quat", "peg_pos", "peg_quat"],
+        )
+        self.assertEqual(task_state["socket_pos"].shape, (3,))
+        self.assertEqual(task_state["socket_quat"].shape, (4,))
+        self.assertEqual(task_state["peg_pos"].shape, (3,))
+        self.assertEqual(task_state["peg_quat"].shape, (4,))
+
+    def test_sample_air_insert_socket_peg_state_uses_fixed_quats_and_left_right_planar_ranges(self):
+        sampler = getattr(act_ex_utils, "sample_air_insert_socket_peg_state", None)
+        self.assertIsNotNone(sampler)
+
+        task_state = sampler()
+
+        np.testing.assert_array_equal(task_state["socket_quat"], np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
+        np.testing.assert_array_equal(task_state["peg_quat"], np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
+        self.assertGreaterEqual(task_state["socket_pos"][0], -0.20)
+        self.assertLessEqual(task_state["socket_pos"][0], -0.05)
+        self.assertGreaterEqual(task_state["socket_pos"][1], 0.70)
+        self.assertLessEqual(task_state["socket_pos"][1], 1.00)
+        self.assertAlmostEqual(float(task_state["socket_pos"][2]), 0.472)
+        self.assertGreaterEqual(task_state["peg_pos"][0], 0.05)
+        self.assertLessEqual(task_state["peg_pos"][0], 0.20)
+        self.assertGreaterEqual(task_state["peg_pos"][1], 0.70)
+        self.assertLessEqual(task_state["peg_pos"][1], 1.00)
+        self.assertAlmostEqual(float(task_state["peg_pos"][2]), 0.46)
+
+    def test_make_sim_env_dispatches_air_insert_socket_peg_headless(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        air_insert_cls = getattr(air_insert_env, "DualDianaMed_Air_Insert", None)
+        self.assertIsNotNone(air_insert_cls)
+
+        diana_med = importlib.import_module("roboimi.assets.robots.diana_med")
+        socket_peg_robot_cls = getattr(diana_med, "BiDianaMedSocketPeg", None)
+        self.assertIsNotNone(
+            socket_peg_robot_cls,
+            "Expected roboimi.assets.robots.diana_med.BiDianaMedSocketPeg",
+        )
+
+        fake_env = object()
+        with mock.patch.object(
+            diana_med,
+            "BiDianaMedSocketPeg",
+            return_value="robot",
+        ), mock.patch.object(
+            air_insert_env,
+            "DualDianaMed_Air_Insert",
+            return_value=fake_env,
+        ) as env_cls:
+            env = make_sim_env(TASK_NAME, headless=True)
+
+        self.assertIs(env, fake_env)
+        env_cls.assert_called_once_with(
+            robot="robot",
+            is_render=False,
+            control_freq=30,
+            is_interpolate=True,
+            cam_view="front",
+        )
+
+    def test_diana_table_scene_exposes_only_top_and_front_scene_cameras(self):
+        xml_path = (
+            pathlib.Path(__file__).resolve().parents[1]
+            / "roboimi/assets/models/manipulators/DianaMed/table_square.xml"
+        )
+        root = ET.parse(xml_path).getroot()
+        cameras = {camera.attrib["name"]: camera.attrib for camera in root.findall(".//camera")}
+
+        self.assertNotIn("angle", cameras, "DianaMed scene should stop exposing the old angle camera")
+        self.assertNotIn("left_side", cameras, "DianaMed scene should no longer expose left_side")
+        self.assertIn("top", cameras)
+        self.assertIn("front", cameras)
+        self.assertEqual(cameras["top"].get("mode"), "targetbody")
+        self.assertEqual(cameras["top"].get("target"), "table")
+
+
+class AirInsertResetAndStateHelpersTest(unittest.TestCase):
+    def test_set_socket_peg_task_state_writes_free_joint_qpos(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        setter = getattr(air_insert_env, "set_socket_peg_task_state", None)
+        self.assertIsNotNone(
+            setter,
+            "Expected roboimi.envs.double_air_insert_env.set_socket_peg_task_state",
+        )
+
+        socket_qpos = np.zeros(7, dtype=np.float64)
+        peg_qpos = np.zeros(7, dtype=np.float64)
+
+        class _FakeJoint:
+            def __init__(self, qpos):
+                self.qpos = qpos
+
+        class _FakeData:
+            def joint(self, name):
+                if name == "blue_socket_joint":
+                    return _FakeJoint(socket_qpos)
+                if name == "red_peg_joint":
+                    return _FakeJoint(peg_qpos)
+                raise AssertionError(f"Unexpected joint name: {name}")
+
+        task_state = {
+            "socket_pos": np.array([-0.12, 0.90, 0.472], dtype=np.float64),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64),
+            "peg_pos": np.array([0.12, 0.91, 0.46], dtype=np.float64),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64),
+        }
+
+        setter(_FakeData(), task_state)
+
+        np.testing.assert_array_equal(
+            socket_qpos,
+            np.array([-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0], dtype=np.float64),
+        )
+        np.testing.assert_array_equal(
+            peg_qpos,
+            np.array([0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0], dtype=np.float64),
+        )
+
+    def test_get_socket_peg_env_state_returns_stable_14d_vector(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        getter = getattr(air_insert_env, "get_socket_peg_env_state", None)
+        self.assertIsNotNone(
+            getter,
+            "Expected roboimi.envs.double_air_insert_env.get_socket_peg_env_state",
+        )
+
+        socket_qpos = np.array([-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)
+        peg_qpos = np.array([0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)
+
+        class _FakeJoint:
+            def __init__(self, qpos):
+                self.qpos = qpos
+
+        class _FakeData:
+            def joint(self, name):
+                if name == "blue_socket_joint":
+                    return _FakeJoint(socket_qpos)
+                if name == "red_peg_joint":
+                    return _FakeJoint(peg_qpos)
+                raise AssertionError(f"Unexpected joint name: {name}")
+
+        env_state = getter(_FakeData())
+
+        self.assertEqual(env_state.shape, (14,))
+        np.testing.assert_array_equal(
+            env_state,
+            np.array(
+                [-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0, 0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0],
+                dtype=np.float64,
+            ),
+        )
+
+    def test_air_insert_env_does_not_script_attach_or_assist_objects_after_reset(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        env_cls = getattr(air_insert_env, "DualDianaMed_Air_Insert", None)
+        self.assertIsNotNone(env_cls)
+
+        source = inspect.getsource(env_cls)
+
+        self.assertNotIn("_update_scripted_grasped_objects", source)
+        self.assertNotIn("_scripted_", source)
+        self.assertNotIn("_stabilize_ring_grasp", source)
+        self.assertNotIn("_ring_grasp_locked", source)
+        get_reward_source = inspect.getsource(env_cls._get_reward)
+        self.assertNotIn("ring_block", get_reward_source)
+        self.assertNotIn("bar_block", get_reward_source)
+
+    def test_socket_peg_xml_defines_active_socket_and_peg_objects(self):
+        asset_dir = pathlib.Path(__file__).resolve().parents[1] / "roboimi/assets/models/manipulators/DianaMed"
+        xml_path = asset_dir / "socket_peg_objects.xml"
+        self.assertTrue(xml_path.exists(), "socket/peg objects should live in socket_peg_objects.xml")
+        self.assertFalse((asset_dir / "ring_bar_objects.xml").exists(), "old ring_bar_objects.xml should be renamed")
+
+        root = ET.parse(xml_path).getroot()
+        body_names = {body.attrib.get("name") for body in root.findall(".//body")}
+        geom_names = {geom.attrib.get("name") for geom in root.findall(".//geom")}
+        joint_names = {joint.attrib.get("name") for joint in root.findall(".//joint")}
+
+        self.assertIn("socket", body_names)
+        self.assertIn("peg", body_names)
+        self.assertNotIn("ring_block", body_names)
+        self.assertNotIn("bar_block", body_names)
+        self.assertIn("blue_socket_joint", joint_names)
+        self.assertIn("red_peg_joint", joint_names)
+        for geom_name in ("socket-1", "socket-2", "socket-3", "socket-4", "pin", "red_peg"):
+            self.assertIn(geom_name, geom_names)
+
+    def test_socket_peg_wrapper_includes_socket_peg_objects(self):
+        xml_path = (
+            pathlib.Path(__file__).resolve().parents[1]
+            / "roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml"
+        )
+        self.assertTrue(xml_path.exists(), "socket/peg wrapper XML should use the new task name")
+        root = ET.parse(xml_path).getroot()
+        includes = [include.attrib.get("file") for include in root.findall(".//include")]
+        self.assertIn("./socket_peg_objects.xml", includes)
+        self.assertNotIn("./ring_bar_objects.xml", includes)
+
+
+class AirInsertRewardAndSuccessTest(unittest.TestCase):
+    @staticmethod
+    def _make_env_state(
+        socket_pos=(0.0, 0.0, 0.472),
+        socket_quat=(1.0, 0.0, 0.0, 0.0),
+        peg_pos=(0.0, 0.0, 0.46),
+        peg_quat=(1.0, 0.0, 0.0, 0.0),
+    ):
+        return np.array([*socket_pos, *socket_quat, *peg_pos, *peg_quat], dtype=np.float64)
+
+    def test_compute_air_insert_reward_counts_left_contact_stage(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+        self.assertIsNotNone(reward_fn)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("socket-1", "table"),
+                ("red_peg", "table"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 1)
+
+    def test_compute_air_insert_reward_counts_right_contact_stage(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("red_peg", "l_finger_right"),
+                ("socket-1", "table"),
+                ("red_peg", "table"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 2)
+
+    def test_compute_air_insert_reward_counts_lift_stages(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("red_peg", "l_finger_right"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 4)
+
+    def test_compute_air_insert_reward_counts_visual_fingertip_contacts_as_gripper_contacts(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-3", "r_fingertip_g0_vis_left"),
+                ("red_peg", "l_fingertip_g0_vis_right"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(
+            reward,
+            4,
+            "visual fingertip geoms are collidable in the Diana XML and should count as gripper-object contacts",
+        )
+
+    def test_peg_inserted_into_socket_uses_pin_contact(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        success_fn = getattr(air_insert_env, "peg_inserted_into_socket", None)
+        self.assertIsNotNone(
+            success_fn,
+            "Expected roboimi.envs.double_air_insert_env.peg_inserted_into_socket",
+        )
+
+        self.assertTrue(success_fn([("red_peg", "pin")]))
+        self.assertTrue(success_fn([("pin", "red_peg")]))
+        self.assertFalse(success_fn([("red_peg", "socket-1")]))
+
+    def test_compute_air_insert_reward_requires_airborne_success_for_final_point(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("red_peg", "l_finger_right"),
+                ("socket-1", "table"),
+                ("red_peg", "pin"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 3)
+
+    def test_compute_air_insert_reward_returns_full_score_on_true_airborne_insert(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("red_peg", "l_finger_right"),
+                ("red_peg", "pin"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 5)
+
+
+class AirInsertPolicyAndSmokeTest(unittest.TestCase):
+    @staticmethod
+    def _canonical_task_state():
+        return {
+            "socket_pos": np.array([-0.12, 0.90, 0.472], dtype=np.float32),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+            "peg_pos": np.array([0.12, 0.90, 0.46], dtype=np.float32),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        }
+
+    def test_air_insert_policy_emits_valid_16d_action(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = act_ex_utils.sample_air_insert_socket_peg_state()
+        policy = policy_cls(inject_noise=False)
+        action = policy.predict(task_state, 0)
+
+        self.assertEqual(action.shape, (16,))
+        np.testing.assert_array_equal(action[-2:], np.array([100, 100]))
+
+    def test_air_insert_policy_inserts_peg_front_view_right_to_left_along_world_x(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = self._canonical_task_state()
+        policy = policy_cls(inject_noise=False)
+        policy.generate_trajectory(task_state)
+
+        start_waypoint = next(wp for wp in policy.right_trajectory if wp["t"] == policy.INSERT_START_T)
+        end_waypoint = next(wp for wp in policy.right_trajectory if wp["t"] == policy.INSERT_END_T)
+
+        self.assertLess(
+            end_waypoint["xyz"][0],
+            start_waypoint["xyz"][0] - 0.10,
+            "front-view right-to-left peg insertion should decrease world x substantially",
+        )
+        self.assertAlmostEqual(float(end_waypoint["xyz"][1]), float(start_waypoint["xyz"][1]), delta=0.02)
+        expected_insert_end_x = float(task_state["socket_pos"][0] + 0.168)
+        self.assertAlmostEqual(float(end_waypoint["xyz"][0]), expected_insert_end_x, delta=0.02)
+        self.assertGreater(float(start_waypoint["xyz"][2]), 0.70)
+
+    def test_air_insert_policy_default_left_grasps_socket_and_right_grasps_peg(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = {
+            "socket_pos": np.array([-0.18, 0.78, 0.472], dtype=np.float32),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+            "peg_pos": np.array([0.16, 0.98, 0.46], dtype=np.float32),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        }
+
+        policy = policy_cls(inject_noise=False)
+        policy.generate_trajectory(task_state)
+        left_close = next(wp for wp in policy.left_trajectory if wp["t"] == 180)
+        right_close = next(wp for wp in policy.right_trajectory if wp["t"] == 180)
+        action_z_offset = getattr(policy_cls, "ACTION_OBJECT_Z_OFFSET", 0.11)
+        expected_socket_pick = task_state["socket_pos"] + np.array([-0.078, 0.0, action_z_offset])
+        expected_peg_pick = task_state["peg_pos"] + np.array([0.078, 0.0, action_z_offset + 0.01])
+
+        np.testing.assert_allclose(left_close["xyz"], expected_socket_pick, atol=1e-6)
+        np.testing.assert_allclose(right_close["xyz"], expected_peg_pick, atol=1e-6)
+        self.assertLess(left_close["gripper"], 0, "default policy should close the left gripper on the socket")
+        self.assertLess(right_close["gripper"], 0, "default policy should close the right gripper on the peg")
+
+    def test_air_insert_policy_socket_hold_tracks_socket_xy_without_sweeping_laterally(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        base_state = {
+            "socket_pos": np.array([-0.20, 0.72, 0.472], dtype=np.float32),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+            "peg_pos": np.array([0.14, 0.76, 0.46], dtype=np.float32),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        }
+        shifted_state = dict(base_state)
+        shifted_state["socket_pos"] = np.array([-0.06, 0.99, 0.472], dtype=np.float32)
+
+        base_policy = policy_cls(inject_noise=False)
+        base_policy.generate_trajectory(base_state)
+        shifted_policy = policy_cls(inject_noise=False)
+        shifted_policy.generate_trajectory(shifted_state)
+
+        base_hold = next(wp for wp in base_policy.left_trajectory if wp["t"] == 450)
+        shifted_hold = next(wp for wp in shifted_policy.left_trajectory if wp["t"] == 450)
+        np.testing.assert_allclose(
+            base_hold["xyz"][:2],
+            base_state["socket_pos"][:2] + np.array([-0.078, 0.0]),
+            atol=1e-6,
+        )
+        np.testing.assert_allclose(
+            shifted_hold["xyz"][:2],
+            shifted_state["socket_pos"][:2] + np.array([-0.078, 0.0]),
+            atol=1e-6,
+        )
+
+    def test_air_insert_policy_predicts_through_full_episode_without_exhausting_waypoints(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = self._canonical_task_state()
+        policy = policy_cls(inject_noise=False)
+
+        for step in range(SIM_TASK_CONFIGS[TASK_NAME]["episode_len"]):
+            action = policy.predict(task_state, step)
+            self.assertEqual(action.shape, (16,))
+
+    def test_scripted_rollout_entrypoint_selects_socket_peg_sampler_and_policy(self):
+        rollout_module = importlib.import_module("roboimi.demos.diana_record_sim_episodes")
+        sampler_fn = getattr(rollout_module, "sample_task_state", None)
+        policy_factory = getattr(rollout_module, "make_policy", None)
+        self.assertIsNotNone(sampler_fn)
+        self.assertIsNotNone(policy_factory)
+
+        task_state = sampler_fn(TASK_NAME)
+        self.assertEqual(list(task_state.keys()), ["socket_pos", "socket_quat", "peg_pos", "peg_quat"])
+
+        policy = policy_factory(TASK_NAME, inject_noise=False)
+        self.assertEqual(policy.__class__.__name__, "TestAirInsertPolicy")
+
+    def test_real_headless_smoke_instantiates_resets_and_steps_new_task_once(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = act_ex_utils.sample_air_insert_socket_peg_state()
+        env = make_sim_env(TASK_NAME, headless=True)
+        policy = policy_cls(inject_noise=False)
+
+        try:
+            env.reset(task_state)
+            action = policy.predict(task_state, 0)
+            env.step(action)
+            self.assertIsNotNone(env.obs)
+            self.assertIn("qpos", env.obs)
+            self.assertIn("images", env.obs)
+        finally:
+            env.exit_flag = True
+            cam_thread = getattr(env, "cam_thread", None)
+            if cam_thread is not None:
+                cam_thread.join(timeout=1.0)
+            viewer = getattr(env, "viewer", None)
+            if viewer is not None:
+                viewer.close()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_eval_vla_execution.py

@@ -1,5 +1,11 @@
 import unittest
+from unittest import mock
 
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+
+from roboimi.demos.vla_scripts import eval_vla
 from roboimi.vla.eval_utils import execute_policy_action
 
 
@@ -14,6 +20,48 @@ class _FakeEnv:
         self.calls.append(("step_jnt", action))
 
 
+class _FakeQueue:
+    def __init__(self, initial_items=None):
+        self.items = list(initial_items or [])
+        self.put_calls = []
+
+    def put(self, item):
+        self.put_calls.append(item)
+        self.items.append(item)
+
+    def get(self, timeout=None):
+        del timeout
+        if not self.items:
+            raise AssertionError("queue unexpectedly empty")
+        return self.items.pop(0)
+
+
+def _make_parallel_cfg(**eval_overrides):
+    eval_cfg = {
+        "ckpt_path": "checkpoints/vla_model_best.pt",
+        "num_episodes": 5,
+        "num_workers": 2,
+        "max_timesteps": 1,
+        "device": "cpu",
+        "task_name": "sim_transfer",
+        "camera_names": ["front"],
+        "use_smoothing": False,
+        "smooth_alpha": 0.3,
+        "verbose_action": False,
+        "headless": True,
+        "artifact_dir": None,
+        "save_artifacts": False,
+        "save_summary_json": False,
+        "save_timing": False,
+        "save_trajectory": False,
+        "save_trajectory_npz": False,
+        "record_video": False,
+        "save_trajectory_image": False,
+    }
+    eval_cfg.update(eval_overrides)
+    return OmegaConf.create({"agent": {}, "eval": eval_cfg})
+
+
 class EvalVLAExecutionTest(unittest.TestCase):
     def test_execute_policy_action_uses_ee_step(self):
         env = _FakeEnv()
@@ -23,6 +71,821 @@ class EvalVLAExecutionTest(unittest.TestCase):
 
         self.assertEqual(env.calls, [("step", action)])
 
+    def test_split_episode_indices_balances_workers(self):
+        self.assertEqual(
+            eval_vla._split_episode_indices(num_episodes=10, num_workers=3),
+            [[0, 1, 2, 3], [4, 5, 6], [7, 8, 9]],
+        )
+
+    def test_normalize_num_workers_caps_worker_count_to_episode_count(self):
+        self.assertEqual(eval_vla._normalize_num_workers(num_workers=5, num_episodes=2), 2)
+
+    def test_plan_episode_box_poses_uses_global_episode_order(self):
+        planned_poses = [
+            np.array([0.1, 0.2, 0.3], dtype=np.float32),
+            np.array([1.1, 1.2, 1.3], dtype=np.float32),
+            np.array([2.1, 2.2, 2.3], dtype=np.float32),
+        ]
+        sampler = mock.Mock(side_effect=planned_poses)
+
+        result = eval_vla._plan_episode_box_poses(num_episodes=3, sampler=sampler)
+
+        self.assertEqual(sampler.call_count, 3)
+        self.assertEqual(len(result), 3)
+        for expected, actual in zip(planned_poses, result):
+            np.testing.assert_array_equal(actual, expected)
+
+    def test_resolve_policy_camera_names_matches_vlaagent_fallback_sorting(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "_target_": "roboimi.vla.agent.VLAAgent",
+                },
+                "eval": {
+                    "camera_names": ["r_vis", "top", "front"],
+                },
+            }
+        )
+
+        self.assertEqual(
+            eval_vla._resolve_policy_camera_names(cfg),
+            ["front", "r_vis", "top"],
+        )
+
+    def test_resolve_policy_camera_names_matches_gr00t_fallback_input_order(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "_target_": "roboimi.vla.agent_gr00t_dit.VLAAgentGr00tDiT",
+                },
+                "eval": {
+                    "camera_names": ["r_vis", "top", "front"],
+                },
+            }
+        )
+
+        self.assertEqual(
+            eval_vla._resolve_policy_camera_names(cfg),
+            ["r_vis", "top", "front"],
+        )
+
+    def test_resolve_eval_image_resize_shape_prefers_agent_top_level_override(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "eval_image_resize_shape": None,
+                    "condition_encoder": {
+                        "eval_image_resize_shape": [256, 256],
+                    },
+                },
+                "data": {
+                    "image_resize_shape": [224, 224],
+                },
+            }
+        )
+
+        self.assertIsNone(eval_vla._resolve_eval_image_resize_shape(cfg))
+
+    def test_resolve_eval_image_resize_shape_prefers_condition_encoder_override(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "condition_encoder": {
+                        "eval_image_resize_shape": None,
+                    },
+                },
+                "data": {
+                    "image_resize_shape": [224, 224],
+                },
+            }
+        )
+
+        self.assertIsNone(eval_vla._resolve_eval_image_resize_shape(cfg))
+
+    def test_resolve_eval_image_resize_shape_prefers_vision_backbone_override(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "vision_backbone": {
+                        "eval_image_resize_shape": [256, 256],
+                    },
+                },
+                "data": {
+                    "image_resize_shape": [224, 224],
+                },
+            }
+        )
+
+        self.assertEqual(eval_vla._resolve_eval_image_resize_shape(cfg), (256, 256))
+
+    def test_build_episode_plans_without_box_poses_keeps_serial_sampling_lazy(self):
+        plans = eval_vla._build_episode_plans(num_episodes=3)
+
+        self.assertEqual(
+            plans,
+            [
+                {"episode_index": 0},
+                {"episode_index": 1},
+                {"episode_index": 2},
+            ],
+        )
+
+    def test_prepare_local_policy_batch_pads_latest_observation_to_obs_horizon(self):
+        queues = eval_vla._new_local_policy_queues(obs_horizon=3)
+        observation = {
+            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
+            "images": {
+                "front": torch.tensor([[[1.0]]], dtype=torch.float32),
+            },
+        }
+
+        eval_vla._populate_local_policy_queues(queues, observation)
+        batch = eval_vla._prepare_local_policy_batch(
+            queues,
+            obs_horizon=3,
+            camera_names=["front"],
+        )
+
+        self.assertEqual(tuple(batch["qpos"].shape), (1, 3, 2))
+        self.assertEqual(tuple(batch["images"]["front"].shape), (1, 3, 1, 1, 1))
+        np.testing.assert_array_equal(
+            batch["qpos"][0].cpu().numpy(),
+            np.array([[1.0, 2.0], [1.0, 2.0], [1.0, 2.0]], dtype=np.float32),
+        )
+        np.testing.assert_array_equal(
+            batch["images"]["front"][0].cpu().numpy(),
+            np.array([[[[1.0]]], [[[1.0]]], [[[1.0]]]], dtype=np.float32),
+        )
+
+    def test_prepare_local_policy_batch_keeps_latest_variable_task_when_present(self):
+        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
+        first_observation = {
+            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
+            "images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
+            "task": "pick the red cube",
+        }
+        second_observation = {
+            "qpos": torch.tensor([3.0, 4.0], dtype=torch.float32),
+            "images": {"front": torch.tensor([[[2.0]]], dtype=torch.float32)},
+            "task": "insert the peg into the socket",
+        }
+
+        eval_vla._populate_local_policy_queues(queues, first_observation)
+        eval_vla._populate_local_policy_queues(queues, second_observation)
+        batch = eval_vla._prepare_local_policy_batch(
+            queues,
+            obs_horizon=2,
+            camera_names=["front"],
+        )
+
+        self.assertEqual(batch["task"], ["insert the peg into the socket"])
+
+    def test_prepare_local_policy_batch_omits_task_for_legacy_observations(self):
+        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
+        observation = {
+            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
+            "images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
+        }
+
+        eval_vla._populate_local_policy_queues(queues, observation)
+        batch = eval_vla._prepare_local_policy_batch(
+            queues,
+            obs_horizon=2,
+            camera_names=["front"],
+        )
+
+        self.assertNotIn("task", batch)
+
+    def test_serialize_deserialize_policy_batch_preserves_task(self):
+        batch = {
+            "qpos": torch.zeros(1, 2, 2, dtype=torch.float32),
+            "images": {"front": torch.zeros(1, 2, 1, 1, 1, dtype=torch.float32)},
+            "task": ["pick the red cube", "insert the peg into the socket"],
+        }
+
+        serialized = eval_vla._serialize_policy_batch(batch)
+        deserialized = eval_vla._deserialize_policy_batch(serialized, device="cpu")
+
+        self.assertEqual(serialized["task"], batch["task"])
+        self.assertEqual(deserialized["task"], batch["task"])
+
+    def test_enqueue_predicted_actions_uses_executable_slice(self):
+        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
+        predicted_actions = torch.tensor(
+            [[[10.0], [20.0], [30.0], [40.0]]],
+            dtype=torch.float32,
+        )
+
+        eval_vla._enqueue_predicted_actions(
+            queues,
+            predicted_actions=predicted_actions,
+            obs_horizon=2,
+            num_action_steps=2,
+        )
+
+        self.assertEqual(len(queues["action"]), 2)
+        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([20.0], dtype=np.float32))
+        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([30.0], dtype=np.float32))
+
+    def test_enqueue_predicted_actions_honors_explicit_chunk_start(self):
+        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
+        predicted_actions = torch.tensor(
+            [[[10.0], [20.0], [30.0], [40.0]]],
+            dtype=torch.float32,
+        )
+
+        eval_vla._enqueue_predicted_actions(
+            queues,
+            predicted_actions=predicted_actions,
+            obs_horizon=2,
+            num_action_steps=2,
+            action_chunk_start=0,
+        )
+
+        self.assertEqual(len(queues["action"]), 2)
+        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([10.0], dtype=np.float32))
+        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([20.0], dtype=np.float32))
+
+    def test_remote_policy_runner_only_requests_server_inference_when_local_action_queue_is_empty(self):
+        request_queue = _FakeQueue()
+        response_queue = _FakeQueue(
+            [
+                {
+                    "type": "predict_chunk_result",
+                    "actions": np.asarray([[[10.0], [20.0], [30.0]]], dtype=np.float32),
+                }
+            ]
+        )
+        runner = eval_vla._RemotePolicyRunner(
+            worker_index=3,
+            server_index=1,
+            request_queue=request_queue,
+            response_queue=response_queue,
+            camera_names=["front"],
+            obs_horizon=2,
+            num_action_steps=2,
+            action_chunk_start=0,
+        )
+        first_observation = {
+            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
+            "images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
+        }
+        second_observation = {
+            "qpos": torch.tensor([3.0, 4.0], dtype=torch.float32),
+            "images": {"front": torch.tensor([[[2.0]]], dtype=torch.float32)},
+        }
+
+        first_action, first_forward = runner.select_action(
+            first_observation,
+            episode_index=7,
+            timestep=0,
+        )
+        second_action, second_forward = runner.select_action(
+            second_observation,
+            episode_index=7,
+            timestep=1,
+        )
+
+        self.assertTrue(first_forward)
+        self.assertFalse(second_forward)
+        self.assertEqual(len(request_queue.put_calls), 1)
+        self.assertEqual(request_queue.put_calls[0]["type"], "predict_chunk")
+        self.assertEqual(request_queue.put_calls[0]["worker_index"], 3)
+        self.assertEqual(request_queue.put_calls[0]["server_index"], 1)
+        np.testing.assert_array_equal(first_action.numpy(), np.array([10.0], dtype=np.float32))
+        np.testing.assert_array_equal(second_action.numpy(), np.array([20.0], dtype=np.float32))
+
+    def test_remote_eval_worker_passes_agent_action_chunk_start_to_remote_runner(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "obs_horizon": 2,
+                    "num_action_steps": 2,
+                    "action_chunk_start": 0,
+                    "camera_names": ["front"],
+                },
+                "eval": {
+                    "obs_horizon": 2,
+                    "num_queries": 2,
+                    "response_timeout_s": 3.0,
+                    "camera_names": ["front"],
+                },
+            }
+        )
+        captured = {}
+
+        class CapturingRunner:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        with mock.patch.object(eval_vla, "_RemotePolicyRunner", CapturingRunner), \
+             mock.patch.object(eval_vla, "_run_eval_episode_plans", return_value={"ok": True}) as run_plans:
+            result = eval_vla._run_remote_eval_worker(
+                cfg,
+                episode_plans=[{"episode_index": 0}],
+                worker_index=1,
+                server_index=2,
+                request_queue=_FakeQueue(),
+                response_queue=_FakeQueue(),
+            )
+
+        self.assertEqual(result, {"ok": True})
+        self.assertEqual(captured["action_chunk_start"], 0)
+        run_plans.assert_called_once()
+
+    def test_merge_worker_summaries_sorts_episodes_and_recomputes_aggregates(self):
+        worker_summaries = [
+            {
+                "avg_inference_fps": 999.0,
+                "avg_control_fps": 999.0,
+                "avg_obs_read_time_ms": 999.0,
+                "avg_total_time_ms": 999.0,
+                "timing_summary": {"count": 999, "model_forward_count": 999},
+                "episodes": [
+                    {
+                        "episode_index": 2,
+                        "episode_reward": 9.0,
+                        "episode_max_reward": 4.0,
+                        "inference_fps": 30.0,
+                        "control_fps": 15.0,
+                    }
+                ],
+                "_merge_state": {
+                    "obs_read_time_ms": [9.0],
+                    "preprocess_time_ms": [1.0],
+                    "inference_time_ms": [3.0],
+                    "env_step_time_ms": [4.0],
+                    "total_time_ms": [10.0],
+                    "model_forward_flags": [False],
+                },
+            },
+            {
+                "avg_inference_fps": 888.0,
+                "avg_control_fps": 888.0,
+                "avg_obs_read_time_ms": 888.0,
+                "avg_total_time_ms": 888.0,
+                "timing_summary": {"count": 888, "model_forward_count": 888},
+                "episodes": [
+                    {
+                        "episode_index": 1,
+                        "episode_reward": 6.0,
+                        "episode_max_reward": 3.0,
+                        "inference_fps": 20.0,
+                        "control_fps": 10.0,
+                    },
+                    {
+                        "episode_index": 0,
+                        "episode_reward": 5.0,
+                        "episode_max_reward": 2.0,
+                        "inference_fps": 10.0,
+                        "control_fps": 5.0,
+                    },
+                ],
+                "_merge_state": {
+                    "obs_read_time_ms": [1.0, 2.0, 12.0],
+                    "preprocess_time_ms": [2.0, 3.0, 4.0],
+                    "inference_time_ms": [4.0, 5.0, 6.0],
+                    "env_step_time_ms": [6.0, 7.0, 8.0],
+                    "total_time_ms": [8.0, 9.0, 20.0],
+                    "model_forward_flags": [True, False, True],
+                },
+            },
+        ]
+        artifact_paths = {
+            "output_dir": "/tmp/merged",
+            "summary_json": "/tmp/merged/rollout_summary.json",
+            "timing_json": "/tmp/merged/timing.json",
+            "trajectory_npz": None,
+            "video_mp4": None,
+            "video_camera_name": None,
+        }
+
+        merged = eval_vla._merge_worker_summaries(worker_summaries, artifact_paths)
+
+        self.assertEqual([episode["episode_index"] for episode in merged["episodes"]], [0, 1, 2])
+        self.assertEqual(merged["episode_rewards"], [5.0, 6.0, 9.0])
+        self.assertEqual(merged["episode_max_rewards"], [2.0, 3.0, 4.0])
+        self.assertAlmostEqual(merged["avg_reward"], 20.0 / 3.0)
+        self.assertAlmostEqual(merged["avg_max_reward"], 3.0)
+        self.assertAlmostEqual(merged["avg_inference_fps"], 20.0)
+        self.assertAlmostEqual(merged["avg_control_fps"], 10.0)
+        self.assertAlmostEqual(merged["avg_obs_read_time_ms"], 6.0)
+        self.assertAlmostEqual(merged["avg_total_time_ms"], 47.0 / 4.0)
+        self.assertEqual(merged["timing_summary"]["count"], 4)
+        self.assertEqual(merged["timing_summary"]["model_forward_count"], 2)
+        self.assertEqual(merged["artifact_dir"], "/tmp/merged")
+        self.assertEqual(merged["artifacts"], artifact_paths)
+
+    def test_build_cuda_server_payloads_uses_round_robin_worker_assignment(self):
+        cfg = _make_parallel_cfg(num_episodes=4, num_workers=4, device="cuda", cuda_devices=[0, 1])
+        artifact_paths = {"output_dir": None}
+
+        with mock.patch.object(
+            eval_vla,
+            "sample_transfer_pose",
+            side_effect=[
+                np.array([0.1, 0.2, 0.3], dtype=np.float32),
+                np.array([0.4, 0.5, 0.6], dtype=np.float32),
+                np.array([0.7, 0.8, 0.9], dtype=np.float32),
+                np.array([1.0, 1.1, 1.2], dtype=np.float32),
+            ],
+        ):
+            worker_payloads, _ = eval_vla._build_parallel_worker_payloads(cfg, artifact_paths)
+
+        server_payloads, assigned_workers = eval_vla._build_cuda_server_payloads(
+            cfg,
+            worker_payloads=worker_payloads,
+            cuda_devices=[0, 1],
+        )
+
+        self.assertEqual([payload["device_index"] for payload in server_payloads], [0, 1])
+        self.assertEqual([payload["worker_index"] for payload in assigned_workers], [0, 1, 2, 3])
+        self.assertEqual([payload["server_index"] for payload in assigned_workers], [0, 1, 0, 1])
+        self.assertEqual(server_payloads[0]["worker_indices"], [0, 2])
+        self.assertEqual(server_payloads[1]["worker_indices"], [1, 3])
+
+    def test_run_eval_parallel_dispatches_episode_splits_and_box_poses(self):
+        cfg = _make_parallel_cfg(num_episodes=5, num_workers=2, artifact_dir="/tmp/parallel-root")
+        planned_poses = [
+            np.array([float(index), float(index) + 0.1, float(index) + 0.2], dtype=np.float32)
+            for index in range(5)
+        ]
+        observed_payloads = []
+
+        def fake_run_spawn_jobs(payloads, max_workers, worker_fn):
+            del worker_fn
+            self.assertEqual(max_workers, 2)
+            observed_payloads.extend(payloads)
+            return [
+                {
+                    "episodes": [
+                        {
+                            "episode_index": 4,
+                            "episode_reward": 5.0,
+                            "episode_max_reward": 5.0,
+                            "inference_fps": 50.0,
+                            "control_fps": 25.0,
+                        },
+                        {
+                            "episode_index": 3,
+                            "episode_reward": 4.0,
+                            "episode_max_reward": 4.0,
+                            "inference_fps": 40.0,
+                            "control_fps": 20.0,
+                        },
+                    ],
+                    "_merge_state": {
+                        "obs_read_time_ms": [4.0, 5.0],
+                        "preprocess_time_ms": [1.0, 1.0],
+                        "inference_time_ms": [2.0, 2.0],
+                        "env_step_time_ms": [3.0, 3.0],
+                        "total_time_ms": [4.0, 5.0],
+                        "model_forward_flags": [True, True],
+                    },
+                },
+                {
+                    "episodes": [
+                        {
+                            "episode_index": 2,
+                            "episode_reward": 3.0,
+                            "episode_max_reward": 3.0,
+                            "inference_fps": 30.0,
+                            "control_fps": 15.0,
+                        },
+                        {
+                            "episode_index": 1,
+                            "episode_reward": 2.0,
+                            "episode_max_reward": 2.0,
+                            "inference_fps": 20.0,
+                            "control_fps": 10.0,
+                        },
+                        {
+                            "episode_index": 0,
+                            "episode_reward": 1.0,
+                            "episode_max_reward": 1.0,
+                            "inference_fps": 10.0,
+                            "control_fps": 5.0,
+                        },
+                    ],
+                    "_merge_state": {
+                        "obs_read_time_ms": [1.0, 2.0, 3.0],
+                        "preprocess_time_ms": [1.0, 1.0, 1.0],
+                        "inference_time_ms": [2.0, 2.0, 2.0],
+                        "env_step_time_ms": [3.0, 3.0, 3.0],
+                        "total_time_ms": [1.0, 2.0, 3.0],
+                        "model_forward_flags": [False, True, False],
+                    },
+                },
+            ]
+
+        with mock.patch.object(
+            eval_vla,
+            "sample_transfer_pose",
+            side_effect=planned_poses,
+        ), mock.patch.object(
+            eval_vla,
+            "_run_spawn_jobs",
+            side_effect=fake_run_spawn_jobs,
+        ):
+            summary = eval_vla._run_eval_parallel(cfg)
+
+        self.assertEqual(len(observed_payloads), 2)
+        self.assertEqual(
+            [[plan["episode_index"] for plan in payload["episode_plans"]] for payload in observed_payloads],
+            [[0, 1, 2], [3, 4]],
+        )
+        for payload in observed_payloads:
+            for plan in payload["episode_plans"]:
+                np.testing.assert_array_equal(
+                    np.asarray(plan["box_pos"], dtype=np.float32),
+                    planned_poses[plan["episode_index"]],
+                )
+        self.assertEqual([episode["episode_index"] for episode in summary["episodes"]], [0, 1, 2, 3, 4])
+        self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0, 4.0, 5.0])
+        self.assertEqual(summary["num_episodes"], 5)
+
+    def test_run_eval_parallel_allows_trajectory_images_and_keeps_worker_artifact_paths(self):
+        cfg = _make_parallel_cfg(
+            num_episodes=2,
+            num_workers=2,
+            artifact_dir="/tmp/parallel-images",
+            save_summary_json=True,
+            save_trajectory_image=True,
+        )
+        observed_payloads = []
+
+        def fake_run_spawn_jobs(payloads, max_workers, worker_fn):
+            del worker_fn
+            self.assertEqual(max_workers, 2)
+            observed_payloads.extend(payloads)
+            return [
+                {
+                    "episodes": [
+                        {
+                            "episode_index": 0,
+                            "episode_reward": 1.0,
+                            "episode_max_reward": 1.0,
+                            "inference_fps": 10.0,
+                            "control_fps": 5.0,
+                            "artifact_paths": {
+                                "trajectory_image": f"{payloads[0]['artifact_dir']}/rollout_front_ep01_trajectory.png",
+                            },
+                        },
+                    ],
+                    "_merge_state": {
+                        "obs_read_time_ms": [1.0],
+                        "preprocess_time_ms": [1.0],
+                        "inference_time_ms": [1.0],
+                        "env_step_time_ms": [1.0],
+                        "total_time_ms": [1.0],
+                        "model_forward_flags": [True],
+                    },
+                },
+                {
+                    "episodes": [
+                        {
+                            "episode_index": 1,
+                            "episode_reward": 2.0,
+                            "episode_max_reward": 2.0,
+                            "inference_fps": 20.0,
+                            "control_fps": 10.0,
+                            "artifact_paths": {
+                                "trajectory_image": f"{payloads[1]['artifact_dir']}/rollout_front_ep02_trajectory.png",
+                            },
+                        },
+                    ],
+                    "_merge_state": {
+                        "obs_read_time_ms": [2.0],
+                        "preprocess_time_ms": [2.0],
+                        "inference_time_ms": [2.0],
+                        "env_step_time_ms": [2.0],
+                        "total_time_ms": [2.0],
+                        "model_forward_flags": [False],
+                    },
+                },
+            ]
+
+        with mock.patch.object(
+            eval_vla,
+            "sample_transfer_pose",
+            side_effect=[
+                np.array([0.1, 0.2, 0.3], dtype=np.float32),
+                np.array([0.4, 0.5, 0.6], dtype=np.float32),
+            ],
+        ), mock.patch.object(
+            eval_vla,
+            "_run_spawn_jobs",
+            side_effect=fake_run_spawn_jobs,
+        ):
+            summary = eval_vla._run_eval_parallel(cfg)
+
+        self.assertEqual(len(observed_payloads), 2)
+        self.assertTrue(observed_payloads[0]["artifact_dir"].endswith("workers/worker_00"))
+        self.assertTrue(observed_payloads[1]["artifact_dir"].endswith("workers/worker_01"))
+        self.assertTrue(
+            summary["episodes"][0]["artifact_paths"]["trajectory_image"].endswith(
+                "workers/worker_00/rollout_front_ep01_trajectory.png"
+            )
+        )
+        self.assertTrue(
+            summary["episodes"][1]["artifact_paths"]["trajectory_image"].endswith(
+                "workers/worker_01/rollout_front_ep02_trajectory.png"
+            )
+        )
+
+    def test_run_eval_parallel_surfaces_worker_failures(self):
+        cfg = _make_parallel_cfg(num_episodes=2, num_workers=2)
+
+        with mock.patch.object(
+            eval_vla,
+            "sample_transfer_pose",
+            side_effect=[
+                np.array([0.1, 0.2, 0.3], dtype=np.float32),
+                np.array([0.4, 0.5, 0.6], dtype=np.float32),
+            ],
+        ), mock.patch.object(
+            eval_vla,
+            "_run_spawn_jobs",
+            side_effect=RuntimeError("boom"),
+        ):
+            with self.assertRaisesRegex(RuntimeError, "Parallel rollout worker failed"):
+                eval_vla._run_eval_parallel(cfg)
+
+    def test_run_eval_parallel_cuda_builds_server_payloads_and_merges_worker_results(self):
+        cfg = _make_parallel_cfg(
+            num_episodes=4,
+            num_workers=4,
+            device="cuda",
+            cuda_devices=[0],
+            artifact_dir="/tmp/cuda-root",
+        )
+        observed_server_payloads = []
+        observed_worker_payloads = []
+
+        def fake_run_cuda_parallel_processes(server_payloads, worker_payloads):
+            observed_server_payloads.extend(server_payloads)
+            observed_worker_payloads.extend(worker_payloads)
+            return [
+                {
+                    "episodes": [
+                        {
+                            "episode_index": 2,
+                            "episode_reward": 3.0,
+                            "episode_max_reward": 3.0,
+                            "inference_fps": 30.0,
+                            "control_fps": 15.0,
+                        },
+                        {
+                            "episode_index": 0,
+                            "episode_reward": 1.0,
+                            "episode_max_reward": 1.0,
+                            "inference_fps": 10.0,
+                            "control_fps": 5.0,
+                        },
+                    ],
+                    "_merge_state": {
+                        "obs_read_time_ms": [1.0, 2.0],
+                        "preprocess_time_ms": [1.0, 1.0],
+                        "inference_time_ms": [2.0, 2.0],
+                        "env_step_time_ms": [3.0, 3.0],
+                        "total_time_ms": [4.0, 4.0],
+                        "model_forward_flags": [True, False],
+                    },
+                },
+                {
+                    "episodes": [
+                        {
+                            "episode_index": 3,
+                            "episode_reward": 4.0,
+                            "episode_max_reward": 4.0,
+                            "inference_fps": 40.0,
+                            "control_fps": 20.0,
+                        },
+                        {
+                            "episode_index": 1,
+                            "episode_reward": 2.0,
+                            "episode_max_reward": 2.0,
+                            "inference_fps": 20.0,
+                            "control_fps": 10.0,
+                        },
+                    ],
+                    "_merge_state": {
+                        "obs_read_time_ms": [3.0, 4.0],
+                        "preprocess_time_ms": [1.0, 1.0],
+                        "inference_time_ms": [2.0, 2.0],
+                        "env_step_time_ms": [3.0, 3.0],
+                        "total_time_ms": [4.0, 4.0],
+                        "model_forward_flags": [True, True],
+                    },
+                },
+            ]
+
+        with mock.patch.object(
+            eval_vla,
+            "sample_transfer_pose",
+            side_effect=[
+                np.array([0.1, 0.2, 0.3], dtype=np.float32),
+                np.array([0.4, 0.5, 0.6], dtype=np.float32),
+                np.array([0.7, 0.8, 0.9], dtype=np.float32),
+                np.array([1.0, 1.1, 1.2], dtype=np.float32),
+            ],
+        ), mock.patch.object(
+            eval_vla,
+            "_run_cuda_parallel_processes",
+            side_effect=fake_run_cuda_parallel_processes,
+            create=True,
+        ):
+            summary = eval_vla._run_eval_parallel_cuda(cfg)
+
+        self.assertEqual(len(observed_server_payloads), 1)
+        self.assertEqual(observed_server_payloads[0]["device_index"], 0)
+        self.assertEqual(len(observed_worker_payloads), 4)
+        self.assertTrue(all(payload["server_index"] == 0 for payload in observed_worker_payloads))
+        self.assertEqual([episode["episode_index"] for episode in summary["episodes"]], [0, 1, 2, 3])
+        self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0, 4.0])
+        self.assertEqual(summary["num_episodes"], 4)
+
+    def test_run_eval_parallel_cuda_surfaces_server_failures(self):
+        cfg = _make_parallel_cfg(num_episodes=2, num_workers=2, device="cuda", cuda_devices=[0])
+
+        with mock.patch.object(
+            eval_vla,
+            "sample_transfer_pose",
+            side_effect=[
+                np.array([0.1, 0.2, 0.3], dtype=np.float32),
+                np.array([0.4, 0.5, 0.6], dtype=np.float32),
+            ],
+        ), mock.patch.object(
+            eval_vla,
+            "_run_cuda_parallel_processes",
+            side_effect=RuntimeError("server boom"),
+            create=True,
+        ):
+            with self.assertRaisesRegex(RuntimeError, "Parallel CUDA rollout failed"):
+                eval_vla._run_eval_parallel_cuda(cfg)
+
+    def test_run_spawn_jobs_supports_real_spawn_with_actual_eval_worker_entry(self):
+        payloads = [
+            {"_spawn_probe": True, "probe_value": 1, "worker_index": 0},
+            {"_spawn_probe": True, "probe_value": 2, "worker_index": 1},
+        ]
+
+        results = eval_vla._run_spawn_jobs(
+            payloads=payloads,
+            max_workers=2,
+            worker_fn=eval_vla._run_eval_worker_entry,
+        )
+
+        self.assertEqual(sorted(result["probe_value"] for result in results), [1, 2])
+        self.assertEqual(sorted(result["worker_index"] for result in results), [0, 1])
+
+    def test_cuda_server_and_env_worker_entrypoints_support_real_spawn_probe(self):
+        ctx = eval_vla.multiprocessing.get_context("spawn")
+        request_queue = ctx.Queue()
+        response_queue = ctx.Queue()
+        result_queue = ctx.Queue()
+
+        server = ctx.Process(
+            target=eval_vla._inference_server_main,
+            args=(
+                {
+                    "_spawn_probe": True,
+                    "server_index": 0,
+                    "request_queue": request_queue,
+                    "response_queues": [response_queue],
+                },
+            ),
+        )
+        worker = ctx.Process(
+            target=eval_vla._env_worker_main,
+            args=(
+                {
+                    "_spawn_probe": True,
+                    "worker_index": 0,
+                    "server_index": 0,
+                    "request_queue": request_queue,
+                    "response_queue": response_queue,
+                    "result_queue": result_queue,
+                },
+            ),
+        )
+
+        server.start()
+        worker.start()
+
+        result = result_queue.get(timeout=10.0)
+
+        worker.join(timeout=10.0)
+        request_queue.put({"type": "shutdown_server"})
+        server.join(timeout=10.0)
+
+        self.assertEqual(result["kind"], "worker_result")
+        self.assertEqual(result["summary"]["probe_worker_index"], 0)
+        self.assertEqual(result["summary"]["probe_server_index"], 0)
+        self.assertEqual(result["summary"]["probe_actions"], [[[11.0], [22.0], [33.0]]])
+        self.assertEqual(worker.exitcode, 0)
+        self.assertEqual(server.exitcode, 0)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_eval_vla_headless.py

@@ -15,6 +15,7 @@ class _FakeAgent:
     def __init__(self):
         self.reset_calls = 0
         self.last_observation = None
+        self.observation_shapes = []
 
     def eval(self):
         return self
@@ -27,6 +28,7 @@ class _FakeAgent:
 
     def select_action(self, observation):
         self.last_observation = observation
+        self.observation_shapes.append(tuple(observation["images"]["front"].shape))
         return torch.zeros(16)
 
 
@@ -36,8 +38,8 @@ class _FakeEnv:
         self.render_calls = 0
         self.reset_calls = []
 
-    def reset(self, box_pos):
-        self.reset_calls.append(np.array(box_pos))
+    def reset(self, task_state):
+        self.reset_calls.append(task_state)
 
     def _get_image_obs(self):
         self.image_obs_calls += 1
@@ -74,7 +76,7 @@ class _FakeRenderer:
         self._env = env
         self._frames = [
             np.full((4, 4, 3), fill_value=index, dtype=np.uint8)
-            for index in range(5)
+            for index in range(8)
         ]
         self._index = 0
 
@@ -108,6 +110,23 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertEqual(tuple(prepared["images"]["front"].shape), (3, 8, 8))
         self.assertEqual(tuple(prepared["qpos"].shape), (16,))
 
+    def test_prepare_observation_preserves_task_when_present(self):
+        obs = {
+            "images": {
+                "front": np.zeros((4, 4, 3), dtype=np.uint8),
+            },
+            "qpos": np.zeros(16, dtype=np.float32),
+            "task": "insert the peg into the socket",
+        }
+
+        prepared = eval_vla.prepare_observation(
+            obs,
+            ["front"],
+            image_resize_shape=None,
+        )
+
+        self.assertEqual(prepared["task"], "insert the peg into the socket")
+
     def test_headless_eval_sets_mujoco_gl_to_egl_when_display_missing(self):
         cfg = OmegaConf.create({"eval": {"headless": True}})
         with mock.patch.dict(eval_vla.os.environ, {}, clear=True):
@@ -125,6 +144,8 @@ class EvalVLAHeadlessTest(unittest.TestCase):
 
         self.assertIn("headless", eval_cfg)
         self.assertFalse(eval_cfg.headless)
+        self.assertIn("task_description", eval_cfg)
+        self.assertIsNone(eval_cfg.task_description)
 
     def test_make_sim_env_accepts_headless_and_disables_render(self):
         fake_env = object()
@@ -144,7 +165,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
             is_render=False,
             control_freq=30,
             is_interpolate=True,
-            cam_view="angle",
+            cam_view="top",
         )
 
     def test_headless_sync_camera_capture_populates_images_without_gui_calls(self):
@@ -153,11 +174,10 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         env.mj_data = object()
         env.exit_flag = False
         env.is_render = False
-        env.cam = 'angle'
+        env.cam = 'top'
         env.r_vis = None
         env.l_vis = None
         env.top = None
-        env.angle = None
         env.front = None
         env._offscreen_renderer = None
 
@@ -176,7 +196,6 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(env.r_vis)
         self.assertIsNotNone(env.l_vis)
         self.assertIsNotNone(env.top)
-        self.assertIsNotNone(env.angle)
         self.assertIsNotNone(env.front)
 
     def test_cam_start_skips_background_thread_when_headless(self):
@@ -196,11 +215,10 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         env.mj_data = object()
         env.exit_flag = False
         env.is_render = False
-        env.cam = "angle"
+        env.cam = "top"
         env.r_vis = None
         env.l_vis = None
         env.top = None
-        env.angle = None
         env.front = None
 
         with mock.patch(
@@ -217,9 +235,33 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(env.r_vis)
         self.assertIsNotNone(env.l_vis)
         self.assertIsNotNone(env.top)
-        self.assertIsNotNone(env.angle)
         self.assertIsNotNone(env.front)
 
+    def test_dual_diana_step_refreshes_obs_after_physics_step(self):
+        env = DualDianaMed.__new__(DualDianaMed)
+        env.compute_qpos = np.zeros(16)
+        env.interpolator_left = None
+        env.interpolator_right = None
+        env.control_timestep = 0.001
+        env.model_timestep = 0.001
+        env.base_time = 0.0
+        events = []
+
+        def fake_get_obs():
+            events.append("obs")
+            return {"images": {}, "qpos": np.zeros(16, dtype=np.float32)}
+
+        env._get_obs = fake_get_obs
+
+        with mock.patch(
+            "roboimi.envs.double_base.MujocoEnv.step",
+            autospec=True,
+            side_effect=lambda _self, _action: events.append("physics"),
+        ):
+            env.step(np.zeros(16))
+
+        self.assertEqual(events, ["physics", "obs"])
+
     def test_eval_main_headless_skips_render_and_still_executes_policy(self):
         fake_env = _FakeEnv()
         fake_agent = _FakeAgent()
@@ -270,6 +312,74 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(fake_agent.last_observation)
         self.assertIn("front", fake_agent.last_observation["images"])
 
+    def test_eval_main_uses_condition_encoder_resize_override(self):
+        fake_env = _FakeEnv()
+        fake_agent = _FakeAgent()
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "condition_encoder": {
+                        "eval_image_resize_shape": None,
+                    },
+                },
+                "data": {
+                    "image_resize_shape": [224, 224],
+                },
+                "eval": {
+                    "ckpt_path": "checkpoints/vla_model_best.pt",
+                    "num_episodes": 1,
+                    "max_timesteps": 1,
+                    "device": "cpu",
+                    "task_name": "sim_transfer",
+                    "camera_names": ["front"],
+                    "use_smoothing": False,
+                    "smooth_alpha": 0.3,
+                    "verbose_action": False,
+                    "headless": True,
+                },
+            }
+        )
+
+        with mock.patch.object(eval_vla, "load_checkpoint", return_value=(fake_agent, None)), \
+             mock.patch.object(eval_vla, "make_sim_env", return_value=fake_env), \
+             mock.patch.object(eval_vla, "sample_transfer_pose", return_value=np.array([0.1, 0.2, 0.3])), \
+             mock.patch.object(eval_vla, "execute_policy_action"), \
+             mock.patch.object(eval_vla, "tqdm", side_effect=lambda iterable, **kwargs: iterable):
+            eval_vla.main.__wrapped__(cfg)
+
+        self.assertEqual(fake_agent.observation_shapes, [(3, 8, 8)])
+
+    def test_eval_main_injects_configured_task_description_when_env_omits_task(self):
+        fake_env = _FakeEnv()
+        fake_agent = _FakeAgent()
+        cfg = OmegaConf.create(
+            {
+                "agent": {},
+                "eval": {
+                    "ckpt_path": "checkpoints/vla_model_best.pt",
+                    "num_episodes": 1,
+                    "max_timesteps": 1,
+                    "device": "cpu",
+                    "task_name": "sim_transfer",
+                    "task_description": "pick the red cube",
+                    "camera_names": ["front"],
+                    "use_smoothing": False,
+                    "smooth_alpha": 0.3,
+                    "verbose_action": False,
+                    "headless": True,
+                },
+            }
+        )
+
+        with mock.patch.object(eval_vla, "load_checkpoint", return_value=(fake_agent, None)), \
+             mock.patch.object(eval_vla, "make_sim_env", return_value=fake_env), \
+             mock.patch.object(eval_vla, "sample_transfer_pose", return_value=np.array([0.1, 0.2, 0.3])), \
+             mock.patch.object(eval_vla, "execute_policy_action"), \
+             mock.patch.object(eval_vla, "tqdm", side_effect=lambda iterable, **kwargs: iterable):
+            eval_vla.main.__wrapped__(cfg)
+
+        self.assertEqual(fake_agent.last_observation["task"], "pick the red cube")
+
     def test_run_eval_returns_average_reward_summary(self):
         reward_sequences = [
             [1.0, 2.0],
@@ -328,5 +438,255 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertEqual(summary["num_episodes"], 2)
 
 
+    def test_eval_config_exposes_num_workers_default(self):
+        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
+
+        self.assertIn("num_workers", eval_cfg)
+        self.assertEqual(eval_cfg.num_workers, 1)
+
+    def test_eval_config_exposes_cuda_devices_default(self):
+        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
+
+        self.assertIn("cuda_devices", eval_cfg)
+        self.assertIsNone(eval_cfg.cuda_devices)
+
+    def test_eval_config_exposes_parallel_timeout_defaults(self):
+        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
+
+        self.assertIn("response_timeout_s", eval_cfg)
+        self.assertIn("server_startup_timeout_s", eval_cfg)
+        self.assertEqual(eval_cfg.response_timeout_s, 300.0)
+        self.assertEqual(eval_cfg.server_startup_timeout_s, 300.0)
+
+    def test_run_eval_uses_serial_path_when_num_workers_is_one(self):
+        cfg = OmegaConf.create(
+            {
+                "eval": {
+                    "num_workers": 1,
+                    "num_episodes": 3,
+                }
+            }
+        )
+
+        with mock.patch.object(
+            eval_vla,
+            "_run_eval_serial",
+            return_value={"mode": "serial"},
+        ) as run_eval_serial, mock.patch.object(
+            eval_vla,
+            "_run_eval_parallel",
+        ) as run_eval_parallel:
+            result = eval_vla._run_eval(cfg)
+
+        self.assertEqual(result, {"mode": "serial"})
+        run_eval_serial.assert_called_once_with(cfg)
+        run_eval_parallel.assert_not_called()
+
+    def test_run_eval_uses_serial_path_when_requested_workers_collapse_to_one(self):
+        cfg = OmegaConf.create(
+            {
+                "eval": {
+                    "num_workers": 8,
+                    "num_episodes": 1,
+                }
+            }
+        )
+
+        with mock.patch.object(
+            eval_vla,
+            "_run_eval_serial",
+            return_value={"mode": "serial"},
+        ) as run_eval_serial, mock.patch.object(
+            eval_vla,
+            "_run_eval_parallel",
+        ) as run_eval_parallel:
+            result = eval_vla._run_eval(cfg)
+
+        self.assertEqual(result, {"mode": "serial"})
+        run_eval_serial.assert_called_once_with(cfg)
+        run_eval_parallel.assert_not_called()
+
+    def test_run_eval_parallel_requires_headless_true(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {},
+                "eval": {
+                    "ckpt_path": "checkpoints/vla_model_best.pt",
+                    "num_episodes": 2,
+                    "num_workers": 2,
+                    "max_timesteps": 1,
+                    "device": "cpu",
+                    "task_name": "sim_transfer",
+                    "camera_names": ["front"],
+                    "use_smoothing": False,
+                    "smooth_alpha": 0.3,
+                    "verbose_action": False,
+                    "headless": False,
+                },
+            }
+        )
+
+        with self.assertRaisesRegex(ValueError, "headless=true"):
+            eval_vla._run_eval_parallel(cfg)
+
+    def test_run_eval_parallel_dispatches_to_cpu_workers_when_device_is_cpu(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {},
+                "eval": {
+                    "ckpt_path": "checkpoints/vla_model_best.pt",
+                    "num_episodes": 2,
+                    "num_workers": 2,
+                    "max_timesteps": 1,
+                    "device": "cpu",
+                    "task_name": "sim_transfer",
+                    "camera_names": ["front"],
+                    "use_smoothing": False,
+                    "smooth_alpha": 0.3,
+                    "verbose_action": False,
+                    "headless": True,
+                    "cuda_devices": None,
+                },
+            }
+        )
+
+        with mock.patch.object(
+            eval_vla,
+            "_run_eval_parallel_cpu",
+            return_value={"mode": "cpu"},
+            create=True,
+        ) as run_cpu_parallel, mock.patch.object(
+            eval_vla,
+            "_run_eval_parallel_cuda",
+            create=True,
+        ) as run_cuda_parallel:
+            result = eval_vla._run_eval_parallel(cfg)
+
+        self.assertEqual(result, {"mode": "cpu"})
+        run_cpu_parallel.assert_called_once_with(cfg)
+        run_cuda_parallel.assert_not_called()
+
+    def test_run_eval_parallel_dispatches_to_cuda_servers_when_device_is_cuda(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {},
+                "eval": {
+                    "ckpt_path": "checkpoints/vla_model_best.pt",
+                    "num_episodes": 2,
+                    "num_workers": 2,
+                    "max_timesteps": 1,
+                    "device": "cuda",
+                    "task_name": "sim_transfer",
+                    "camera_names": ["front"],
+                    "use_smoothing": False,
+                    "smooth_alpha": 0.3,
+                    "verbose_action": False,
+                    "headless": True,
+                    "cuda_devices": [0],
+                },
+            }
+        )
+
+        with mock.patch.object(
+            eval_vla,
+            "_run_eval_parallel_cpu",
+            create=True,
+        ) as run_cpu_parallel, mock.patch.object(
+            eval_vla,
+            "_run_eval_parallel_cuda",
+            return_value={"mode": "cuda"},
+            create=True,
+        ) as run_cuda_parallel:
+            result = eval_vla._run_eval_parallel(cfg)
+
+        self.assertEqual(result, {"mode": "cuda"})
+        run_cpu_parallel.assert_not_called()
+        run_cuda_parallel.assert_called_once_with(cfg)
+
+    def test_resolve_cuda_devices_defaults_to_single_logical_gpu(self):
+        cfg = OmegaConf.create(
+            {
+                "device": "cuda",
+                "cuda_devices": None,
+            }
+        )
+
+        self.assertEqual(eval_vla._resolve_cuda_devices(cfg), [0])
+
+    def test_resolve_cuda_devices_rejects_empty_selection(self):
+        cfg = OmegaConf.create(
+            {
+                "device": "cuda",
+                "cuda_devices": [],
+            }
+        )
+
+        with self.assertRaisesRegex(ValueError, "cuda_devices"):
+            eval_vla._resolve_cuda_devices(cfg)
+
+    def test_run_eval_uses_air_insert_sampler_for_socket_peg_task(self):
+        self.assertTrue(
+            hasattr(eval_vla, "sample_air_insert_socket_peg_state"),
+            "Expected eval_vla to expose the new socket/peg reset sampler",
+        )
+
+        fake_env = _FakeEnv()
+        fake_agent = _FakeAgent()
+        sampled_task_state = {
+            "socket_pos": np.array([-0.10, 0.80, 0.47], dtype=np.float32),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+            "peg_pos": np.array([0.10, 0.82, 0.47], dtype=np.float32),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        }
+        cfg = OmegaConf.create(
+            {
+                "agent": {},
+                "eval": {
+                    "ckpt_path": "checkpoints/vla_model_best.pt",
+                    "num_episodes": 1,
+                    "max_timesteps": 1,
+                    "device": "cpu",
+                    "task_name": "sim_air_insert_socket_peg",
+                    "camera_names": ["front"],
+                    "use_smoothing": False,
+                    "smooth_alpha": 0.3,
+                    "verbose_action": False,
+                    "headless": True,
+                },
+            }
+        )
+
+        with mock.patch.object(
+            eval_vla,
+            "load_checkpoint",
+            return_value=(fake_agent, None),
+        ), mock.patch.object(
+            eval_vla,
+            "make_sim_env",
+            return_value=fake_env,
+        ) as make_env, mock.patch.object(
+            eval_vla,
+            "sample_air_insert_socket_peg_state",
+            return_value=sampled_task_state,
+        ) as socket_peg_sampler, mock.patch.object(
+            eval_vla,
+            "sample_transfer_pose",
+            side_effect=AssertionError("sample_transfer_pose should not be used for sim_air_insert_socket_peg"),
+        ), mock.patch.object(
+            eval_vla,
+            "execute_policy_action",
+        ) as execute_policy_action, mock.patch.object(
+            eval_vla,
+            "tqdm",
+            side_effect=lambda iterable, **kwargs: iterable,
+        ):
+            eval_vla._run_eval(cfg)
+
+        make_env.assert_called_once_with("sim_air_insert_socket_peg", headless=True)
+        socket_peg_sampler.assert_called_once_with()
+        execute_policy_action.assert_called_once()
+        self.assertEqual(fake_env.reset_calls, [sampled_task_state])
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_eval_vla_rollout_artifacts.py

@@ -102,10 +102,8 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
         self.assertIn('artifact_dir', eval_cfg)
         self.assertFalse(eval_cfg.save_summary_json)
         self.assertFalse(eval_cfg.save_trajectory_npz)
-        self.assertFalse(eval_cfg.save_trajectory_image)
         self.assertFalse(eval_cfg.record_video)
         self.assertIsNone(eval_cfg.artifact_dir)
-        self.assertIsNone(eval_cfg.trajectory_image_camera_name)
         self.assertIsNone(eval_cfg.video_camera_name)
         self.assertEqual(eval_cfg.video_fps, 30)
 
@@ -135,8 +133,6 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
                         'artifact_dir': tmpdir,
                         'save_summary_json': True,
                         'save_trajectory_npz': True,
-                        'save_trajectory_image': True,
-                        'trajectory_image_camera_name': 'front',
                         'record_video': True,
                         'video_camera_name': 'front',
                         'video_fps': 12,
@@ -180,14 +176,12 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
             trajectory_path = Path(artifacts['trajectory_npz'])
             summary_path = Path(artifacts['summary_json'])
             video_path = Path(artifacts['video_mp4'])
-            trajectory_image_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
 
             self.assertEqual(Path(artifacts['output_dir']), Path(tmpdir))
             self.assertEqual(artifacts['video_camera_name'], 'front')
             self.assertTrue(trajectory_path.exists())
             self.assertTrue(summary_path.exists())
             self.assertTrue(video_path.exists())
-            self.assertTrue(trajectory_image_path.exists())
 
             rollout_npz = np.load(trajectory_path)
             np.testing.assert_array_equal(rollout_npz['episode_index'], np.array([0, 0]))
@@ -224,120 +218,267 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
                 saved_summary = json.load(fh)
             self.assertEqual(saved_summary['artifacts']['trajectory_npz'], str(trajectory_path))
             self.assertEqual(saved_summary['artifacts']['video_mp4'], str(video_path))
-            self.assertEqual(
-                saved_summary['episodes'][0]['artifact_paths']['trajectory_image'],
-                str(trajectory_image_path),
-            )
             self.assertEqual(saved_summary['episode_rewards'], [3.0])
             self.assertAlmostEqual(summary['avg_reward'], 3.0)
             self.assertIn('avg_obs_read_time_ms', summary)
             self.assertIn('avg_env_step_time_ms', summary)
 
-    def test_run_eval_exports_front_trajectory_images_without_video_dependency(self):
-        actions = [
-            np.arange(16, dtype=np.float32),
-            np.arange(16, dtype=np.float32) + 10.0,
-            np.arange(16, dtype=np.float32) + 100.0,
-            np.arange(16, dtype=np.float32) + 110.0,
+    def test_run_eval_parallel_rejects_trajectory_and_video_exports(self):
+        unsupported_flags = [
+            "record_video",
+            "save_trajectory",
+            "save_trajectory_npz",
         ]
-        fake_agent = _FakeAgent(actions)
-        fake_env = _FakeEnv()
 
+        for flag_name in unsupported_flags:
+            with self.subTest(flag_name=flag_name):
+                cfg = OmegaConf.create(
+                    {
+                        "agent": {},
+                        "eval": {
+                            "ckpt_path": "checkpoints/vla_model_best.pt",
+                            "num_episodes": 2,
+                            "num_workers": 2,
+                            "max_timesteps": 1,
+                            "device": "cpu",
+                            "task_name": "sim_transfer",
+                            "camera_names": ["front"],
+                            "use_smoothing": False,
+                            "smooth_alpha": 0.3,
+                            "verbose_action": False,
+                            "headless": True,
+                            "save_artifacts": True,
+                            flag_name: True,
+                        },
+                    }
+                )
+
+                with self.assertRaisesRegex(ValueError, flag_name):
+                    eval_vla._run_eval_parallel(cfg)
+
+    def test_run_eval_parallel_writes_merged_summary_timing_and_worker_dirs(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             cfg = OmegaConf.create(
                 {
-                    'agent': {},
-                    'eval': {
-                        'ckpt_path': 'checkpoints/vla_model_best.pt',
-                        'num_episodes': 2,
-                        'max_timesteps': 2,
-                        'device': 'cpu',
-                        'task_name': 'sim_transfer',
-                        'camera_names': ['top', 'front'],
-                        'use_smoothing': True,
-                        'smooth_alpha': 0.5,
-                        'verbose_action': False,
-                        'headless': True,
-                        'artifact_dir': tmpdir,
-                        'save_trajectory_image': True,
-                        'record_video': False,
+                    "agent": {},
+                    "eval": {
+                        "ckpt_path": "checkpoints/vla_model_best.pt",
+                        "num_episodes": 3,
+                        "num_workers": 2,
+                        "max_timesteps": 1,
+                        "device": "cpu",
+                        "task_name": "sim_transfer",
+                        "camera_names": ["front"],
+                        "use_smoothing": False,
+                        "smooth_alpha": 0.3,
+                        "verbose_action": False,
+                        "headless": True,
+                        "artifact_dir": tmpdir,
+                        "save_summary_json": True,
+                        "save_timing": True,
                     },
                 }
             )
 
-            trajectory_image_calls = []
-
-            def fake_save_rollout_trajectory_image(
-                env,
-                output_path,
-                raw_actions,
-                camera_name,
-                *,
-                line_radius=0.004,
-                max_markers=1500,
-            ):
-                del env, line_radius, max_markers
-                trajectory_image_calls.append(
+            def fake_run_spawn_jobs(payloads, max_workers, worker_fn):
+                del max_workers, worker_fn
+                return [
                     {
-                        'output_path': output_path,
-                        'camera_name': camera_name,
-                        'raw_actions': [np.array(action, copy=True) for action in raw_actions],
-                    }
-                )
-                if output_path is None:
-                    return None
-                output_path = Path(output_path)
-                output_path.parent.mkdir(parents=True, exist_ok=True)
-                output_path.write_bytes(b'fake-png')
-                return str(output_path)
+                        "episodes": [
+                            {
+                                "episode_index": 2,
+                                "episode_reward": 3.0,
+                                "episode_max_reward": 3.0,
+                                "inference_fps": 30.0,
+                                "control_fps": 15.0,
+                            }
+                        ],
+                        "_merge_state": {
+                            "obs_read_time_ms": [3.0],
+                            "preprocess_time_ms": [1.0],
+                            "inference_time_ms": [2.0],
+                            "env_step_time_ms": [4.0],
+                            "total_time_ms": [5.0],
+                            "model_forward_flags": [True],
+                        },
+                    },
+                    {
+                        "episodes": [
+                            {
+                                "episode_index": 1,
+                                "episode_reward": 2.0,
+                                "episode_max_reward": 2.0,
+                                "inference_fps": 20.0,
+                                "control_fps": 10.0,
+                            },
+                            {
+                                "episode_index": 0,
+                                "episode_reward": 1.0,
+                                "episode_max_reward": 1.0,
+                                "inference_fps": 10.0,
+                                "control_fps": 5.0,
+                            },
+                        ],
+                        "_merge_state": {
+                            "obs_read_time_ms": [1.0, 2.0],
+                            "preprocess_time_ms": [1.0, 1.0],
+                            "inference_time_ms": [2.0, 2.0],
+                            "env_step_time_ms": [4.0, 4.0],
+                            "total_time_ms": [5.0, 5.0],
+                            "model_forward_flags": [False, True],
+                        },
+                    },
+                ]
 
             with mock.patch.object(
                 eval_vla,
-                'load_checkpoint',
-                return_value=(fake_agent, None),
+                "sample_transfer_pose",
+                side_effect=[
+                    np.array([0.1, 0.2, 0.3], dtype=np.float32),
+                    np.array([0.4, 0.5, 0.6], dtype=np.float32),
+                    np.array([0.7, 0.8, 0.9], dtype=np.float32),
+                ],
             ), mock.patch.object(
                 eval_vla,
-                'make_sim_env',
-                return_value=fake_env,
-            ), mock.patch.object(
-                eval_vla,
-                'sample_transfer_pose',
-                return_value=np.array([0.1, 0.2, 0.3], dtype=np.float32),
-            ), mock.patch.object(
-                eval_vla,
-                'tqdm',
-                side_effect=lambda iterable, **kwargs: iterable,
-            ), mock.patch.object(
-                eval_vla,
-                '_save_rollout_trajectory_image',
-                side_effect=fake_save_rollout_trajectory_image,
-            ) as save_trajectory_image_mock, mock.patch.object(
-                eval_vla,
-                '_open_video_writer',
-            ) as open_video_writer_mock:
-                summary = eval_vla._run_eval(cfg)
+                "_run_spawn_jobs",
+                side_effect=fake_run_spawn_jobs,
+            ):
+                summary = eval_vla._run_eval_parallel(cfg)
 
-            self.assertEqual(save_trajectory_image_mock.call_count, 2)
-            open_video_writer_mock.assert_not_called()
-            self.assertIsNone(summary['artifacts']['video_mp4'])
-            self.assertEqual(summary['artifacts']['trajectory_image_camera_name'], 'front')
-            self.assertEqual(
-                [call['camera_name'] for call in trajectory_image_calls],
-                ['front', 'front'],
+            summary_path = Path(tmpdir) / "rollout_summary.json"
+            timing_path = Path(tmpdir) / "timing.json"
+            worker_00_dir = Path(tmpdir) / "workers" / "worker_00"
+            worker_01_dir = Path(tmpdir) / "workers" / "worker_01"
+
+            self.assertTrue(summary_path.exists())
+            self.assertTrue(timing_path.exists())
+            self.assertTrue(worker_00_dir.is_dir())
+            self.assertTrue(worker_01_dir.is_dir())
+            self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0])
+
+            with summary_path.open("r", encoding="utf-8") as fh:
+                saved_summary = json.load(fh)
+            with timing_path.open("r", encoding="utf-8") as fh:
+                saved_timing = json.load(fh)
+
+            self.assertEqual(saved_summary["episode_rewards"], [1.0, 2.0, 3.0])
+            self.assertEqual(saved_summary["artifact_dir"], tmpdir)
+            self.assertEqual(saved_timing["count"], 3)
+            self.assertEqual(saved_timing["model_forward_count"], 2)
+
+    def test_run_eval_parallel_cuda_writes_merged_summary_timing_and_worker_dirs(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = OmegaConf.create(
+                {
+                    "agent": {},
+                    "eval": {
+                        "ckpt_path": "checkpoints/vla_model_best.pt",
+                        "num_episodes": 3,
+                        "num_workers": 2,
+                        "cuda_devices": [0],
+                        "max_timesteps": 1,
+                        "device": "cuda",
+                        "task_name": "sim_transfer",
+                        "camera_names": ["front"],
+                        "use_smoothing": False,
+                        "smooth_alpha": 0.3,
+                        "verbose_action": False,
+                        "headless": True,
+                        "artifact_dir": tmpdir,
+                        "save_summary_json": True,
+                        "save_timing": True,
+                    },
+                }
             )
 
-            first_episode_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
-            second_episode_path = Path(summary['episodes'][1]['artifact_paths']['trajectory_image'])
-            self.assertTrue(first_episode_path.exists())
-            self.assertTrue(second_episode_path.exists())
-            self.assertNotEqual(first_episode_path, second_episode_path)
-            self.assertEqual(first_episode_path.parent, Path(tmpdir))
-            self.assertEqual(second_episode_path.parent, Path(tmpdir))
+            def fake_run_cuda_parallel_processes(server_payloads, worker_payloads):
+                self.assertEqual(len(server_payloads), 1)
+                self.assertEqual(server_payloads[0]["device_index"], 0)
+                self.assertEqual([payload["server_index"] for payload in worker_payloads], [0, 0])
+                return [
+                    {
+                        "episodes": [
+                            {
+                                "episode_index": 2,
+                                "episode_reward": 3.0,
+                                "episode_max_reward": 3.0,
+                                "inference_fps": 30.0,
+                                "control_fps": 15.0,
+                            }
+                        ],
+                        "_merge_state": {
+                            "obs_read_time_ms": [3.0],
+                            "preprocess_time_ms": [1.0],
+                            "inference_time_ms": [2.0],
+                            "env_step_time_ms": [4.0],
+                            "total_time_ms": [5.0],
+                            "model_forward_flags": [True],
+                        },
+                    },
+                    {
+                        "episodes": [
+                            {
+                                "episode_index": 1,
+                                "episode_reward": 2.0,
+                                "episode_max_reward": 2.0,
+                                "inference_fps": 20.0,
+                                "control_fps": 10.0,
+                            },
+                            {
+                                "episode_index": 0,
+                                "episode_reward": 1.0,
+                                "episode_max_reward": 1.0,
+                                "inference_fps": 10.0,
+                                "control_fps": 5.0,
+                            },
+                        ],
+                        "_merge_state": {
+                            "obs_read_time_ms": [1.0, 2.0],
+                            "preprocess_time_ms": [1.0, 1.0],
+                            "inference_time_ms": [2.0, 2.0],
+                            "env_step_time_ms": [4.0, 4.0],
+                            "total_time_ms": [5.0, 5.0],
+                            "model_forward_flags": [False, True],
+                        },
+                    },
+                ]
 
-            np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][0], actions[0])
-            np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][1], actions[1])
-            np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][0], actions[2])
-            np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][1], actions[3])
+            with mock.patch.object(
+                eval_vla,
+                "sample_transfer_pose",
+                side_effect=[
+                    np.array([0.1, 0.2, 0.3], dtype=np.float32),
+                    np.array([0.4, 0.5, 0.6], dtype=np.float32),
+                    np.array([0.7, 0.8, 0.9], dtype=np.float32),
+                ],
+            ), mock.patch.object(
+                eval_vla,
+                "_run_cuda_parallel_processes",
+                side_effect=fake_run_cuda_parallel_processes,
+                create=True,
+            ):
+                summary = eval_vla._run_eval_parallel_cuda(cfg)
+
+            summary_path = Path(tmpdir) / "rollout_summary.json"
+            timing_path = Path(tmpdir) / "timing.json"
+            worker_00_dir = Path(tmpdir) / "workers" / "worker_00"
+            worker_01_dir = Path(tmpdir) / "workers" / "worker_01"
+
+            self.assertTrue(summary_path.exists())
+            self.assertTrue(timing_path.exists())
+            self.assertTrue(worker_00_dir.is_dir())
+            self.assertTrue(worker_01_dir.is_dir())
+            self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0])
+
+            with summary_path.open("r", encoding="utf-8") as fh:
+                saved_summary = json.load(fh)
+            with timing_path.open("r", encoding="utf-8") as fh:
+                saved_timing = json.load(fh)
+
+            self.assertEqual(saved_summary["episode_rewards"], [1.0, 2.0, 3.0])
+            self.assertEqual(saved_summary["artifact_dir"], tmpdir)
+            self.assertEqual(saved_timing["count"], 3)
+            self.assertEqual(saved_timing["model_forward_count"], 2)
 
 
 if __name__ == '__main__':

tests/test_robot_asset_paths.py

@@ -4,7 +4,7 @@ import unittest
 from pathlib import Path
 from unittest import mock
 
-from roboimi.assets.robots.diana_med import BiDianaMed
+from roboimi.assets.robots import diana_med
 
 
 class _FakeKDL:
@@ -24,6 +24,7 @@ class RobotAssetPathResolutionTest(unittest.TestCase):
         _FakeKDL.reset_calls = []
 
     def test_bidianamed_resolves_robot_asset_paths_independent_of_cwd(self):
+        BiDianaMed = diana_med.BiDianaMed
         repo_root = Path(__file__).resolve().parents[1]
         expected_xml = repo_root / 'roboimi/assets/models/manipulators/DianaMed/bi_diana_transfer_ee.xml'
         expected_urdf = repo_root / 'roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf'
@@ -58,6 +59,47 @@ class RobotAssetPathResolutionTest(unittest.TestCase):
         self.assertEqual({Path(path) for path in _FakeKDL.init_calls}, {expected_urdf})
         self.assertTrue(all(Path(path).is_absolute() for path in _FakeKDL.init_calls))
 
+    def test_bidianamed_socket_peg_resolves_robot_asset_paths_independent_of_cwd(self):
+        BiDianaMedSocketPeg = getattr(diana_med, 'BiDianaMedSocketPeg', None)
+        self.assertIsNotNone(
+            BiDianaMedSocketPeg,
+            'Expected roboimi.assets.robots.diana_med.BiDianaMedSocketPeg',
+        )
+
+        repo_root = Path(__file__).resolve().parents[1]
+        expected_xml = repo_root / 'roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml'
+        expected_urdf = repo_root / 'roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf'
+        xml_calls = []
+
+        def fake_from_xml_path(*, filename, assets=None):
+            xml_calls.append((filename, assets))
+            return object()
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                with mock.patch(
+                    'roboimi.assets.robots.arm_base.mujoco.MjModel.from_xml_path',
+                    side_effect=fake_from_xml_path,
+                ), mock.patch(
+                    'roboimi.assets.robots.arm_base.mujoco.MjData',
+                    return_value=object(),
+                ), mock.patch(
+                    'roboimi.assets.robots.arm_base.KDL_utils',
+                    _FakeKDL,
+                ):
+                    BiDianaMedSocketPeg()
+            finally:
+                os.chdir(previous_cwd)
+
+        self.assertEqual(len(xml_calls), 1)
+        self.assertEqual(Path(xml_calls[0][0]), expected_xml)
+        self.assertTrue(Path(xml_calls[0][0]).is_absolute())
+        self.assertGreaterEqual(len(_FakeKDL.init_calls), 2)
+        self.assertEqual({Path(path) for path in _FakeKDL.init_calls}, {expected_urdf})
+        self.assertTrue(all(Path(path).is_absolute() for path in _FakeKDL.init_calls))
+
 
 if __name__ == '__main__':
     unittest.main()

tests/test_simple_robot_dataset_episode_filter.py

@@ -0,0 +1,42 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+import h5py
+import numpy as np
+
+
+def _write_episode(path: Path, length: int = 2):
+    with h5py.File(path, 'w') as f:
+        f.create_dataset('action', data=np.zeros((length, 16), dtype=np.float32))
+        obs = f.create_group('observations')
+        obs.create_dataset('qpos', data=np.zeros((length, 16), dtype=np.float32))
+        images = obs.create_group('images')
+        for cam_name in ('l_vis', 'r_vis', 'front'):
+            images.create_dataset(cam_name, data=np.zeros((length, 4, 4, 3), dtype=np.uint8))
+
+
+class SimpleRobotDatasetEpisodeFilterTest(unittest.TestCase):
+    def test_filters_by_original_episode_indices_and_exposes_available_indices(self):
+        from roboimi.vla.data.simpe_robot_dataset import SimpleRobotDataset
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            root = Path(tmpdir)
+            _write_episode(root / 'episode_0.hdf5')
+            _write_episode(root / 'episode_2.hdf5')
+            _write_episode(root / 'episode_10.hdf5')
+
+            dataset = SimpleRobotDataset(
+                root,
+                camera_names=['l_vis', 'r_vis', 'front'],
+                image_resize_shape=None,
+                episode_indices=[10, 2],
+            )
+
+            self.assertEqual(dataset.available_episode_indices, [2, 10])
+            self.assertEqual(set(dataset.episodes.keys()), {2, 10})
+            self.assertEqual(len(dataset), 4)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_smolvla_imf_agent.py

@@ -0,0 +1,311 @@
+import contextlib
+import importlib
+import importlib.machinery
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest import mock
+
+import torch
+from hydra import compose, initialize_config_dir
+from hydra.core.global_hydra import GlobalHydra
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+from torch import nn
+
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve())
+_CAMERA_NAMES = ('r_vis', 'top', 'front')
+_MISSING = object()
+
+
+class _RecordingHead(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.scale = nn.Parameter(torch.tensor(0.5))
+        self.calls = []
+
+    @staticmethod
+    def _broadcast(value, reference):
+        while value.ndim < reference.ndim:
+            value = value.unsqueeze(-1)
+        return value
+
+    def forward(self, sample, r, t, cond=None):
+        self.calls.append({
+            'sample': sample.detach().clone(),
+            'r': r.detach().clone(),
+            't': t.detach().clone(),
+            'cond': None if cond is None else cond.detach().clone(),
+        })
+        cond_term = 0.0 if cond is None else cond.mean(dim=(1, 2), keepdim=True)
+        return self.scale * sample + self._broadcast(r, sample) + 2.0 * self._broadcast(t, sample) + cond_term
+
+
+class _TaskAwareConditionEncoder(nn.Module):
+    output_dim = 4
+    condition_sequence_length = 3
+    tokens_per_step = 3
+    joint_output_dim = 4
+    camera_names = _CAMERA_NAMES
+    num_cameras = 3
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.constructor_kwargs = dict(kwargs)
+        self.bias = nn.Parameter(torch.tensor(0.0))
+        self.calls = []
+
+    def forward(self, images, state, task=None):
+        self.calls.append({'task': task, 'state': state.detach().clone(), 'image_keys': tuple(images.keys())})
+        batch_size = state.shape[0]
+        state_last = state[:, -1, 0]
+        if isinstance(task, str):
+            task_lengths = torch.full((batch_size,), float(len(task)), dtype=state.dtype, device=state.device)
+        else:
+            task_lengths = torch.tensor([float(len(item)) for item in task], dtype=state.dtype, device=state.device)
+        image_marker = images['r_vis'][:, -1].mean(dim=(1, 2, 3))
+        token0 = torch.stack([state_last, task_lengths, image_marker, torch.ones_like(state_last)], dim=-1)
+        token1 = token0 + 1.0
+        token2 = token0 + 2.0
+        return torch.stack([token0, token1, token2], dim=1) + self.bias
+
+
+class _BF16TaskAwareConditionEncoder(_TaskAwareConditionEncoder):
+    def forward(self, images, state, task=None):
+        return super().forward(images, state, task=task).to(dtype=torch.bfloat16)
+
+
+class _StubIMFHead(nn.Module):
+    def __init__(self, input_dim, output_dim, horizon, n_obs_steps, cond_dim, **kwargs):
+        super().__init__()
+        self.constructor_kwargs = {
+            'input_dim': input_dim,
+            'output_dim': output_dim,
+            'horizon': horizon,
+            'n_obs_steps': n_obs_steps,
+            'cond_dim': cond_dim,
+            **kwargs,
+        }
+        self.proj = nn.Linear(input_dim, output_dim)
+        self.cond_obs_emb = nn.Linear(cond_dim, max(cond_dim, 1))
+
+    def forward(self, sample, r, t, cond=None):
+        return torch.zeros_like(sample)
+
+    def get_optim_groups(self, weight_decay):
+        return [
+            {'params': [self.proj.weight], 'weight_decay': weight_decay},
+            {'params': [self.proj.bias, self.cond_obs_emb.weight, self.cond_obs_emb.bias], 'weight_decay': 0.0},
+        ]
+
+
+@contextlib.contextmanager
+def _stub_optional_modules(include_head=False, include_condition_encoder=False):
+    previous = {}
+
+    def inject(name, module):
+        if name not in previous:
+            previous[name] = sys.modules.get(name, _MISSING)
+        sys.modules[name] = module
+
+    diffusers_module = types.ModuleType('diffusers')
+    schedulers_module = types.ModuleType('diffusers.schedulers')
+    ddpm_module = types.ModuleType('diffusers.schedulers.scheduling_ddpm')
+    ddim_module = types.ModuleType('diffusers.schedulers.scheduling_ddim')
+
+    class _FakeScheduler:
+        def __init__(self, num_train_timesteps=100, **kwargs):
+            self.config = types.SimpleNamespace(num_train_timesteps=num_train_timesteps)
+
+    ddpm_module.DDPMScheduler = _FakeScheduler
+    ddim_module.DDIMScheduler = _FakeScheduler
+    diffusers_module.DDPMScheduler = _FakeScheduler
+    diffusers_module.DDIMScheduler = _FakeScheduler
+    diffusers_module.schedulers = schedulers_module
+
+    try:
+        inject('diffusers', diffusers_module)
+        inject('diffusers.schedulers', schedulers_module)
+        inject('diffusers.schedulers.scheduling_ddpm', ddpm_module)
+        inject('diffusers.schedulers.scheduling_ddim', ddim_module)
+        if include_head:
+            import roboimi.vla.models.heads as heads_package
+            head_module = types.ModuleType('roboimi.vla.models.heads.imf_transformer1d')
+            head_module.IMFTransformer1D = _StubIMFHead
+            inject('roboimi.vla.models.heads.imf_transformer1d', head_module)
+            setattr(heads_package, 'imf_transformer1d', head_module)
+        if include_condition_encoder:
+            module = types.ModuleType('tests.fake_smolvla_condition_encoder')
+            module.TaskAwareConditionEncoder = _TaskAwareConditionEncoder
+            module.BF16TaskAwareConditionEncoder = _BF16TaskAwareConditionEncoder
+            inject('tests.fake_smolvla_condition_encoder', module)
+        yield
+    finally:
+        for name, old in reversed(list(previous.items())):
+            if old is _MISSING:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = old
+
+
+def _compose_cfg(overrides=None):
+    if not OmegaConf.has_resolver('len'):
+        OmegaConf.register_new_resolver('len', lambda x: len(x))
+    GlobalHydra.instance().clear()
+    with initialize_config_dir(version_base=None, config_dir=_CONFIG_DIR):
+        return compose(config_name='config', overrides=list(overrides or []))
+
+
+class SmolVLAIMFAgentTest(unittest.TestCase):
+    def test_compute_loss_and_predict_action_pass_variable_task_to_condition_encoder(self):
+        from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
+
+        condition_encoder = _BF16TaskAwareConditionEncoder()
+        head = _RecordingHead()
+        agent = SmolVLAIMFAttnResAgent(
+            condition_encoder=condition_encoder,
+            action_encoder=nn.Identity(),
+            head=head,
+            action_dim=2,
+            obs_dim=1,
+            pred_horizon=3,
+            obs_horizon=2,
+            diffusion_steps=10,
+            inference_steps=1,
+            num_cams=3,
+            camera_names=_CAMERA_NAMES,
+            num_action_steps=2,
+            head_type='transformer',
+        )
+        images = {
+            'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
+            'top': torch.full((2, 2, 1, 2, 2), 2.0),
+            'front': torch.full((2, 2, 1, 2, 2), 3.0),
+        }
+        qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]])
+        actions = torch.zeros(2, 3, 2)
+        tasks = ['short', 'longer task']
+        loss = agent.compute_loss({'images': images, 'qpos': qpos, 'action': actions, 'task': tasks})
+        self.assertTrue(torch.isfinite(loss))
+        self.assertEqual(condition_encoder.calls[-1]['task'], tasks)
+        self.assertEqual(head.calls[-1]['cond'].shape, (2, 3, 4))
+        self.assertTrue(torch.allclose(head.calls[-1]['cond'][:, 0, 1], torch.tensor([5.0, 11.0])))
+
+        with mock.patch('roboimi.vla.agent_imf.torch.randn', return_value=torch.zeros(2, 3, 2)):
+            pred = agent.predict_action(images, qpos, task=tasks)
+        self.assertEqual(pred.shape, (2, 3, 2))
+        self.assertEqual(condition_encoder.calls[-1]['task'], tasks)
+
+    def test_condition_tokens_are_cast_to_action_head_dtype(self):
+        from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
+
+        condition_encoder = _TaskAwareConditionEncoder()
+        head = _RecordingHead()
+        agent = SmolVLAIMFAttnResAgent(
+            condition_encoder=condition_encoder,
+            action_encoder=nn.Identity(),
+            head=head,
+            action_dim=2,
+            obs_dim=1,
+            pred_horizon=3,
+            obs_horizon=2,
+            diffusion_steps=10,
+            inference_steps=1,
+            num_cams=3,
+            camera_names=_CAMERA_NAMES,
+            num_action_steps=2,
+            head_type='transformer',
+        )
+        images = {
+            'r_vis': torch.full((1, 2, 1, 2, 2), 1.0),
+            'top': torch.full((1, 2, 1, 2, 2), 2.0),
+            'front': torch.full((1, 2, 1, 2, 2), 3.0),
+        }
+        qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32)
+
+        cond = agent._build_cond(images, qpos, task=['pick'])
+
+        self.assertEqual(cond.dtype, head.scale.dtype)
+
+    def test_unknown_dataset_task_uses_configured_task_description(self):
+        from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
+
+        condition_encoder = _TaskAwareConditionEncoder()
+        head = _RecordingHead()
+        agent = SmolVLAIMFAttnResAgent(
+            condition_encoder=condition_encoder,
+            action_encoder=nn.Identity(),
+            head=head,
+            action_dim=2,
+            obs_dim=1,
+            pred_horizon=3,
+            obs_horizon=2,
+            diffusion_steps=10,
+            inference_steps=1,
+            num_cams=3,
+            camera_names=_CAMERA_NAMES,
+            num_action_steps=2,
+            head_type='transformer',
+            task_description='insert the peg into the socket',
+        )
+        images = {
+            'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
+            'top': torch.full((2, 2, 1, 2, 2), 2.0),
+            'front': torch.full((2, 2, 1, 2, 2), 3.0),
+        }
+        qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]], dtype=torch.float32)
+
+        agent._build_cond(images, qpos, task=['unknown', ''])
+
+        self.assertEqual(
+            condition_encoder.calls[-1]['task'],
+            ['insert the peg into the socket', 'insert the peg into the socket'],
+        )
+
+    def test_hydra_config_instantiates_smolvla_imf_attnres_with_condition_encoder_contract(self):
+        cfg = _compose_cfg(overrides=[
+            'agent=smolvla_imf_attnres',
+            'agent.condition_encoder._target_=tests.fake_smolvla_condition_encoder.TaskAwareConditionEncoder',
+            'agent.condition_dim=4',
+            'agent.condition_sequence_length=3',
+            'agent.head.n_layer=1',
+            'agent.head.n_emb=16',
+        ])
+        self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_smolvla_conditioned.SmolVLAIMFAttnResAgent')
+        self.assertEqual(cfg.agent.head.cond_dim, cfg.agent.condition_dim)
+        self.assertEqual(cfg.agent.head.n_obs_steps, cfg.agent.condition_sequence_length)
+
+        with _stub_optional_modules(include_head=True, include_condition_encoder=True):
+            agent = instantiate(cfg.agent)
+
+        self.assertEqual(agent.per_step_cond_dim, 4)
+        self.assertEqual(agent.condition_sequence_length, 3)
+        self.assertIsInstance(agent.noise_pred_net, _StubIMFHead)
+        self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 4)
+        self.assertEqual(agent.noise_pred_net.constructor_kwargs['n_obs_steps'], 3)
+
+    def test_hydra_config_exposes_smolvla_pretrained_vlm_defaults(self):
+        cfg = _compose_cfg(overrides=[
+            'agent=smolvla_imf_attnres',
+        ])
+
+        self.assertEqual(cfg.agent.condition_encoder.model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct')
+        self.assertTrue(cfg.agent.condition_encoder.load_vlm_weights)
+        self.assertEqual(cfg.agent.condition_encoder.num_vlm_layers, 16)
+        self.assertTrue(cfg.agent.condition_encoder.freeze_vlm)
+        self.assertTrue(cfg.agent.condition_encoder.freeze_vision_encoder)
+        self.assertTrue(cfg.agent.condition_encoder.run_text_model)
+        self.assertEqual(cfg.agent.condition_encoder.max_state_dim, 32)
+        self.assertIsNone(cfg.agent.condition_encoder.dataset_image_resize_shape)
+        self.assertIsNone(cfg.agent.condition_encoder.eval_image_resize_shape)
+        self.assertEqual(cfg.agent.condition_dim, 960)
+        self.assertEqual(cfg.agent.condition_sequence_length, 241)
+        self.assertEqual(cfg.agent.head.cond_dim, 960)
+        self.assertEqual(cfg.agent.head.n_obs_steps, 241)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_smolvla_native_agent.py

@@ -0,0 +1,445 @@
+import contextlib
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+from hydra import compose, initialize_config_dir
+from hydra.core.global_hydra import GlobalHydra
+from omegaconf import OmegaConf
+from torch import nn
+
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve())
+_CAMERA_NAMES = ('r_vis', 'top', 'front')
+_MISSING = object()
+
+
+def _stats():
+    return {
+        'qpos_min': [0.0, 10.0],
+        'qpos_max': [10.0, 20.0],
+        'action_min': [-10.0, 10.0],
+        'action_max': [10.0, 30.0],
+    }
+
+
+class FakeTokenizer:
+    def __init__(self):
+        self.calls = []
+
+    def __call__(self, texts, **kwargs):
+        self.calls.append({'texts': list(texts), 'kwargs': dict(kwargs)})
+        batch = len(texts)
+        if kwargs.get('padding') == 'max_length' and kwargs.get('max_length') is not None:
+            max_len = int(kwargs['max_length'])
+        else:
+            max_len = max(len(text) for text in texts) if texts else 0
+        input_ids = torch.zeros(batch, max_len, dtype=torch.long)
+        attention_mask = torch.zeros(batch, max_len, dtype=torch.long)
+        for i, text in enumerate(texts):
+            encoded = [ord(ch) % 127 for ch in text]
+            if kwargs.get('truncation') and len(encoded) > max_len:
+                encoded = encoded[:max_len]
+            if encoded:
+                input_ids[i, : len(encoded)] = torch.tensor(encoded, dtype=torch.long)
+                attention_mask[i, : len(encoded)] = 1
+        return {'input_ids': input_ids, 'attention_mask': attention_mask}
+
+
+class FakeNativeModel(nn.Module):
+    def __init__(self, action_dim=2, chunk_size=3):
+        super().__init__()
+        self.anchor = nn.Parameter(torch.tensor(0.0))
+        self.action_dim = action_dim
+        self.chunk_size = chunk_size
+        self.forward_calls = []
+        self.sample_calls = []
+        self.sample_return = torch.tensor([[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]]])
+
+    def forward(self, **kwargs):
+        self.forward_calls.append({k: _clone(v) for k, v in kwargs.items()})
+        actions = kwargs['actions']
+        loss = (actions**2).sum(dim=-1)
+        if 'action_is_pad' in kwargs and kwargs['action_is_pad'] is not None:
+            mask = (~kwargs['action_is_pad']).to(loss.dtype)
+            loss = (loss * mask).sum() / mask.sum().clamp_min(1.0)
+        else:
+            loss = loss.mean()
+        return {'loss': loss + self.anchor}
+
+    def sample_actions(self, **kwargs):
+        self.sample_calls.append({k: _clone(v) for k, v in kwargs.items()})
+        return self.sample_return.to(device=kwargs['state'].device, dtype=kwargs['state'].dtype)
+
+
+def _clone(value):
+    if torch.is_tensor(value):
+        return value.detach().clone()
+    if isinstance(value, dict):
+        return {k: _clone(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return list(value)
+    return value
+
+
+def _make_agent(**kwargs):
+    from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
+
+    params = dict(
+        model=FakeNativeModel(),
+        tokenizer=FakeTokenizer(),
+        action_dim=2,
+        obs_dim=2,
+        chunk_size=3,
+        n_action_steps=2,
+        obs_horizon=2,
+        camera_names=_CAMERA_NAMES,
+        num_cams=3,
+        dataset_stats=_stats(),
+        normalization_type='min_max',
+        task_description='fallback task',
+        model_config={'resize_imgs_with_padding': None},
+    )
+    params.update(kwargs)
+    return SmolVLANativeAgent(**params)
+
+
+def _batch(task=None):
+    images = {
+        'front': torch.full((2, 2, 1, 2, 2), 3.0),
+        'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
+        'top': torch.full((2, 2, 1, 2, 2), 2.0),
+    }
+    batch = {
+        'images': images,
+        'qpos': torch.tensor([[[0.0, 10.0], [10.0, 20.0]], [[5.0, 15.0], [10.0, 10.0]]]),
+        'action': torch.tensor([[[-10.0, 10.0], [0.0, 20.0], [10.0, 30.0]], [[0.0, 20.0], [10.0, 10.0], [-10.0, 30.0]]]),
+        'action_is_pad': torch.tensor([[False, False, True], [False, True, True]]),
+    }
+    if task is not None:
+        batch['task'] = task
+    return batch
+
+
+@contextlib.contextmanager
+def _stub_native_modules():
+    old_pkg = sys.modules.get('roboimi.vla.models.smolvla', _MISSING)
+    old_conf = sys.modules.get('roboimi.vla.models.smolvla.configuration', _MISSING)
+    try:
+        pkg = types.ModuleType('roboimi.vla.models.smolvla')
+        conf = types.ModuleType('roboimi.vla.models.smolvla.configuration')
+
+        class NativeSmolVLAConfig:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+        conf.NativeSmolVLAConfig = NativeSmolVLAConfig
+        sys.modules['roboimi.vla.models.smolvla'] = pkg
+        sys.modules['roboimi.vla.models.smolvla.configuration'] = conf
+        yield
+    finally:
+        for name, old in [
+            ('roboimi.vla.models.smolvla.configuration', old_conf),
+            ('roboimi.vla.models.smolvla', old_pkg),
+        ]:
+            if old is _MISSING:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = old
+
+
+def _compose_cfg(overrides=None):
+    if not OmegaConf.has_resolver('len'):
+        OmegaConf.register_new_resolver('len', lambda x: len(x))
+    GlobalHydra.instance().clear()
+    with initialize_config_dir(version_base=None, config_dir=_CONFIG_DIR):
+        return compose(config_name='config', overrides=list(overrides or []))
+
+
+
+
+class StrictSignatureNativeModel(nn.Module):
+    def __init__(self, action_dim=2, max_action_dim=4, chunk_size=3):
+        super().__init__()
+        self.anchor = nn.Parameter(torch.tensor(0.0))
+        self.action_dim = action_dim
+        self.max_action_dim = max_action_dim
+        self.chunk_size = chunk_size
+        self.forward_calls = []
+        self.sample_calls = []
+
+    def forward(self, images, img_masks, lang_tokens, lang_masks, state, actions):
+        self.forward_calls.append({
+            'images': _clone(images),
+            'img_masks': _clone(img_masks),
+            'lang_tokens': _clone(lang_tokens),
+            'lang_masks': _clone(lang_masks),
+            'state': _clone(state),
+            'actions': _clone(actions),
+        })
+        # Per-element losses in the native core's max_action_dim space.
+        return actions.pow(2) + self.anchor
+
+    def sample_actions(self, images, img_masks, lang_tokens, lang_masks, state):
+        self.sample_calls.append({
+            'images': _clone(images),
+            'img_masks': _clone(img_masks),
+            'lang_tokens': _clone(lang_tokens),
+            'lang_masks': _clone(lang_masks),
+            'state': _clone(state),
+        })
+        batch_size = state.shape[0]
+        out = torch.zeros(batch_size, self.chunk_size, self.max_action_dim, device=state.device, dtype=state.dtype)
+        out[..., : self.action_dim] = torch.tensor(
+            [[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]]],
+            device=state.device,
+            dtype=state.dtype,
+        ).expand(batch_size, -1, -1)
+        out[..., self.action_dim :] = 123.0
+        return out
+
+
+class SmolVLANativeAgentTest(unittest.TestCase):
+    def test_compute_loss_orders_normalizes_tokenizes_and_passes_action_pad(self):
+        agent = _make_agent()
+        batch = _batch(task=['pick', 'place'])
+
+        loss = agent.compute_loss(batch)
+
+        self.assertTrue(torch.isfinite(loss))
+        call = agent.model.forward_calls[-1]
+        self.assertIsInstance(call['images'], list)
+        self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
+        self.assertTrue(torch.allclose(call['images'][0], torch.full((2, 1, 2, 2), 1.0)))
+        self.assertTrue(torch.allclose(call['state'], torch.tensor([[1.0, 1.0], [1.0, -1.0]])))
+        self.assertTrue(torch.allclose(call['actions'][0], torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])))
+        self.assertNotIn('action_is_pad', call)
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['pick\n', 'place\n'])
+        self.assertIn('lang_tokens', call)
+        self.assertIn('lang_masks', call)
+        self.assertEqual(call['lang_tokens'].dtype, torch.long)
+        self.assertEqual(call['lang_masks'].dtype, torch.bool)
+
+    def test_task_fallback_for_missing_empty_and_unknown_but_valid_task_wins(self):
+        agent = _make_agent(task_description='default instruction')
+
+        agent.compute_loss(_batch())
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'default instruction\n'])
+
+        agent.compute_loss(_batch(task=[None, '']))
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'default instruction\n'])
+
+        agent.compute_loss(_batch(task=['unknown', 'valid task']))
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'valid task\n'])
+
+    def test_missing_camera_raises_value_error(self):
+        agent = _make_agent()
+        batch = _batch(task=['pick', 'place'])
+        del batch['images']['top']
+
+        with self.assertRaisesRegex(ValueError, 'missing.*top'):
+            agent.compute_loss(batch)
+
+    def test_predict_action_chunk_samples_and_denormalizes_actions(self):
+        agent = _make_agent()
+        batch = _batch(task=['pick', 'place'])
+        batch.pop('action')
+        batch.pop('action_is_pad')
+        agent.model.sample_return = torch.tensor([[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]], [[0.5, -0.5], [-0.5, 0.5], [0.0, 1.0]]])
+
+        actions = agent.predict_action_chunk(batch)
+
+        self.assertTrue(torch.allclose(actions[0], torch.tensor([[-10.0, 30.0], [0.0, 20.0], [10.0, 10.0]])))
+        self.assertEqual(actions.shape, (2, 3, 2))
+        self.assertEqual(len(agent.model.sample_calls), 1)
+        self.assertTrue(torch.allclose(agent.model.sample_calls[-1]['state'][0], torch.tensor([1.0, 1.0])))
+
+    def test_select_action_queues_actions_and_defaults_chunk_start_zero(self):
+        agent = _make_agent(n_action_steps=2, action_chunk_start=0)
+        agent.model.sample_return = torch.tensor([[[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]])
+        obs = {
+            'images': {name: torch.full((1, 2, 2), float(i + 1)) for i, name in enumerate(_CAMERA_NAMES)},
+            'qpos': torch.tensor([0.0, 10.0]),
+            'task': 'rollout task',
+        }
+
+        action0 = agent.select_action(obs)
+        action1 = agent.select_action(obs)
+
+        self.assertEqual(len(agent.model.sample_calls), 1)
+        self.assertTrue(torch.allclose(action0, torch.tensor([-10.0, 10.0])))
+        self.assertTrue(torch.allclose(action1, torch.tensor([0.0, 20.0])))
+        self.assertEqual(agent.action_chunk_start, 0)
+
+    def test_get_normalization_stats_returns_stats(self):
+        stats = _make_agent().get_normalization_stats()
+        self.assertEqual(stats['normalization_type'], 'min_max')
+        self.assertEqual(stats['qpos_min'], [0.0, 10.0])
+
+
+    def test_agent_adapts_roboimi_batch_to_native_vla_signature_and_reduces_loss(self):
+        from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
+
+        model = StrictSignatureNativeModel(action_dim=2, max_action_dim=4, chunk_size=3)
+        agent = SmolVLANativeAgent(
+            model=model,
+            tokenizer=FakeTokenizer(),
+            action_dim=2,
+            obs_dim=2,
+            chunk_size=3,
+            n_action_steps=2,
+            obs_horizon=2,
+            camera_names=_CAMERA_NAMES,
+            num_cams=3,
+            dataset_stats=_stats(),
+            normalization_type='min_max',
+            model_config={'max_state_dim': 4, 'max_action_dim': 4, 'resize_imgs_with_padding': None},
+        )
+
+        batch = _batch(task=['pick', 'place'])
+        batch['images'] = {
+            'front': torch.full((2, 2, 1, 2, 2), 0.6),
+            'r_vis': torch.full((2, 2, 1, 2, 2), 0.2),
+            'top': torch.full((2, 2, 1, 2, 2), 0.4),
+        }
+
+        loss = agent.compute_loss(batch)
+
+        self.assertEqual(loss.ndim, 0)
+        self.assertTrue(torch.isfinite(loss))
+        call = model.forward_calls[-1]
+        self.assertIsInstance(call['images'], list)
+        self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
+        self.assertEqual(tuple(call['images'][0].shape), (2, 1, 2, 2))
+        self.assertTrue(torch.allclose(call['images'][0], torch.full((2, 1, 2, 2), -0.6)))
+        self.assertTrue(torch.allclose(call['images'][1], torch.full((2, 1, 2, 2), -0.2)))
+        self.assertTrue(torch.allclose(call['images'][2], torch.full((2, 1, 2, 2), 0.2)))
+        self.assertEqual(len(call['img_masks']), len(_CAMERA_NAMES))
+        self.assertTrue(torch.equal(call['img_masks'][0], torch.ones(2, dtype=torch.bool)))
+        self.assertTrue(torch.equal(call['lang_tokens'], model.forward_calls[-1]['lang_tokens']))
+        self.assertEqual(call['lang_tokens'].dtype, torch.long)
+        self.assertEqual(call['lang_masks'].dtype, torch.bool)
+        self.assertEqual(tuple(call['state'].shape), (2, 4))
+        self.assertTrue(torch.allclose(call['state'][0], torch.tensor([1.0, 1.0, 0.0, 0.0])))
+        self.assertEqual(tuple(call['actions'].shape), (2, 3, 4))
+        self.assertTrue(torch.allclose(call['actions'][0, :, :2], torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])))
+        self.assertTrue(torch.allclose(call['actions'][..., 2:], torch.zeros(2, 3, 2)))
+        # Valid entries: sample0 steps 0,1 and sample1 step 0, only first action_dim losses are reduced.
+        expected = (2.0 + 0.0 + 0.0) / (3 * 2)
+        self.assertTrue(torch.allclose(loss, torch.tensor(expected)))
+
+    def test_predict_action_chunk_accepts_native_max_action_dim_and_crops_before_denorm(self):
+        from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
+
+        model = StrictSignatureNativeModel(action_dim=2, max_action_dim=4, chunk_size=3)
+        agent = SmolVLANativeAgent(
+            model=model,
+            tokenizer=FakeTokenizer(),
+            action_dim=2,
+            obs_dim=2,
+            chunk_size=3,
+            n_action_steps=2,
+            obs_horizon=2,
+            camera_names=_CAMERA_NAMES,
+            num_cams=3,
+            dataset_stats=_stats(),
+            normalization_type='min_max',
+            model_config={'max_state_dim': 4, 'max_action_dim': 4, 'resize_imgs_with_padding': None},
+        )
+        batch = _batch(task=['pick', 'place'])
+        batch.pop('action')
+        batch.pop('action_is_pad')
+
+        actions = agent.predict_action_chunk(batch)
+
+        self.assertEqual(actions.shape, (2, 3, 2))
+        self.assertTrue(torch.allclose(actions[0], torch.tensor([[-10.0, 30.0], [0.0, 20.0], [10.0, 10.0]])))
+        call = model.sample_calls[-1]
+        self.assertEqual(tuple(call['state'].shape), (2, 4))
+        self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
+
+    def test_build_model_filters_and_maps_non_core_config_fields(self):
+        from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
+        from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
+
+        captured = {}
+
+        class BuildOnlyAgent(SmolVLANativeAgent):
+            def _build_tokenizer(self, tokenizer_name):
+                return FakeTokenizer()
+
+            def _build_model(self):
+                cfg_kwargs = self._native_config_kwargs()
+                captured.update(cfg_kwargs)
+                return FakeNativeModel(action_dim=2, chunk_size=3)
+
+        agent = BuildOnlyAgent(
+            action_dim=2,
+            obs_dim=2,
+            chunk_size=3,
+            n_action_steps=2,
+            obs_horizon=2,
+            camera_names=_CAMERA_NAMES,
+            num_cams=3,
+            dataset_stats=_stats(),
+            normalization_type='min_max',
+            model_config={
+                'state_dim': 2,
+                'action_dim': 2,
+                'max_state_dim': 4,
+                'max_action_dim': 4,
+                'tokenizer_name': 'dummy-tokenizer',
+                'freeze_vlm': True,
+                'image_resize_shape': [512, 320],
+                'num_cameras': 3,
+                'load_vlm_weights': False,
+            },
+        )
+        self.assertIsNotNone(agent.model)
+        config = NativeSmolVLAConfig(**captured)
+        self.assertEqual(config.resize_imgs_with_padding, (512, 320))
+        self.assertEqual(config.max_state_dim, 4)
+        self.assertEqual(config.max_action_dim, 4)
+        self.assertNotIn('state_dim', captured)
+        self.assertNotIn('action_dim', captured)
+        self.assertNotIn('tokenizer_name', captured)
+        self.assertNotIn('freeze_vlm', captured)
+        self.assertNotIn('image_resize_shape', captured)
+        self.assertNotIn('num_cameras', captured)
+
+    def test_hydra_config_target_and_key_fields(self):
+        with _stub_native_modules():
+            cfg = _compose_cfg(overrides=['agent=smolvla_native'])
+        self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_smolvla_native.SmolVLANativeAgent')
+        self.assertEqual(cfg.agent.action_dim, 16)
+        self.assertEqual(cfg.agent.obs_dim, 16)
+        self.assertEqual(cfg.agent.normalization_type, 'gaussian')
+        self.assertEqual(list(cfg.agent.camera_names), list(cfg.data.camera_names))
+        self.assertEqual(cfg.agent.num_cams, len(cfg.data.camera_names))
+        self.assertEqual(cfg.agent.chunk_size, 32)
+        self.assertEqual(cfg.agent.n_action_steps, 16)
+        self.assertEqual(cfg.agent.model_config.max_state_dim, 32)
+        self.assertEqual(cfg.agent.model_config.max_action_dim, 32)
+        self.assertEqual(cfg.agent.model_config.chunk_size, 32)
+        self.assertEqual(cfg.agent.model_config.n_action_steps, 16)
+        self.assertIsNone(cfg.agent.dataset_image_resize_shape)
+        self.assertIsNone(cfg.agent.eval_image_resize_shape)
+        self.assertEqual(list(cfg.agent.model_config.resize_imgs_with_padding), [512, 512])
+        self.assertNotIn('state_dim', cfg.agent.model_config)
+        self.assertNotIn('action_dim', cfg.agent.model_config)
+        self.assertNotIn('tokenizer_name', cfg.agent.model_config)
+        self.assertEqual(cfg.agent.model_config.vlm_model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct')
+
+    def test_tokenize_tasks_preserves_existing_single_newline(self):
+        agent = _make_agent(model_config={'resize_imgs_with_padding': None, 'pad_language_to': 'max_length', 'tokenizer_max_length': 12})
+
+        agent._tokenize_tasks(['already newline\n', 'needs newline'], batch_size=2, device=torch.device('cpu'))
+
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['already newline\n', 'needs newline\n'])
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_smolvla_native_modeling.py

@@ -0,0 +1,121 @@
+import unittest
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from roboimi.vla.models.smolvla import NativeSmolVLAConfig
+from roboimi.vla.models.smolvla.modeling import (
+    VLAFlowMatching,
+    make_att_2d_masks,
+    pad_vector,
+    resize_with_pad,
+)
+
+
+class FakeTokenizer:
+    fake_image_token_id = 32000
+    global_image_token_id = 32001
+
+
+class FakeProcessor:
+    tokenizer = FakeTokenizer()
+
+
+class FakeVLMWithExpert(nn.Module):
+    def __init__(self, vlm_hidden_size=8, expert_hidden_size=6, image_tokens=3, vocab_size=64):
+        super().__init__()
+        self.config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=vlm_hidden_size))
+        self.expert_hidden_size = expert_hidden_size
+        self.image_tokens = image_tokens
+        self.processor = FakeProcessor()
+        self.vlm = SimpleNamespace(device=torch.device("cpu"))
+        self.image_proj = nn.Linear(3, vlm_hidden_size)
+        self.token_emb = nn.Embedding(vocab_size, vlm_hidden_size)
+        self.suffix_proj = nn.Linear(expert_hidden_size, expert_hidden_size)
+
+    def embed_image(self, image):
+        # Deterministic lightweight image embedding: pool pixels, then repeat.
+        pooled = image.mean(dim=(-1, -2)).to(dtype=torch.float32)
+        return self.image_proj(pooled).unsqueeze(1).expand(-1, self.image_tokens, -1)
+
+    def embed_language_tokens(self, tokens):
+        return self.token_emb(tokens)
+
+    def forward(self, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, fill_kv_cache):
+        prefix_embs, suffix_embs = inputs_embeds
+        prefix_out = prefix_embs if prefix_embs is not None else None
+        suffix_out = self.suffix_proj(suffix_embs) if suffix_embs is not None else None
+        cache = ("fake-cache",) if fill_kv_cache else past_key_values
+        return (prefix_out, suffix_out), cache
+
+
+class NativeSmolVLAModelingTest(unittest.TestCase):
+    def test_config_rejects_action_steps_greater_than_chunk_size(self):
+        with self.assertRaisesRegex(ValueError, "n_action_steps"):
+            NativeSmolVLAConfig(chunk_size=2, n_action_steps=3)
+
+    def test_pad_vector_pads_returns_original_for_equal_and_rejects_truncation(self):
+        vector = torch.tensor([[1.0, 2.0, 3.0]])
+        padded = pad_vector(vector, 5)
+        self.assertEqual(tuple(padded.shape), (1, 5))
+        torch.testing.assert_close(padded, torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0]]))
+
+        same = pad_vector(vector, 3)
+        self.assertIs(same, vector)
+
+        with self.assertRaisesRegex(ValueError, "target dimension"):
+            pad_vector(vector, 2)
+
+    def test_resize_with_pad_returns_requested_spatial_size(self):
+        img = torch.arange(2 * 3 * 4 * 8, dtype=torch.float32).reshape(2, 3, 4, 8)
+        resized = resize_with_pad(img, width=10, height=10, pad_value=-1)
+        self.assertEqual(tuple(resized.shape), (2, 3, 10, 10))
+
+    def test_make_att_2d_masks_implements_prefix_lm_semantics(self):
+        pad_masks = torch.tensor([[True, True, True, True, False]])
+        # First two tokens are bidirectional prefix, later valid tokens are causal.
+        att_masks = torch.tensor([[False, False, True, True, True]])
+        mask = make_att_2d_masks(pad_masks, att_masks)
+        expected = torch.tensor(
+            [[
+                [True, True, False, False, False],
+                [True, True, False, False, False],
+                [True, True, True, False, False],
+                [True, True, True, True, False],
+                [False, False, False, False, False],
+            ]]
+        )
+        torch.testing.assert_close(mask, expected)
+
+    def test_vla_flow_matching_forward_and_sample_actions_with_fake_vlm(self):
+        torch.manual_seed(0)
+        config = NativeSmolVLAConfig(
+            chunk_size=4,
+            n_action_steps=4,
+            max_state_dim=5,
+            max_action_dim=3,
+            num_steps=2,
+            prefix_length=8,
+            add_image_special_tokens=False,
+        )
+        fake_vlm = FakeVLMWithExpert(vlm_hidden_size=8, expert_hidden_size=6)
+        model = VLAFlowMatching(config, vlm_with_expert=fake_vlm)
+
+        bsize = 2
+        images = [torch.randn(bsize, 3, 6, 6)]
+        img_masks = [torch.tensor([True, False])]
+        lang_tokens = torch.tensor([[1, 2, 3], [4, 5, 0]])
+        lang_masks = torch.tensor([[True, True, True], [True, True, False]])
+        state = torch.randn(bsize, config.max_state_dim)
+        actions = torch.randn(bsize, config.chunk_size, config.max_action_dim)
+
+        losses = model(images, img_masks, lang_tokens, lang_masks, state, actions)
+        self.assertEqual(tuple(losses.shape), (bsize, config.chunk_size, config.max_action_dim))
+
+        sampled = model.sample_actions(images, img_masks, lang_tokens, lang_masks, state)
+        self.assertEqual(tuple(sampled.shape), (bsize, config.chunk_size, config.max_action_dim))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_smolvla_prefix_encoder.py

@@ -0,0 +1,328 @@
+import types
+import unittest
+from unittest import mock
+
+import torch
+from torch import nn
+
+
+class _FakeVisionOutput:
+    def __init__(self, last_hidden_state):
+        self.last_hidden_state = last_hidden_state
+
+
+class _FakeVisionModel(nn.Module):
+    def __init__(self, hidden_size=4):
+        super().__init__()
+        self.dtype = torch.float32
+        self.scale = nn.Parameter(torch.tensor(1.0))
+        self.calls = []
+        self.hidden_size = hidden_size
+
+    def forward(self, pixel_values=None, patch_attention_mask=None):
+        self.calls.append({
+            'pixel_values': pixel_values.detach().clone(),
+            'patch_attention_mask': patch_attention_mask,
+        })
+        pooled = pixel_values.mean(dim=(2, 3)) * self.scale
+        tokens = torch.stack([pooled, pooled + 1.0], dim=1)
+        return _FakeVisionOutput(tokens)
+
+
+class _FakeConnector(nn.Module):
+    def __init__(self, in_dim=3, out_dim=4):
+        super().__init__()
+        self.proj = nn.Linear(in_dim, out_dim, bias=False)
+        with torch.no_grad():
+            self.proj.weight.copy_(
+                torch.tensor(
+                    [
+                        [1.0, 0.0, 0.0],
+                        [0.0, 1.0, 0.0],
+                        [0.0, 0.0, 1.0],
+                        [1.0, 1.0, 1.0],
+                    ]
+                )
+            )
+
+    def forward(self, x):
+        return self.proj(x)
+
+
+class _FakeTextModel(nn.Module):
+    def __init__(self, vocab_size=32, hidden_size=4, num_layers=6):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, hidden_size)
+        self.layers = nn.ModuleList([nn.Linear(hidden_size, hidden_size) for _ in range(num_layers)])
+        self.norm = nn.Identity()
+        self.forward_calls = []
+        with torch.no_grad():
+            for idx in range(vocab_size):
+                self.embed.weight[idx].fill_(float(idx))
+
+    def get_input_embeddings(self):
+        return self.embed
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        use_cache=None,
+        return_dict=True,
+        cache_position=None,
+        **kwargs,
+    ):
+        del input_ids, past_key_values, use_cache, cache_position, kwargs
+        self.forward_calls.append({
+            'attention_mask': None if attention_mask is None else attention_mask.detach().clone(),
+            'position_ids': None if position_ids is None else position_ids.detach().clone(),
+            'inputs_embeds': inputs_embeds.detach().clone(),
+        })
+        hidden = inputs_embeds
+        for layer in self.layers:
+            hidden = layer(hidden)
+        hidden = self.norm(hidden)
+        if return_dict:
+            return types.SimpleNamespace(last_hidden_state=hidden)
+        return (hidden,)
+
+
+class _FakeVLM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        text_config = types.SimpleNamespace(hidden_size=4, head_dim=2, num_attention_heads=2, num_key_value_heads=1)
+        self.config = types.SimpleNamespace(text_config=text_config)
+        self.model = types.SimpleNamespace(
+            vision_model=_FakeVisionModel(hidden_size=4),
+            connector=_FakeConnector(in_dim=3, out_dim=4),
+            text_model=_FakeTextModel(hidden_size=4, num_layers=6),
+        )
+
+
+class _FakeTokenizer:
+    fake_image_token_id = 29
+    global_image_token_id = 30
+
+    def __init__(self):
+        self.padding_side = 'left'
+        self.calls = []
+
+    def __call__(self, text, *, padding, max_length, return_tensors, truncation):
+        self.calls.append({
+            'text': list(text),
+            'padding': padding,
+            'max_length': max_length,
+            'return_tensors': return_tensors,
+            'truncation': truncation,
+            'padding_side': self.padding_side,
+        })
+        batch = len(text)
+        ids = torch.zeros(batch, max_length, dtype=torch.long)
+        mask = torch.zeros(batch, max_length, dtype=torch.bool)
+        for row, item in enumerate(text):
+            del item
+            ids[row, :3] = torch.tensor([1, 2, 3])
+            mask[row, :3] = True
+        return {'input_ids': ids, 'attention_mask': mask}
+
+
+class SmolVLAPrefixEncoderTest(unittest.TestCase):
+    def test_loads_pretrained_vlm_with_local_files_only_crops_layers_and_freezes_vlm(self):
+        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
+
+        fake_vlm = _FakeVLM()
+        fake_tokenizer = _FakeTokenizer()
+        with mock.patch(
+            'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoModelForImageTextToText.from_pretrained',
+            return_value=fake_vlm,
+        ) as model_loader, mock.patch(
+            'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoTokenizer.from_pretrained',
+            return_value=fake_tokenizer,
+        ) as tokenizer_loader:
+            encoder = SmolVLAPrefixEncoder(
+                model_name='HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
+                local_files_only=True,
+                num_vlm_layers=2,
+                freeze_vlm=True,
+                max_state_dim=5,
+                tokenizer_max_length=4,
+                camera_names=('r_vis', 'top'),
+                resize_imgs_with_padding=None,
+            )
+
+        model_loader.assert_called_once()
+        self.assertEqual(model_loader.call_args.kwargs['local_files_only'], True)
+        self.assertEqual(model_loader.call_args.kwargs['torch_dtype'], 'bfloat16')
+        tokenizer_loader.assert_called_once_with(
+            'HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
+            local_files_only=True,
+        )
+        self.assertEqual(len(encoder.vlm.model.text_model.layers), 2)
+        self.assertTrue(all(not p.requires_grad for p in encoder.vlm.parameters()))
+        self.assertFalse(encoder.vlm.training)
+        encoder.train()
+        self.assertFalse(encoder.vlm.training)
+
+    def test_accepts_train_eval_resize_compatibility_fields(self):
+        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
+
+        encoder = SmolVLAPrefixEncoder(
+            vlm=_FakeVLM(),
+            tokenizer=_FakeTokenizer(),
+            num_vlm_layers=3,
+            camera_names=('r_vis', 'top'),
+            resize_imgs_with_padding=None,
+            dataset_image_resize_shape=None,
+            eval_image_resize_shape=(640, 480),
+        )
+
+        self.assertIsNone(encoder.dataset_image_resize_shape)
+        self.assertEqual(encoder.eval_image_resize_shape, (640, 480))
+
+    def test_embed_prefix_uses_variable_tasks_camera_order_last_state_and_smolvla_masks(self):
+        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
+
+        fake_vlm = _FakeVLM()
+        fake_tokenizer = _FakeTokenizer()
+        encoder = SmolVLAPrefixEncoder(
+            vlm=fake_vlm,
+            tokenizer=fake_tokenizer,
+            num_vlm_layers=3,
+            freeze_vlm=True,
+            train_state_proj=True,
+            max_state_dim=5,
+            tokenizer_max_length=4,
+            camera_names=('r_vis', 'top'),
+            resize_imgs_with_padding=None,
+        )
+        with torch.no_grad():
+            encoder.state_proj.weight.zero_()
+            encoder.state_proj.bias.zero_()
+            encoder.state_proj.weight[:, :4] = torch.eye(4)
+
+        images = {
+            'top': torch.full((2, 2, 3, 2, 2), 0.75),
+            'r_vis': torch.full((2, 2, 3, 2, 2), 0.25),
+        }
+        state = torch.tensor(
+            [
+                [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+                [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]],
+            ]
+        )
+        tasks = ['pick red cube', 'open drawer']
+        out = encoder.embed_prefix(images=images, state=state, task=tasks)
+
+        # 2 cameras * 2 image tokens + 4 language tokens + 1 state token
+        self.assertEqual(out.shape, (2, 9, 4))
+        self.assertEqual(encoder.output_dim, 4)
+        self.assertEqual(encoder.tokens_per_step, 9)
+        self.assertEqual(encoder.condition_sequence_length, 9)
+
+        self.assertEqual(fake_tokenizer.calls[-1]['text'], ['pick red cube\n', 'open drawer\n'])
+        self.assertEqual(fake_tokenizer.calls[-1]['padding'], 'max_length')
+        self.assertEqual(fake_tokenizer.calls[-1]['max_length'], 4)
+        self.assertEqual(fake_tokenizer.calls[-1]['padding_side'], 'right')
+
+        # Camera order is r_vis then top, and pixels are mapped [0,1] -> [-1,1].
+        first_camera_pixels = fake_vlm.model.vision_model.calls[0]['pixel_values']
+        second_camera_pixels = fake_vlm.model.vision_model.calls[1]['pixel_values']
+        self.assertTrue(torch.allclose(first_camera_pixels, torch.full((2, 3, 2, 2), -0.5)))
+        self.assertTrue(torch.allclose(second_camera_pixels, torch.full((2, 3, 2, 2), 0.5)))
+
+        # Last token is padded last state projected from [4,5,6,0,0] and [10,11,12,0,0].
+        self.assertTrue(torch.allclose(out[0, -1], torch.tensor([4.0, 5.0, 6.0, 0.0])))
+        self.assertTrue(torch.allclose(out[1, -1], torch.tensor([10.0, 11.0, 12.0, 0.0])))
+
+        pad_mask, att_mask = encoder.last_prefix_pad_mask, encoder.last_prefix_att_mask
+        self.assertEqual(pad_mask.shape, (2, 9))
+        self.assertEqual(att_mask.shape, (2, 9))
+        self.assertTrue(torch.all(att_mask[:, :8] == 0))
+        self.assertTrue(torch.all(att_mask[:, -1] == 1))
+
+    def test_forward_can_run_cropped_frozen_text_model_over_prefix_tokens(self):
+        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
+
+        fake_vlm = _FakeVLM()
+        fake_tokenizer = _FakeTokenizer()
+        encoder = SmolVLAPrefixEncoder(
+            vlm=fake_vlm,
+            tokenizer=fake_tokenizer,
+            num_vlm_layers=2,
+            freeze_vlm=True,
+            max_state_dim=5,
+            tokenizer_max_length=4,
+            camera_names=('r_vis',),
+            resize_imgs_with_padding=None,
+            run_text_model=True,
+        )
+        images = {
+            'r_vis': torch.full((2, 1, 3, 2, 2), 0.25),
+        }
+        state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]])
+
+        out = encoder(images, state=state, task=['pick', 'place'])
+
+        # 1 camera * 2 image tokens + 4 language tokens + 1 state token.
+        self.assertEqual(out.shape, (2, 7, 4))
+        self.assertEqual(len(fake_vlm.model.text_model.layers), 2)
+        self.assertEqual(len(fake_vlm.model.text_model.forward_calls), 1)
+        text_call = fake_vlm.model.text_model.forward_calls[-1]
+        self.assertEqual(tuple(text_call['attention_mask'].shape), (2, 1, 7, 7))
+        self.assertEqual(tuple(text_call['position_ids'].shape), (2, 7))
+        self.assertTrue(torch.all(out != 0))
+
+    def test_frozen_vlm_still_backpropagates_through_text_model_to_state_projection(self):
+        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
+
+        fake_vlm = _FakeVLM()
+        encoder = SmolVLAPrefixEncoder(
+            vlm=fake_vlm,
+            tokenizer=_FakeTokenizer(),
+            num_vlm_layers=3,
+            freeze_vlm=True,
+            train_state_proj=True,
+            max_state_dim=5,
+            tokenizer_max_length=4,
+            camera_names=('r_vis',),
+            resize_imgs_with_padding=None,
+            run_text_model=True,
+        )
+        images = {
+            'r_vis': torch.full((2, 1, 3, 2, 2), 0.25),
+        }
+        state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]])
+
+        out = encoder(images, state=state, task=['pick', 'place'])
+        out[:, -1].sum().backward()
+
+        self.assertIsNotNone(encoder.state_proj.weight.grad)
+        self.assertGreater(float(encoder.state_proj.weight.grad.abs().sum()), 0.0)
+        self.assertTrue(all(param.grad is None for param in fake_vlm.parameters()))
+
+    def test_embed_prefix_rejects_missing_camera_and_task_batch_mismatch(self):
+        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
+
+        encoder = SmolVLAPrefixEncoder(
+            vlm=_FakeVLM(),
+            tokenizer=_FakeTokenizer(),
+            num_vlm_layers=3,
+            camera_names=('r_vis', 'top'),
+            resize_imgs_with_padding=None,
+            max_state_dim=5,
+            tokenizer_max_length=4,
+        )
+        images = {'r_vis': torch.rand(2, 1, 3, 2, 2)}
+        state = torch.rand(2, 1, 3)
+        with self.assertRaisesRegex(ValueError, 'missing.*top'):
+            encoder.embed_prefix(images=images, state=state, task=['a', 'b'])
+        images['top'] = torch.rand(2, 1, 3, 2, 2)
+        with self.assertRaisesRegex(ValueError, 'task batch'):
+            encoder.embed_prefix(images=images, state=state, task=['only one'])
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_train_vla_rollout_validation.py

@@ -14,6 +14,8 @@ from roboimi.demos.vla_scripts import eval_vla, train_vla
 
 
 class _FakeDataset:
+    available_episode_indices = [0, 1]
+
     def __len__(self):
         return 4
 
@@ -29,6 +31,13 @@ class _FakeLoader:
         return iter(self._batches)
 
 
+class _FakeValDataset(_FakeDataset):
+    available_episode_indices = [1]
+
+    def __len__(self):
+        return 2
+
+
 class _FakeOptimizer:
     def __init__(self, lr=1e-3):
         self.param_groups = [{'lr': lr}]
@@ -91,6 +100,16 @@ class _FakeAgent(nn.Module):
         return {}
 
 
+class _CapturingAgent(_FakeAgent):
+    def __init__(self):
+        super().__init__()
+        self.compute_loss_inputs = []
+
+    def compute_loss(self, agent_input):
+        self.compute_loss_inputs.append(agent_input)
+        return (self.weight - torch.tensor(0.5)).pow(2)
+
+
 class _SequentialLossAgent(nn.Module):
     def __init__(self, losses):
         super().__init__()
@@ -150,6 +169,94 @@ class _FakeEvalEnv:
 
 
 class TrainVLARolloutValidationTest(unittest.TestCase):
+    def test_run_training_passes_variable_batch_task_to_agent_input(self):
+        cfg = OmegaConf.create(
+            {
+                'train': {
+                    'device': 'cpu',
+                    'batch_size': 2,
+                    'num_workers': 0,
+                    'val_split': 0.0,
+                    'seed': 0,
+                    'lr': 1e-3,
+                    'max_steps': 1,
+                    'log_freq': 100,
+                    'save_freq': 1000,
+                    'warmup_steps': 1,
+                    'scheduler_type': 'constant',
+                    'min_lr': 0.0,
+                    'grad_clip': 1.0,
+                    'weight_decay': 0.0,
+                    'pretrained_ckpt': None,
+                    'resume_ckpt': None,
+                    'use_swanlab': False,
+                    'rollout_val_freq_epochs': 0,
+                    'rollout_validate_on_checkpoint': False,
+                    'rollout_num_episodes': 1,
+                },
+                'data': {
+                    'camera_names': ['front'],
+                    'dataset_dir': 'unused',
+                },
+                'agent': {
+                    '_target_': 'fake.agent',
+                    'normalization_type': 'min_max',
+                },
+                'eval': {
+                    'ckpt_path': 'unused.pt',
+                    'num_episodes': 1,
+                    'max_timesteps': 1,
+                    'device': 'cpu',
+                    'task_name': 'sim_transfer',
+                    'camera_names': ['front'],
+                    'use_smoothing': False,
+                    'smooth_alpha': 0.3,
+                    'verbose_action': False,
+                    'headless': True,
+                },
+                'experiment': {},
+            }
+        )
+        agent = _CapturingAgent()
+        batch_task = ['pick the red cube', 'insert the peg into the socket']
+
+        def fake_instantiate(config_node, **_kwargs):
+            if config_node is cfg.data:
+                return _FakeDataset()
+            if config_node is cfg.agent:
+                return agent
+            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
+
+        def fake_dataloader(_dataset, *, shuffle, **_kwargs):
+            del shuffle, _kwargs
+            return _FakeLoader(
+                {
+                    'observation.front': torch.zeros(2, 2, 3, 4, 4),
+                    'observation.state': torch.zeros(2, 2, 4),
+                    'action': torch.zeros(2, 8, 2),
+                    'action_is_pad': torch.zeros(2, 8, dtype=torch.bool),
+                    'task': list(batch_task),
+                },
+                length=1,
+            )
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \
+                     mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \
+                     mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \
+                     mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \
+                     mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \
+                     mock.patch.object(train_vla.torch, 'save', return_value=None):
+                    train_vla._run_training(cfg)
+            finally:
+                os.chdir(previous_cwd)
+
+        self.assertEqual(len(agent.compute_loss_inputs), 1)
+        self.assertEqual(agent.compute_loss_inputs[0]['task'], batch_task)
+
     def test_default_train_config_uses_full_dataset_and_epoch_rollout_validation(self):
         cfg = OmegaConf.load(Path('roboimi/vla/conf/config.yaml'))
 
@@ -158,6 +265,151 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
         self.assertGreater(float(cfg.train.lr), 5e-5)
         self.assertGreater(cfg.train.num_workers, 8)
         self.assertEqual(cfg.train.rollout_val_freq_epochs, 50)
+        self.assertEqual(cfg.train.rollout_device, cfg.train.device)
+        self.assertIsNone(cfg.train.rollout_num_workers)
+        self.assertIsNone(cfg.train.rollout_cuda_devices)
+
+    def test_explicit_val_episode_indices_builds_held_out_dataset(self):
+        cfg = OmegaConf.create(
+            {
+                'train': {
+                    'val_episode_indices': [1],
+                    'val_split': 0.0,
+                    'seed': 42,
+                },
+                'data': {},
+            }
+        )
+        instantiate_calls = []
+
+        def fake_instantiate(config_node, **kwargs):
+            del config_node
+            instantiate_calls.append(dict(kwargs))
+            if kwargs.get('episode_indices') == [1]:
+                return _FakeValDataset()
+            return _FakeDataset()
+
+        with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate):
+            dataset, train_dataset, val_dataset, explicit = train_vla.build_train_val_datasets(
+                cfg,
+                dataset_image_resize_shape=None,
+            )
+
+        self.assertIsInstance(dataset, _FakeDataset)
+        self.assertIsInstance(train_dataset, _FakeDataset)
+        self.assertIsInstance(val_dataset, _FakeValDataset)
+        self.assertEqual(explicit, [1])
+        self.assertEqual(instantiate_calls[1]['episode_indices'], [0])
+        self.assertEqual(instantiate_calls[2]['episode_indices'], [1])
+
+
+    def test_run_training_rollout_validation_propagates_gpu_parallel_settings(self):
+        cfg = OmegaConf.create(
+            {
+                'train': {
+                    'device': 'cpu',
+                    'batch_size': 1,
+                    'num_workers': 0,
+                    'val_split': 0.0,
+                    'seed': 0,
+                    'lr': 1e-3,
+                    'max_steps': 2,
+                    'log_freq': 1,
+                    'save_freq': 1000,
+                    'warmup_steps': 1,
+                    'scheduler_type': 'constant',
+                    'min_lr': 0.0,
+                    'grad_clip': 1.0,
+                    'weight_decay': 0.0,
+                    'pretrained_ckpt': None,
+                    'resume_ckpt': None,
+                    'use_swanlab': False,
+                    'rollout_val_freq_epochs': 2,
+                    'rollout_num_episodes': 5,
+                    'rollout_device': 'cuda',
+                    'rollout_num_workers': 4,
+                    'rollout_cuda_devices': [0, 1],
+                    'rollout_response_timeout_s': 123.0,
+                    'rollout_server_startup_timeout_s': 456.0,
+                },
+                'data': {
+                    'camera_names': ['front'],
+                },
+                'agent': {
+                    '_target_': 'fake.agent',
+                },
+                'eval': {
+                    'ckpt_path': 'unused.pt',
+                    'num_episodes': 99,
+                    'max_timesteps': 1,
+                    'device': 'cpu',
+                    'task_name': 'sim_transfer',
+                    'camera_names': ['front'],
+                    'use_smoothing': False,
+                    'smooth_alpha': 0.3,
+                    'verbose_action': False,
+                    'headless': False,
+                },
+            }
+        )
+        rollout_mock = mock.Mock(return_value={'avg_reward': 1.0})
+
+        def fake_instantiate(config_node, **_kwargs):
+            if config_node is cfg.data:
+                return _FakeDataset()
+            if config_node is cfg.agent:
+                return _FakeAgent()
+            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
+
+        def fake_dataloader(_dataset, *, shuffle, **_kwargs):
+            del shuffle, _kwargs
+            return _FakeLoader(
+                {
+                    'observation.front': torch.zeros(1, 3, 2, 2),
+                    'observation.state': torch.zeros(1, 4),
+                    'action': torch.zeros(1, 2),
+                    'action_is_pad': torch.zeros(1, 1, dtype=torch.bool),
+                },
+                length=1,
+            )
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate),                      mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader),                      mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)),                      mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()),                      mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)),                      mock.patch.object(train_vla.torch, 'save', return_value=None),                      mock.patch.object(eval_vla, '_run_eval', rollout_mock, create=True):
+                    train_vla._run_training(cfg)
+            finally:
+                os.chdir(previous_cwd)
+
+        rollout_cfg = rollout_mock.call_args.args[0]
+        self.assertEqual(rollout_cfg.eval.device, 'cuda')
+        self.assertEqual(rollout_cfg.eval.num_workers, 4)
+        self.assertEqual(list(rollout_cfg.eval.cuda_devices), [0, 1])
+        self.assertEqual(float(rollout_cfg.eval.response_timeout_s), 123.0)
+        self.assertEqual(float(rollout_cfg.eval.server_startup_timeout_s), 456.0)
+        self.assertTrue(rollout_cfg.eval.headless)
+        self.assertEqual(rollout_cfg.eval.num_episodes, 5)
+        self.assertFalse(rollout_cfg.eval.record_video)
+        self.assertTrue(rollout_cfg.eval.save_summary_json)
+        self.assertTrue(rollout_cfg.eval.save_trajectory_image)
+
+    def test_resolve_dataset_image_resize_shape_prefers_agent_top_level_override(self):
+        cfg = OmegaConf.create(
+            {
+                'agent': {
+                    'dataset_image_resize_shape': None,
+                    'vision_backbone': {
+                        'dataset_image_resize_shape': [256, 256],
+                    },
+                },
+                'data': {
+                    'image_resize_shape': [224, 224],
+                },
+            }
+        )
+
+        self.assertIsNone(train_vla._resolve_dataset_image_resize_shape(cfg))
 
     def test_training_passes_backbone_image_resize_override_to_dataset_instantiation(self):
         cfg = OmegaConf.create(
@@ -245,6 +497,93 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
         self.assertIn('image_resize_shape', captured_dataset_kwargs)
         self.assertIsNone(captured_dataset_kwargs['image_resize_shape'])
 
+    def test_training_passes_condition_encoder_image_resize_override_to_dataset_instantiation(self):
+        cfg = OmegaConf.create(
+            {
+                'agent': {
+                    'condition_encoder': {
+                        'dataset_image_resize_shape': None,
+                    },
+                    'normalization_type': 'min_max',
+                },
+                'data': {
+                    'dataset_dir': 'unused',
+                    'camera_names': ['front'],
+                    'image_resize_shape': [224, 224],
+                },
+                'train': {
+                    'batch_size': 2,
+                    'lr': 1e-4,
+                    'max_steps': 0,
+                    'device': 'cpu',
+                    'disable_cudnn': False,
+                    'num_workers': 0,
+                    'val_split': 0.0,
+                    'seed': 42,
+                    'log_freq': 1,
+                    'save_freq': 10,
+                    'use_swanlab': False,
+                    'rollout_val_freq_epochs': 0,
+                    'rollout_validate_on_checkpoint': False,
+                    'rollout_num_episodes': 1,
+                    'warmup_steps': 1,
+                    'scheduler_type': 'constant',
+                    'min_lr': 1e-6,
+                    'weight_decay': 1e-5,
+                    'grad_clip': 1.0,
+                    'pretrained_ckpt': None,
+                },
+                'eval': {
+                    'ckpt_path': 'unused.pt',
+                    'num_episodes': 1,
+                    'headless': True,
+                    'device': 'cpu',
+                    'verbose_action': False,
+                },
+                'experiment': {},
+            }
+        )
+        captured_dataset_kwargs = {}
+
+        def fake_instantiate(config_node, **kwargs):
+            if config_node is cfg.data:
+                captured_dataset_kwargs.update(kwargs)
+                return _FakeDataset()
+            if config_node is cfg.agent:
+                return _FakeAgent()
+            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
+
+        def fake_dataloader(_dataset, *, shuffle, **_kwargs):
+            del shuffle, _kwargs
+            return _FakeLoader(
+                {
+                    'observation.front': torch.zeros(1, 3, 2, 2),
+                    'observation.state': torch.zeros(1, 4),
+                    'action': torch.zeros(1, 2),
+                    'action_is_pad': torch.zeros(1, 1, dtype=torch.bool),
+                },
+                length=1,
+            )
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \
+                     mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \
+                     mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \
+                     mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \
+                     mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \
+                     mock.patch.object(train_vla, '_init_swanlab', return_value=None), \
+                     mock.patch.object(train_vla, '_finish_swanlab', return_value=None), \
+                     mock.patch.object(train_vla.torch, 'save', return_value=None):
+                    train_vla._run_training(cfg)
+            finally:
+                os.chdir(previous_cwd)
+
+        self.assertIn('image_resize_shape', captured_dataset_kwargs)
+        self.assertIsNone(captured_dataset_kwargs['image_resize_shape'])
+
     def test_eval_main_delegates_to_plain_run_eval_helper(self):
         cfg = OmegaConf.create(
             {

tests/test_train_vla_transformer_optimizer.py

@@ -39,6 +39,17 @@ class FakeLoader:
         return iter(())
 
 
+class FakeTqdm:
+    def __init__(self, iterable, **_kwargs):
+        self.iterable = iterable
+
+    def __iter__(self):
+        return iter(self.iterable)
+
+    def set_postfix(self, *_args, **_kwargs):
+        return None
+
+
 class FakeScheduler:
     def state_dict(self):
         return {}
@@ -46,13 +57,18 @@ class FakeScheduler:
     def load_state_dict(self, state_dict):
         return None
 
+    def step(self):
+        return None
+
 
 class RecordingAdamW:
     created = []
 
-    def __init__(self, params, lr, weight_decay):
+    def __init__(self, params, lr, weight_decay, betas=(0.9, 0.999), eps=1e-8):
         self.lr = lr
         self.weight_decay = weight_decay
+        self.betas = betas
+        self.eps = eps
         self.param_groups = self._normalize_param_groups(params, lr, weight_decay)
         RecordingAdamW.created.append(self)
 
@@ -79,6 +95,12 @@ class RecordingAdamW:
     def load_state_dict(self, state_dict):
         return None
 
+    def zero_grad(self):
+        return None
+
+    def step(self):
+        return None
+
 
 class RecordingTransformerHead(nn.Module):
     def __init__(self):
@@ -324,7 +346,7 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase):
                      mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
                      mock.patch.object(module, 'AdamW', RecordingAdamW), \
                      mock.patch.object(module.torch, 'save', return_value=None), \
-                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: iterable):
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
                     module.main(cfg)
             finally:
                 os.chdir(previous_cwd)
@@ -404,7 +426,7 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase):
                      mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
                      mock.patch.object(module, 'AdamW', RecordingAdamW), \
                      mock.patch.object(module.torch, 'save', return_value=None), \
-                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: iterable):
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
                     module.main(cfg)
             finally:
                 os.chdir(previous_cwd)
@@ -422,3 +444,159 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase):
 
 if __name__ == '__main__':
     unittest.main()
+
+class TrainVLASmolVLAOptimizerTest(unittest.TestCase):
+    def test_build_training_optimizer_excludes_frozen_vlm_parameters_and_keeps_state_proj_and_head(self):
+        module = TrainVLATransformerOptimizerTest()._load_train_vla_module()
+
+        class _Head(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.proj = nn.Linear(2, 2)
+
+            def get_optim_groups(self, weight_decay):
+                return [{'params': list(self.parameters()), 'weight_decay': weight_decay}]
+
+        class _ConditionEncoder(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.vlm = nn.Linear(2, 2)
+                for param in self.vlm.parameters():
+                    param.requires_grad = False
+                self.state_proj = nn.Linear(2, 2)
+
+        class _Agent(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.noise_pred_net = _Head()
+                self.condition_encoder = _ConditionEncoder()
+
+        agent = _Agent()
+        with mock.patch.object(module, 'AdamW', RecordingAdamW):
+            optimizer = module.build_training_optimizer(agent, lr=1e-4, weight_decay=0.01)
+
+        names_by_param_id = {id(param): name for name, param in agent.named_parameters()}
+        optimizer_names = {
+            names_by_param_id[id(param)]
+            for group in optimizer.param_groups
+            for param in group['params']
+        }
+        self.assertIn('condition_encoder.state_proj.weight', optimizer_names)
+        self.assertIn('condition_encoder.state_proj.bias', optimizer_names)
+        self.assertIn('noise_pred_net.proj.weight', optimizer_names)
+        self.assertNotIn('condition_encoder.vlm.weight', optimizer_names)
+        self.assertNotIn('condition_encoder.vlm.bias', optimizer_names)
+
+    def test_smolvla_native_training_preset_overrides_optimizer_scheduler_and_grad_clip(self):
+        module = TrainVLATransformerOptimizerTest()._load_train_vla_module()
+
+        class _NativeAgent(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.model = nn.Linear(2, 2)
+
+            def to(self, device):
+                return self
+
+            def get_normalization_stats(self):
+                return {}
+
+        agent = _NativeAgent()
+        cfg = TrainVLATransformerOptimizerTest()._make_cfg()
+        cfg.agent = AttrDict(_target_='roboimi.vla.agent_smolvla_native.SmolVLANativeAgent')
+        cfg.train.lr = 9e-4
+        cfg.train.max_steps = 1
+        cfg.train.weight_decay = 0.123
+        cfg.train.grad_clip = 1.0
+        cfg.train.warmup_steps = 7
+        cfg.train.scheduler_type = 'constant'
+        cfg.train.min_lr = 1e-7
+
+        scheduler_calls = []
+        clip_calls = []
+
+        def fake_instantiate(config_node, **_kwargs):
+            if config_node is cfg.data:
+                return FakeDataset()
+            if config_node is cfg.agent:
+                return agent
+            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
+
+        def fake_scheduler(*args, **kwargs):
+            scheduler_calls.append(kwargs)
+            return FakeScheduler()
+
+        def fake_clip(parameters, max_norm):
+            clip_calls.append(float(max_norm))
+            return torch.tensor(0.0)
+
+        class OneBatchLoader:
+            def __len__(self):
+                return 1
+
+            def __iter__(self):
+                batch = {
+                    'observation.front': torch.zeros(1, 1, 1, 2, 2),
+                    'observation.state': torch.zeros(1, 1, 2),
+                    'action': torch.zeros(1, 1, 2),
+                }
+                return iter([batch])
+
+        def fake_compute_loss(_batch):
+            return agent.model.weight.sum() * 0.0 + torch.tensor(1.0, requires_grad=True)
+
+        agent.compute_loss = fake_compute_loss
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \
+                     mock.patch.object(module, 'DataLoader', side_effect=lambda *args, **kwargs: OneBatchLoader()), \
+                     mock.patch.object(module, 'get_lr_schedule_with_warmup', side_effect=fake_scheduler), \
+                     mock.patch.object(module, 'AdamW', RecordingAdamW), \
+                     mock.patch.object(module.torch.nn.utils, 'clip_grad_norm_', side_effect=fake_clip), \
+                     mock.patch.object(module.torch, 'save', return_value=None), \
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
+                    module.main(cfg)
+            finally:
+                os.chdir(previous_cwd)
+
+        optimizer = RecordingAdamW.created[-1]
+        self.assertEqual(optimizer.lr, 1e-4)
+        self.assertEqual(optimizer.weight_decay, 1e-10)
+        self.assertEqual(optimizer.betas, (0.9, 0.95))
+        self.assertEqual(optimizer.eps, 1e-8)
+        self.assertEqual(scheduler_calls[-1], {
+            'warmup_steps': 1000,
+            'max_steps': 1,
+            'scheduler_type': 'cosine',
+            'min_lr': 2.5e-6,
+        })
+        self.assertEqual(clip_calls, [10.0])
+
+    def test_cosine_scheduler_spans_requested_training_steps_then_clamps(self):
+        module = TrainVLATransformerOptimizerTest()._load_train_vla_module()
+        param = nn.Parameter(torch.tensor(1.0))
+        optimizer = torch.optim.SGD([param], lr=1e-4)
+        base_lr = optimizer.param_groups[0]['lr']
+
+        scheduler = module.get_lr_schedule_with_warmup(
+            optimizer,
+            warmup_steps=1000,
+            max_steps=150000,
+            scheduler_type='cosine',
+            min_lr=2.5e-6,
+        )
+
+        lr_lambda = scheduler.lr_lambdas[0]
+        observed = {}
+        for step in (0, 1, 1000, 30000, 40000, 150000, 160000):
+            observed[step] = base_lr * lr_lambda(step)
+
+        self.assertGreater(observed[1], observed[0])
+        self.assertLess(observed[1000], 1e-4)
+        self.assertGreater(observed[30000], observed[40000])
+        self.assertGreater(observed[40000], observed[150000])
+        self.assertAlmostEqual(observed[150000], 2.5e-6, places=12)
+        self.assertAlmostEqual(observed[160000], 2.5e-6, places=12)