Add inverse dynamics support

cosmos-framework/cosmos_framework/data/vfm/action/action_viz/app.py

@@ -39,6 +39,7 @@ from cosmos_framework.data.vfm.action.action_viz.state import (
     GenerationResult,
     control_points_from_action,
     make_generation_id,
+    model_mode_generates_action,
     read_generated_action,
 )
 from cosmos_framework.data.vfm.action.urdf_visualizer.unified_action import ActionFormat, get_video_from_sample
@@ -284,6 +285,7 @@ def launch_action_viz(config: AppConfig) -> None:
             num_steps_input = client.gui.add_number("Sampling steps", initial_value=DEFAULT_SAMPLING_STEPS, min=1, step=1)
             guidance_input = client.gui.add_number("Guidance", initial_value=DEFAULT_GUIDANCE, min=0.0, max=7.0, step=0.1)
             generate_button = client.gui.add_button("Run forward dynamics")
+            inverse_dynamics_button = client.gui.add_button("Run inverse dynamics")
             policy_button = client.gui.add_button("Run policy")
             status_text = client.gui.add_markdown(f"*{worker_client.status_message}*")
             info_text = client.gui.add_markdown("")
@@ -300,6 +302,7 @@ def launch_action_viz(config: AppConfig) -> None:
         with client.gui.add_folder("Video"):
             gt_panel = client.gui.add_image(np.zeros((64, 64, 3), dtype=np.uint8), label="Ground truth")
             generated_panel = client.gui.add_image(np.zeros((64, 64, 3), dtype=np.uint8), label="Generated")
+            generated_panel.visible = False
 
         with client.gui.add_folder("Trajectory editor"):
             target_dropdown = client.gui.add_dropdown("Trajectory", options=("right", "left"), initial_value="right")
@@ -696,8 +699,12 @@ def launch_action_viz(config: AppConfig) -> None:
                 return
             if loaded.video is not None and t < len(loaded.video):
                 gt_panel.image = loaded.video[t]
-            if loaded.generated_video:
-                generated_panel.image = loaded.generated_video[min(t, len(loaded.generated_video) - 1)]
+            generated_frame = _select_generated_video_frame(loaded.generated_video, t)
+            if generated_frame is None:
+                generated_panel.visible = False
+            else:
+                generated_panel.image = generated_frame
+                generated_panel.visible = True
 
         def _update_action_text(t: int) -> None:
             loaded = sample_state["loaded"]
@@ -750,6 +757,7 @@ def launch_action_viz(config: AppConfig) -> None:
             _flush_scene_update()
             loaded.generated_video = []
             generated_panel.image = np.zeros((64, 64, 3), dtype=np.uint8)
+            generated_panel.visible = False
 
             generation_id = make_generation_id()
             client_generation_root = config.output_root / "active" / f"client_{client.client_id}"
@@ -802,13 +810,16 @@ def launch_action_viz(config: AppConfig) -> None:
                     )
 
                 try:
-                    if result.status == "success" and result.video_path:
-                        loaded.generated_video = _load_video_frames(Path(result.video_path))
-                        if model_mode == "policy" and result.generated_action_path:
+                    if result.status == "success":
+                        if result.video_path:
+                            loaded.generated_video = _load_video_frames(Path(result.video_path))
+                        else:
+                            loaded.generated_video = []
+                        if model_mode_generates_action(model_mode):
                             try:
-                                _replace_loaded_action(loaded, read_generated_action(Path(result.generated_action_path)))
-                            except Exception as exc:
-                                status_text.content = f"*Generation {generation_id} loaded video; action ignored: {exc}*"
+                                _replace_loaded_action(loaded, _read_generated_result_action(result))
+                            except ValueError as exc:
+                                status_text.content = f"*Generation {generation_id} failed: {exc}*"
                             else:
                                 status_text.content = f"*Generation {generation_id} complete.*"
                         else:
@@ -846,6 +857,10 @@ def launch_action_viz(config: AppConfig) -> None:
         def _(_) -> None:
             run_generation("forward_dynamics")
 
+        @inverse_dynamics_button.on_click
+        def _(_) -> None:
+            run_generation("inverse_dynamics")
+
         @policy_button.on_click
         def _(_) -> None:
             run_generation("policy")
@@ -1021,6 +1036,22 @@ def _set_loaded_action(loaded: LoadedSample, raw_action: np.ndarray) -> None:
     loaded.baked_action = raw_action.astype(np.float32).copy()
 
 
+def _read_generated_result_action(result: GenerationResult) -> np.ndarray:
+    """Read the raw action emitted by an action-generating model result."""
+
+    if result.generated_action_path is None:
+        raise ValueError(f"Generation {result.generation_id} has no saved generated action")
+    return read_generated_action(Path(result.generated_action_path))
+
+
+def _select_generated_video_frame(generated_video: list[np.ndarray], t: int) -> np.ndarray | None:
+    """Return the displayed generated-video frame, or None when no video exists."""
+
+    if not generated_video:
+        return None
+    return generated_video[min(max(int(t), 0), len(generated_video) - 1)]
+
+
 def _build_render_scene(loaded: LoadedSample, *, source_robot_animation: bool) -> Any:
     """Build a render scene from the loaded sample's current editor state."""
 

cosmos-framework/cosmos_framework/data/vfm/action/action_viz/app_test.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from cosmos_framework.data.vfm.action.action_viz.app import (
+    _read_generated_result_action,
+    _select_generated_video_frame,
+)
+from cosmos_framework.data.vfm.action.action_viz.state import (
+    BRIDGE_ACTION_DIM,
+    GenerationResult,
+    write_generated_action,
+)
+
+
+@pytest.mark.L0
+def test_read_generated_result_action_supports_inverse_dynamics(tmp_path) -> None:
+    action = np.ones((2, BRIDGE_ACTION_DIM), dtype=np.float32)
+    action_path = write_generated_action(action, tmp_path)
+    result = GenerationResult(
+        generation_id="gen",
+        model_mode="inverse_dynamics",
+        status="success",
+        generated_action_path=str(action_path),
+    )
+
+    loaded = _read_generated_result_action(result)
+
+    np.testing.assert_allclose(loaded, action, atol=1e-6)
+
+
+@pytest.mark.L0
+def test_read_generated_result_action_requires_saved_action_path() -> None:
+    result = GenerationResult(generation_id="gen", model_mode="inverse_dynamics", status="success")
+
+    with pytest.raises(ValueError, match="Generation gen has no saved generated action"):
+        _read_generated_result_action(result)
+
+
+@pytest.mark.L0
+def test_select_generated_video_frame_returns_none_for_action_only_generations() -> None:
+    assert _select_generated_video_frame([], 0) is None
+
+
+@pytest.mark.L0
+def test_select_generated_video_frame_clamps_to_available_frames() -> None:
+    first = np.zeros((2, 2, 3), dtype=np.uint8)
+    last = np.ones((2, 2, 3), dtype=np.uint8)
+
+    np.testing.assert_array_equal(_select_generated_video_frame([first, last], -1), first)
+    np.testing.assert_array_equal(_select_generated_video_frame([first, last], 99), last)

cosmos-framework/cosmos_framework/data/vfm/action/action_viz/local_worker.py

@@ -21,6 +21,7 @@ import numpy as np
 from cosmos_framework.data.vfm.action.action_viz.state import (
     GenerationRequest,
     GenerationResult,
+    model_mode_generates_action,
     request_to_json_dict,
     write_generated_action,
     write_generation_request,
@@ -957,7 +958,7 @@ def _copy_generation_outputs(
 
     generated_video_path = None
     vision_path = sample_output_dir / "vision.mp4"
-    if vision_path.is_file():
+    if request.model_mode != "inverse_dynamics" and vision_path.is_file():
         generated_video_path = output_dir / "generated.mp4"
         shutil.copyfile(vision_path, generated_video_path)
 
@@ -975,6 +976,8 @@ def _copy_generation_outputs(
                 to_model_space=False,
             )
             generated_action_path = write_generated_action(np.asarray(raw_action, dtype=np.float32), output_dir)
+    if model_mode_generates_action(request.model_mode) and generated_action_path is None:
+        raise ValueError(f"{request.model_mode} generation did not return an action")
 
     result_payload: dict[str, Any] = {
         "ok": True,

cosmos-framework/cosmos_framework/data/vfm/action/action_viz/local_worker_test.py

@@ -1,10 +1,22 @@
 from __future__ import annotations
 
+import json
+from pathlib import Path
 from typing import Any
 
+import numpy as np
 import pytest
 
-from cosmos_framework.data.vfm.action.action_viz.local_worker import _patched_unipc_progress
+from cosmos_framework.data.vfm.action.action_viz.local_worker import (
+    _copy_generation_outputs,
+    _patched_unipc_progress,
+)
+from cosmos_framework.data.vfm.action.action_viz.state import (
+    BRIDGE_ACTION_DIM,
+    ControlPoint,
+    GenerationRequest,
+    read_generated_action,
+)
 
 
 @pytest.mark.L0
@@ -24,3 +36,96 @@ def test_patched_unipc_progress_reports_sampler_percentages(monkeypatch: pytest.
 
     assert updates == [(37, "sampling"), (63, "sampling"), (90, "sampling")]
     assert unipc.progress_bar is passthrough_progress_bar
+
+
+@pytest.mark.L0
+def test_copy_generation_outputs_writes_inverse_dynamics_action_without_copying_auxiliary_video(tmp_path) -> None:
+    request = _request(tmp_path, "inverse_dynamics")
+    inference_dir = tmp_path / "framework_output"
+    sample_output_dir = inference_dir / request.generation_id
+    sample_output_dir.mkdir(parents=True)
+    sample_outputs = {
+        "outputs": [
+            {
+                "content": {
+                    "action": [
+                        [0.0] * BRIDGE_ACTION_DIM,
+                        [0.0] * BRIDGE_ACTION_DIM,
+                    ]
+                }
+            }
+        ]
+    }
+    (sample_output_dir / "sample_outputs.json").write_text(json.dumps(sample_outputs), encoding="utf-8")
+    (sample_output_dir / "vision.mp4").write_bytes(b"auxiliary video")
+
+    payload, video_path, action_path = _copy_generation_outputs(request, inference_dir, Path(request.output_dir))
+
+    assert payload["artifacts"]["video_filename"] is None
+    assert payload["artifacts"]["action_filename"] == "generated_action.json"
+    assert video_path is None
+    assert action_path is not None
+    action = read_generated_action(Path(action_path))
+    assert action.shape == (2, BRIDGE_ACTION_DIM)
+
+
+@pytest.mark.L0
+def test_copy_generation_outputs_keeps_policy_video(tmp_path) -> None:
+    request = _request(tmp_path, "policy")
+    output_dir = Path(request.output_dir)
+    output_dir.mkdir(parents=True)
+    inference_dir = tmp_path / "framework_output"
+    sample_output_dir = inference_dir / request.generation_id
+    sample_output_dir.mkdir(parents=True)
+    sample_outputs = {
+        "outputs": [
+            {
+                "content": {
+                    "action": [
+                        [0.0] * BRIDGE_ACTION_DIM,
+                    ]
+                }
+            }
+        ]
+    }
+    (sample_output_dir / "sample_outputs.json").write_text(json.dumps(sample_outputs), encoding="utf-8")
+    (sample_output_dir / "vision.mp4").write_bytes(b"generated video")
+
+    payload, video_path, action_path = _copy_generation_outputs(request, inference_dir, output_dir)
+
+    assert payload["artifacts"]["video_filename"] == "generated.mp4"
+    assert video_path == str(output_dir / "generated.mp4")
+    assert Path(video_path).read_bytes() == b"generated video"
+    assert action_path is not None
+
+
+@pytest.mark.L0
+def test_copy_generation_outputs_requires_action_for_inverse_dynamics(tmp_path) -> None:
+    request = _request(tmp_path, "inverse_dynamics")
+    inference_dir = tmp_path / "framework_output"
+    sample_output_dir = inference_dir / request.generation_id
+    sample_output_dir.mkdir(parents=True)
+    (sample_output_dir / "sample_outputs.json").write_text(
+        json.dumps({"outputs": [{"content": {}}]}),
+        encoding="utf-8",
+    )
+
+    with pytest.raises(ValueError, match="inverse_dynamics generation did not return an action"):
+        _copy_generation_outputs(request, inference_dir, Path(request.output_dir))
+
+
+def _request(tmp_path, model_mode: str) -> GenerationRequest:
+    return GenerationRequest(
+        generation_id="gen",
+        model_mode=model_mode,
+        dataset="bridge",
+        sample_index=0,
+        experiment_name="",
+        s3_checkpoint_dir="nvidia/Cosmos3-Nano",
+        checkpoint_cache_dir=None,
+        output_dir=str(tmp_path / "output"),
+        seed=0,
+        num_steps=1,
+        control_points=[ControlPoint(frame=0, values=[0.0] * BRIDGE_ACTION_DIM)],
+        baked_action=np.zeros((1, BRIDGE_ACTION_DIM), dtype=np.float32).tolist(),
+    )

cosmos-framework/cosmos_framework/data/vfm/action/action_viz/state.py

@@ -24,7 +24,8 @@ BRIDGE_ACTION_CHANNELS = (
     "gripper",
 )
 BRIDGE_DATASET_SELECTOR = "bridge_20260501"
-GENERATION_MODEL_MODES = ("forward_dynamics", "policy")
+GENERATION_MODEL_MODES = ("forward_dynamics", "inverse_dynamics", "policy")
+ACTION_GENERATING_MODEL_MODES = ("inverse_dynamics", "policy")
 
 
 @dataclass(frozen=True)
@@ -189,7 +190,7 @@ def read_generation_request(path: Path) -> GenerationRequest:
     return request_from_json_dict(json.loads(path.read_text()))
 
 
-def write_generated_action(action: np.ndarray, output_dir: Path, filename: str = "policy_action.json") -> Path:
+def write_generated_action(action: np.ndarray, output_dir: Path, filename: str = "generated_action.json") -> Path:
     """Write a generated raw action array and return its path."""
 
     action = _validate_generated_action_array(action)
@@ -231,6 +232,12 @@ def read_generated_action(path: Path) -> np.ndarray:
     return _validate_generated_action_array(np.asarray(json.loads(path.read_text()), dtype=np.float32))
 
 
+def model_mode_generates_action(model_mode: str) -> bool:
+    """Return whether a supported generation mode emits a raw action trajectory."""
+
+    return _validate_model_mode(model_mode) in ACTION_GENERATING_MODEL_MODES
+
+
 def _validate_action_array(action: np.ndarray, action_dim: int | None = None) -> np.ndarray:
     action = np.asarray(action, dtype=np.float32)
     expected_dim = BRIDGE_ACTION_DIM if action_dim is None else int(action_dim)

cosmos-framework/cosmos_framework/data/vfm/action/action_viz/state_test.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from cosmos_framework.data.vfm.action.action_viz.state import (
+    BRIDGE_ACTION_DIM,
+    model_mode_generates_action,
+    read_generated_action,
+    write_generated_action,
+)
+
+
+@pytest.mark.L0
+@pytest.mark.parametrize(
+    ("model_mode", "generates_action"),
+    [
+        ("forward_dynamics", False),
+        ("inverse_dynamics", True),
+        ("policy", True),
+    ],
+)
+def test_model_mode_generates_action(model_mode: str, generates_action: bool) -> None:
+    assert model_mode_generates_action(model_mode) is generates_action
+
+
+@pytest.mark.L0
+def test_generated_action_round_trip_uses_generic_filename(tmp_path) -> None:
+    action = np.arange(2 * BRIDGE_ACTION_DIM, dtype=np.float32).reshape(2, BRIDGE_ACTION_DIM)
+
+    path = write_generated_action(action, tmp_path)
+
+    assert path.name == "generated_action.json"
+    np.testing.assert_allclose(read_generated_action(path), action, atol=1e-6)