world_model.py

#!/usr/bin/env python3
r"""WorldModel: pluggable next-frame + reward predictor for MARL in Isaac Sim.

Backends:
- ``mock``     : tiny learned Conv/MLP dynamics in PyTorch (falls back to a numpy
  random-feature ridge regressor when torch is unavailable).
- ``cosmos3``  : NVIDIA Cosmos3 world foundation model.  Three execution paths,
  tried in order:
  1. **NIM HTTP API** – if ``cosmos_nim_url`` is set (e.g.
     ``http://localhost:8000/v1/infer``) a POST request with the conditioning
     image + prompt is sent to the running Cosmos3-Generator NIM server.
  2. **Direct Python import** – if ``cosmos_framework`` is importable in the
     current interpreter, the inference pipeline is loaded in-process.
  3. **Subprocess** – falls back to ``python -m
     cosmos_framework.scripts.inference`` with a JSON input file using
     ``forward_dynamics`` (action-conditioned) or ``image2video`` mode.
  If none of the paths succeed (missing checkpoint / deps / server), the
  backend transparently falls back to ``mock``.
- ``dreamdojo``: DreamDojo / Cosmos-Predict2.5 action-conditioned robot world
  model.  Two execution paths:
  1. **Direct Python import** – if ``cosmos_predict2`` is importable, the
     ``ActionStreamingInference`` class is used in-process.
  2. **Subprocess** – ``python -m cosmos_predict2._src.predict2.interactive.
     inference.action_video2world_teleop`` with ``--input_frame`` and
     ``--action_file``.
  Falls back to ``mock`` if deps / checkpoints are missing.

The class is designed to run from ``H:\Robotics\ISSAC_SIM_5.1.0\python.bat``
(where torch is present) as well as from a minimal Python with only numpy/cv2.
"""

from __future__ import annotations

import argparse
import base64
import json
import logging
import math
import os
import shutil
import subprocess
import sys
import tempfile
import time
import warnings
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np

logger = logging.getLogger("world_model")

# ---------------------------------------------------------------------------
# Optional dependencies (kept conditional so the file is importable everywhere)
# ---------------------------------------------------------------------------
try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    _HAS_TORCH = True
except Exception:
    _HAS_TORCH = False

    # Provide harmless stand-ins so the torch-only class can still be defined
    # when torch is not installed. The class is never instantiated in that case.
    class _DummyObject:
        def __init__(self, *args, **kwargs):
            pass

        def __call__(self, *args, **kwargs):
            return self

        def __getattr__(self, name):
            return _DummyObject

        def __iter__(self):
            return iter([])

    class _DummyNN:
        class Module:
            def __init__(self):
                pass

            def to(self, device):
                return self

            def parameters(self):
                return []

        def __getattr__(self, name):
            return _DummyObject

    nn = _DummyNN()
    F = _DummyObject()

try:
    import cv2

    _HAS_CV2 = True
except Exception:
    _HAS_CV2 = False

try:
    from PIL import Image as PILImage

    _HAS_PIL = True
except Exception:
    _HAS_PIL = False

try:
    import imageio

    _HAS_IMAGEIO = True
except Exception:
    _HAS_IMAGEIO = False

try:
    import requests

    _HAS_REQUESTS = True
except Exception:
    _HAS_REQUESTS = False

_HAS_UV = shutil.which("uv") is not None

# ---------------------------------------------------------------------------
# Paths for the local workspace (see H:\Robotics\AGENTS.md)
# ---------------------------------------------------------------------------
ROOT = Path(r"H:\Robotics")
COSMOS_ROOT = ROOT / "cosmos-framework"
DREAMDOJO_ROOT = ROOT / "DreamDojo"
DEFAULT_OUTPUT_DIR = ROOT / "marl_world_model" / "outputs"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _to_numpy(a: Any) -> np.ndarray:
    """Convert an observation / action / reward to a numpy array."""
    if isinstance(a, np.ndarray):
        return a
    if _HAS_TORCH and isinstance(a, torch.Tensor):
        return a.detach().cpu().numpy()
    if isinstance(a, (list, tuple)):
        return np.array(a)
    return np.array(a)


def _split_observation(observation: Any) -> Tuple[np.ndarray, np.ndarray, bool]:
    """Return (image, state, is_dict) from an observation.

    Supports either a raw image array (H, W, C) or a dict with
    ``image`` and optional ``state`` keys.
    """
    if isinstance(observation, dict):
        image = _to_numpy(observation["image"])
        state = _to_numpy(observation.get("state", np.zeros(0, dtype=np.float32)))
        is_dict = True
    else:
        image = _to_numpy(observation)
        state = np.zeros(0, dtype=np.float32)
        is_dict = False

    if image.ndim == 2:
        image = np.stack([image] * 3, axis=-1)
    if image.dtype != np.uint8:
        if image.max() <= 1.0:
            image = (image * 255).astype(np.uint8)
        else:
            image = image.astype(np.uint8)
    return image, state, is_dict


def _reconstruct_observation(
    image: np.ndarray,
    state: np.ndarray,
    is_dict: bool,
    action: np.ndarray,
    state_integration: str,
) -> Union[np.ndarray, Dict[str, Any]]:
    """Build the next observation in the same format as the input."""
    next_state = state
    if is_dict and state.size > 0 and state_integration == "add" and action.shape[0] >= state.shape[0]:
        next_state = state + action[: state.shape[0]].astype(state.dtype)
    if is_dict:
        return {"image": image, "state": next_state}
    return image


def _save_image(image: np.ndarray, path: Union[str, Path]) -> str:
    """Save an RGB uint8 image to disk."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    if _HAS_CV2:
        bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) if image.shape[-1] == 3 else image
        cv2.imwrite(str(path), bgr)
    elif _HAS_PIL:
        PILImage.fromarray(image).save(str(path))
    else:
        raise RuntimeError("No cv2/PIL available; cannot save image")
    return str(path)


def _load_image(path: Union[str, Path]) -> np.ndarray:
    """Load an image as RGB uint8."""
    if _HAS_CV2:
        img = cv2.imread(str(path))
        if img is None:
            raise RuntimeError(f"cv2 could not read {path}")
        if img.ndim == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        elif img.shape[-1] == 3:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return img
    if _HAS_PIL:
        return np.array(PILImage.open(str(path)).convert("RGB"))
    raise RuntimeError("No cv2/PIL available; cannot load image")


def _resize_image(image: np.ndarray, size: Tuple[int, int]) -> np.ndarray:
    """Resize to ``size=(width, height)``."""
    if _HAS_CV2:
        return cv2.resize(image, size)
    if _HAS_PIL:
        return np.array(PILImage.fromarray(image).resize(size))
    raise RuntimeError("No cv2/PIL available; cannot resize image")


def _read_video_frames(path: Union[str, Path]) -> List[np.ndarray]:
    """Return all frames of a video as a list of RGB uint8 images."""
    if _HAS_CV2:
        cap = cv2.VideoCapture(str(path))
        if not cap.isOpened():
            raise RuntimeError(f"Could not open video {path}")
        frames = []
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            if frame.ndim == 2:
                frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
            elif frame.shape[-1] == 3:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)
        cap.release()
        return frames
    if _HAS_IMAGEIO:
        reader = imageio.get_reader(str(path))
        return [np.array(f) for f in reader]
    raise RuntimeError("No video reader available")


def _run_command(cmd: List[str], cwd: Optional[Union[str, Path]], timeout: float) -> str:
    """Run a command and return its stdout, raising on non-zero exit."""
    logger.info("Running command: %s in %s", " ".join(cmd), cwd or ".")
    try:
        result = subprocess.run(
            cmd,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            timeout=timeout,
        )
    except subprocess.TimeoutExpired as e:
        raise RuntimeError(f"Command timed out after {timeout}s: {' '.join(cmd)}\n{e.output or ''}")
    if result.returncode != 0:
        raise RuntimeError(f"Command failed (code {result.returncode}): {result.stdout[:2000]}")
    logger.debug(result.stdout[:1000])
    return result.stdout


# ---------------------------------------------------------------------------
# Tiny learned dynamics models used by the ``mock`` backend
# ---------------------------------------------------------------------------
class TinyTorchWorldModel(nn.Module):
    """Small Conv/GRU/MLP world model for fast testing.

    Architecture:
        - 2-layer conv encoder (avg pooled to 64-d)
        - action MLP projection
        - single-layer GRU core
        - MLP that directly predicts the next image (C, H, W)
        - reward and done heads
    """

    def __init__(self, image_shape: Tuple[int, int, int], action_dim: int, hidden: int = 64, device: str = "cpu"):
        super().__init__()
        c, h, w = image_shape
        self.image_shape = image_shape
        self.action_dim = action_dim
        self.hidden = hidden
        self.device = device

        self.conv1 = nn.Conv2d(c, 32, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
        self.pool = nn.AdaptiveAvgPool2d(1)

        self.action_encoder = nn.Linear(action_dim, hidden)
        self.core = nn.GRUCell(64 + hidden, hidden)

        self.next_image_head = nn.Sequential(
            nn.Linear(hidden, 256),
            nn.ReLU(),
            nn.Linear(256, c * h * w),
            nn.Sigmoid(),
        )
        self.reward_head = nn.Linear(hidden, 1)
        self.done_head = nn.Linear(hidden, 1)

        # Per-agent hidden state so multi-agent rollouts do not cross-contaminate.
        self._h: Dict[Optional[int], Optional[torch.Tensor]] = {}
        self._last_done: Dict[Optional[int], bool] = {}
        self.to(device)

    def forward(
        self,
        image: torch.Tensor,
        action: torch.Tensor,
        h: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        b = image.shape[0]
        x = F.relu(self.conv1(image))
        x = F.relu(self.conv2(x))
        x = self.pool(x).view(b, -1)
        a = F.relu(self.action_encoder(action))
        cat = torch.cat([x, a], dim=1)
        if h is None:
            h = torch.zeros(b, self.hidden, device=image.device)
        h = self.core(cat, h)
        next_image = self.next_image_head(h).view(b, *self.image_shape)
        reward = self.reward_head(h).squeeze(-1)
        done_logit = self.done_head(h).squeeze(-1)
        return next_image, reward, done_logit, h

    def imagine(
        self, image: np.ndarray, action: np.ndarray, agent_id: Optional[int] = None
    ) -> Tuple[np.ndarray, float, bool]:
        if agent_id not in self._last_done or self._last_done.get(agent_id, True):
            self._h[agent_id] = None
            self._last_done[agent_id] = False

        h = self._h.get(agent_id)
        img_t = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float().to(self.device) / 255.0
        act_t = torch.from_numpy(action).float().unsqueeze(0).to(self.device)

        with torch.no_grad():
            next_t, reward_t, done_logit, h = self.forward(img_t, act_t, h)

        self._h[agent_id] = h
        next_image = (next_t.squeeze(0).permute(1, 2, 0).detach().cpu().numpy() * 255).astype(np.uint8)
        reward = float(reward_t.detach().cpu().numpy()[0])
        done = float(torch.sigmoid(done_logit).detach().cpu().numpy()[0]) > 0.5
        if done:
            self._last_done[agent_id] = True
        return next_image, reward, done

    def train(
        self,
        obs_list: List[np.ndarray],
        act_list: List[np.ndarray],
        next_list: List[np.ndarray],
        reward_list: List[float],
        done_list: List[float],
        n_epochs: int = 5,
        lr: float = 1e-3,
    ) -> None:
        n = len(obs_list)
        if n == 0:
            return

        obs = np.stack(obs_list).astype(np.float32) / 255.0
        nxt = np.stack(next_list).astype(np.float32) / 255.0
        act = np.stack(act_list).astype(np.float32)
        rew = np.array(reward_list, dtype=np.float32)
        done = np.array(done_list, dtype=np.float32)

        obs_t = torch.from_numpy(obs).permute(0, 3, 1, 2).to(self.device)
        nxt_t = torch.from_numpy(nxt).permute(0, 3, 1, 2).to(self.device)
        act_t = torch.from_numpy(act).to(self.device)
        rew_t = torch.from_numpy(rew).to(self.device)
        done_t = torch.from_numpy(done).to(self.device)

        optimizer = torch.optim.Adam(self.parameters(), lr=lr)
        batch_size = min(8, n)
        for _ in range(n_epochs):
            perm = torch.randperm(n)
            for i in range(0, n, batch_size):
                idx = perm[i : i + batch_size]
                pred, r, d, _ = self.forward(obs_t[idx], act_t[idx], None)
                loss_img = F.mse_loss(pred, nxt_t[idx])
                loss_r = F.mse_loss(r, rew_t[idx])
                loss_d = F.binary_cross_entropy_with_logits(d, done_t[idx])
                loss = loss_img + 0.1 * loss_r + 0.1 * loss_d
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

        self._h.clear()
        self._last_done.clear()


class TinyNumpyWorldModel:
    """Random-feature ridge-regression world model used when torch is absent.

    The image is flattened and passed through a fixed random projection; a small
    ridge regression is solved on the projected features to predict the next
    image residual, reward, and done probability.
    """

    def __init__(
        self,
        image_shape: Tuple[int, int, int],
        action_dim: int,
        feature_dim: int = 64,
        alpha: float = 1e-3,
        seed: int = 42,
    ):
        self.image_shape = image_shape
        self.action_dim = action_dim
        self.feature_dim = feature_dim
        self.alpha = alpha
        self.d_img = int(np.prod(image_shape))

        rng = np.random.default_rng(seed)
        scale_i = 1.0 / math.sqrt(self.d_img)
        scale_a = 1.0 / math.sqrt(max(action_dim, 1))
        self.W_i = rng.normal(0.0, scale_i, (self.d_img, feature_dim)).astype(np.float32)
        self.W_a = rng.normal(0.0, scale_a, (action_dim, feature_dim)).astype(np.float32)
        self.b = rng.normal(0.0, 0.1, feature_dim).astype(np.float32)

        self.W_img: Optional[np.ndarray] = None
        self.W_reward: Optional[np.ndarray] = None
        self.W_done: Optional[np.ndarray] = None

    def _features(self, image: np.ndarray, action: np.ndarray) -> np.ndarray:
        img_flat = image.reshape(1, self.d_img).astype(np.float32) / 255.0
        act = action.reshape(1, self.action_dim).astype(np.float32)
        z = img_flat @ self.W_i + act @ self.W_a + self.b
        return np.maximum(z, 0.0)

    def imagine(self, image: np.ndarray, action: np.ndarray, agent_id: Optional[int] = None) -> Tuple[np.ndarray, float, bool]:
        # agent_id is accepted for API symmetry but the numpy mock is stateless.
        phi = self._features(image, action)
        img_flat = image.reshape(1, self.d_img).astype(np.float32) / 255.0

        if self.W_img is None:
            residual = np.zeros_like(img_flat)
        else:
            residual = phi @ self.W_img

        next_img = ((img_flat + residual) * 255.0).clip(0, 255).reshape(self.image_shape).astype(np.uint8)

        if self.W_reward is None:
            reward = 0.0
        else:
            reward = float((phi @ self.W_reward)[0, 0])

        if self.W_done is None:
            done_logit = -5.0
        else:
            done_logit = float((phi @ self.W_done)[0, 0])
        done = (1.0 / (1.0 + math.exp(-done_logit))) > 0.5
        return next_img, reward, done

    @staticmethod
    def _ridge_solve(A: np.ndarray, B: np.ndarray) -> np.ndarray:
        try:
            return np.linalg.solve(A, B)
        except np.linalg.LinAlgError:
            return np.linalg.pinv(A) @ B

    def train(
        self,
        obs_list: List[np.ndarray],
        act_list: List[np.ndarray],
        next_list: List[np.ndarray],
        reward_list: List[float],
        done_list: List[float],
    ) -> None:
        n = len(obs_list)
        if n == 0:
            return

        obs = np.stack(obs_list).astype(np.float32)
        nxt = np.stack(next_list).astype(np.float32)
        act = np.stack(act_list).astype(np.float32)

        X = np.zeros((n, self.feature_dim), dtype=np.float32)
        for i in range(n):
            X[i] = self._features(obs[i], act[i]).ravel()

        Y_img = ((nxt - obs).reshape(n, -1).astype(np.float32)) / 255.0
        Y_reward = np.array(reward_list, dtype=np.float32).reshape(n, 1)
        Y_done = np.array(done_list, dtype=np.float32).reshape(n, 1)

        A = X.T @ X + self.alpha * np.eye(self.feature_dim, dtype=np.float32)

        self.W_img = self._ridge_solve(A, X.T @ Y_img)
        self.W_reward = self._ridge_solve(A, X.T @ Y_reward)
        self.W_done = self._ridge_solve(A, X.T @ Y_done)


# ---------------------------------------------------------------------------
# Main wrapper
# ---------------------------------------------------------------------------
class WorldModel:
    """Pluggable world model wrapper.

    Parameters
    ----------
    backend:
        One of ``"mock"``, ``"cosmos3"``, ``"dreamdojo"``.
    cosmos_root / dreamdojo_root / output_dir:
        Local workspace paths.
    device:
        ``"auto"`` (torch-only), ``"cpu"``, or ``"cuda"``. Ignored by the numpy
        fallback.
    feature_dim:
        Hidden feature size for the numpy mock model.
    fallback_to_mock:
        If True, any failure in ``cosmos3`` or ``dreamdojo`` silently switches to
        the ``mock`` backend and retries ``imagine()``.
    state_integration:
        How to update a dict ``state`` component: ``"add"`` adds the first
        ``len(state)`` action components, ``"none"`` leaves it unchanged.
    **config:
        Backend-specific overrides, e.g. ``cosmos_checkpoint``,
        ``cosmos_prompt_template``, ``dreamdojo_experiment``,
        ``dreamdojo_ckpt``, ``cosmos_timeout``.
    """

    VALID_BACKENDS = {"mock", "cosmos3", "dreamdojo"}

    def __init__(
        self,
        backend: str = "mock",
        cosmos_root: Union[str, Path] = COSMOS_ROOT,
        dreamdojo_root: Union[str, Path] = DREAMDOJO_ROOT,
        output_dir: Union[str, Path] = DEFAULT_OUTPUT_DIR,
        device: str = "auto",
        feature_dim: int = 64,
        fallback_to_mock: bool = True,
        state_integration: str = "add",
        **config: Any,
    ):
        if backend not in self.VALID_BACKENDS:
            raise ValueError(f"Unknown backend {backend!r}; expected one of {self.VALID_BACKENDS}")

        self.backend = backend
        self.original_backend = backend
        self.cosmos_root = Path(cosmos_root)
        self.dreamdojo_root = Path(dreamdojo_root)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.config = config
        self.fallback_to_mock = fallback_to_mock
        self.state_integration = state_integration
        self.feature_dim = feature_dim

        if device == "auto" and _HAS_TORCH:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        self._mock: Optional[Union[TinyTorchWorldModel, TinyNumpyWorldModel]] = None
        self._log_counter = 0

    # -----------------------------------------------------------------------
    # Public API
    # -----------------------------------------------------------------------
    def train(
        self,
        real_observations: List[Any],
        real_actions: List[Any],
        real_rewards: List[float],
    ) -> None:
        """Train / log real transition data.

        For ``cosmos3`` and ``dreamdojo`` this only writes the data to disk
        (checkpoints are not downloaded). For ``mock`` this updates the tiny
        learned dynamics model.
        """
        if self.backend == "mock":
            self._mock_train(real_observations, real_actions, real_rewards)
        elif self.backend == "cosmos3":
            self._cosmos_train(real_observations, real_actions, real_rewards)
        elif self.backend == "dreamdojo":
            self._dreamdojo_train(real_observations, real_actions, real_rewards)
        else:
            raise ValueError(f"Unknown backend: {self.backend}")

    def imagine(
        self, observation: Any, action: Any, agent_id: Optional[int] = None
    ) -> Tuple[Any, float, bool]:
        """Predict ``(next_observation, reward, done)`` from (obs, action).

        If the configured backend fails and ``fallback_to_mock`` is True, the
        backend is switched to ``mock`` and the call is retried.

        ``agent_id`` is forwarded to stateful backends so multi-agent rollouts
        keep separate hidden states.
        """
        # Proactive availability check — switch to mock early with a clear log
        # message so we don't waste time on a subprocess that will fail.
        if self.backend == "cosmos3" and not self._cosmos_available():
            if self.fallback_to_mock:
                logger.warning(
                    "Cosmos3 backend unavailable (no NIM URL, cosmos_framework not "
                    "importable, and cosmos-framework root not found); falling back to mock."
                )
                self.backend = "mock"
            else:
                raise RuntimeError("Cosmos3 backend unavailable and fallback_to_mock=False")
        elif self.backend == "dreamdojo" and not self._dreamdojo_available():
            if self.fallback_to_mock:
                logger.warning(
                    "DreamDojo backend unavailable (cosmos_predict2 not importable "
                    "and DreamDojo root not found); falling back to mock."
                )
                self.backend = "mock"
            else:
                raise RuntimeError("DreamDojo backend unavailable and fallback_to_mock=False")

        try:
            if self.backend == "mock":
                return self._mock_imagine(observation, action, agent_id=agent_id)
            if self.backend == "cosmos3":
                return self._cosmos_imagine(observation, action, agent_id=agent_id)
            if self.backend == "dreamdojo":
                return self._dreamdojo_imagine(observation, action, agent_id=agent_id)
            raise ValueError(f"Unknown backend: {self.backend}")
        except Exception as exc:
            if self.fallback_to_mock and self.backend != "mock":
                warnings.warn(
                    f"WorldModel backend '{self.backend}' failed ({exc}); falling back to 'mock'.",
                    stacklevel=2,
                )
                logger.warning("Backend %s failed; falling back to mock", self.backend)
                self.backend = "mock"
                try:
                    return self._mock_imagine(observation, action, agent_id=agent_id)
                except Exception as exc2:
                    raise RuntimeError(f"Fallback to mock also failed: {exc2}") from exc
            raise

    # -----------------------------------------------------------------------
    # Mock backend
    # -----------------------------------------------------------------------
    def _ensure_mock(self, image: np.ndarray, action: np.ndarray) -> None:
        """Lazily create / recreate the tiny mock network when shapes change."""
        c, h, w = image.shape[2], image.shape[0], image.shape[1]
        action_dim = int(action.shape[0])

        if self._mock is not None:
            if getattr(self._mock, "image_shape", None) == (c, h, w) and getattr(self._mock, "action_dim", None) == action_dim:
                return

        if _HAS_TORCH:
            self._mock = TinyTorchWorldModel((c, h, w), action_dim, hidden=64, device=self.device)
        else:
            self._mock = TinyNumpyWorldModel((h, w, c), action_dim, feature_dim=self.feature_dim)

    def _mock_imagine(
        self, observation: Any, action: Any, agent_id: Optional[int] = None
    ) -> Tuple[Any, float, bool]:
        image, state, is_dict = _split_observation(observation)
        action_arr = _to_numpy(action).reshape(-1)
        self._ensure_mock(image, action_arr)
        next_image, reward, done = self._mock.imagine(image, action_arr, agent_id=agent_id)
        return _reconstruct_observation(next_image, state, is_dict, action_arr, self.state_integration), reward, done

    def _mock_train(
        self,
        observations: List[Any],
        actions: List[Any],
        rewards: List[float],
    ) -> None:
        n = min(len(observations) - 1, len(actions), len(rewards))
        if n <= 0:
            logger.warning("Not enough transitions to train the mock model (need obs, action, reward, next_obs)")
            return

        obs_list: List[np.ndarray] = []
        next_list: List[np.ndarray] = []
        act_list: List[np.ndarray] = []
        rew_list: List[float] = []
        done_list: List[float] = []

        for i in range(n):
            img, _, _ = _split_observation(observations[i])
            nxt, _, _ = _split_observation(observations[i + 1])
            act = _to_numpy(actions[i]).reshape(-1)
            rew = float(rewards[i])
            obs_list.append(img)
            next_list.append(nxt)
            act_list.append(act)
            rew_list.append(rew)
            done_list.append(0.0)

        img0, _, _ = _split_observation(observations[0])
        act0 = _to_numpy(actions[0]).reshape(-1)
        self._ensure_mock(img0, act0)
        self._mock.train(obs_list, act_list, next_list, rew_list, done_list)

    # -----------------------------------------------------------------------
    # Cosmos3 backend
    # -----------------------------------------------------------------------
    def _cosmos_available(self) -> bool:
        """Return True if any Cosmos3 execution path is available."""
        # NIM HTTP server
        if self.config.get("cosmos_nim_url"):
            return True
        # Direct import
        try:
            import cosmos_framework  # noqa: F401

            return True
        except Exception:
            pass
        # Subprocess via uv or the cosmos-framework root
        if self.cosmos_root.is_dir():
            return True
        return False

    def _build_cosmos_prompt(self, action: np.ndarray) -> str:
        action_desc = ", ".join(f"{a:.4f}" for a in action)
        template = self.config.get(
            "cosmos_prompt_template",
            "A robot camera view of the current scene. The next action is [{action}]. Predict the next video frame.",
        )
        return template.replace("{action}", action_desc)

    def _build_cosmos_input(
        self,
        image_path: Union[str, Path],
        prompt: str,
        output_dir: Union[str, Path],
        action_path: Optional[Union[str, Path]] = None,
    ) -> Dict[str, Any]:
        """Build the JSON input for ``cosmos_framework.scripts.inference``.

        When *action_path* is provided the ``forward_dynamics`` (action-conditioned)
        mode is used; otherwise ``image2video`` is used.
        """
        if action_path is not None:
            return {
                "model_mode": "forward_dynamics",
                "name": "world_model_step",
                "prompt": prompt,
                "vision_path": str(image_path),
                "action_path": str(action_path),
                "action_chunk_size": int(self.config.get("cosmos_action_chunk_size", 1)),
                "domain_name": str(self.config.get("cosmos_domain_name", "bridge_orig_lerobot")),
                "image_size": int(self.config.get("cosmos_image_size", 256)),
                "fps": int(self.config.get("cosmos_fps", 10)),
                "view_point": str(self.config.get("cosmos_view_point", "ego_view")),
                "seed": int(self.config.get("cosmos_seed", 0)),
            }
        return {
            "model_mode": self.config.get("cosmos_model_mode", "image2video"),
            "name": "world_model_step",
            "prompt": prompt,
            "vision_path": str(image_path),
            "resolution": self.config.get("cosmos_resolution", "256"),
            "aspect_ratio": self.config.get("cosmos_aspect_ratio", "1,1"),
            "num_frames": int(self.config.get("cosmos_num_frames", 24)),
            "fps": int(self.config.get("cosmos_fps", 10)),
            "seed": int(self.config.get("cosmos_seed", 0)),
        }

    def _build_cosmos_command(self, input_path: Union[str, Path], output_dir: Union[str, Path]) -> List[str]:
        checkpoint = str(self.config.get("cosmos_checkpoint", "Cosmos3-Nano"))
        seed = int(self.config.get("cosmos_seed", 0))

        if _HAS_UV:
            cmd = [
                "uv",
                "run",
                "--all-extras",
                "--group=cu130",
                "python",
                "-m",
                "cosmos_framework.scripts.inference",
            ]
        else:
            cmd = [sys.executable, "-m", "cosmos_framework.scripts.inference"]

        cmd.extend(
            [
                "-i",
                str(input_path),
                "-o",
                str(output_dir),
                "--checkpoint-path",
                checkpoint,
                "--seed",
                str(seed),
            ]
        )
        return cmd

    def _save_cosmos_action(self, action: np.ndarray, path: Path) -> None:
        """Write an action JSON file in the format expected by forward_dynamics.

        The Cosmos3 ``forward_dynamics`` mode reads a JSON file containing a
        list of action dictionaries.  Each dict has an ``action`` key with the
        raw action vector.  We emit a single-step action chunk.
        """
        chunk_size = int(self.config.get("cosmos_action_chunk_size", 1))
        action_list = []
        for i in range(chunk_size):
            entry = {"action": action.tolist(), "frame_idx": i}
            action_list.append(entry)
        path.write_text(json.dumps(action_list, indent=2))

    def _cosmos_nim_imagine(
        self, image: np.ndarray, action_arr: np.ndarray
    ) -> Optional[np.ndarray]:
        """Call a running Cosmos3-Generator NIM server via HTTP.

        Returns the predicted next frame as an RGB uint8 array, or ``None`` if
        the server is unreachable.
        """
        if not _HAS_REQUESTS:
            logger.debug("requests not available; skipping NIM path")
            return None
        nim_url = self.config.get("cosmos_nim_url")
        if not nim_url:
            return None

        import io

        prompt = self._build_cosmos_prompt(action_arr)
        buf = io.BytesIO()
        if _HAS_PIL:
            PILImage.fromarray(image).save(buf, format="PNG")
        elif _HAS_CV2:
            bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            ok, buf_bytes = cv2.imencode(".png", bgr)
            if not ok:
                return None
            buf = io.BytesIO(buf_bytes.tobytes())
        else:
            return None
        img_b64 = base64.b64encode(buf.getvalue()).decode("ascii")

        payload: Dict[str, Any] = {
            "prompt": prompt,
            "image": f"data:image/png;base64,{img_b64}",
            "seed": int(self.config.get("cosmos_seed", 0)),
        }
        try:
            logger.info("Calling Cosmos3 NIM at %s", nim_url)
            resp = requests.post(nim_url, json=payload, timeout=float(self.config.get("cosmos_timeout", 120.0)))
            resp.raise_for_status()
            data = resp.json()
        except Exception as exc:
            logger.warning("Cosmos3 NIM request failed: %s", exc)
            return None

        video_b64 = data.get("b64_video") or data.get("video")
        if not video_b64:
            logger.warning("Cosmos3 NIM returned no video data")
            return None
        video_bytes = base64.b64decode(video_b64)
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as vf:
            vf.write(video_bytes)
            vf_path = vf.name
        try:
            frames = _read_video_frames(vf_path)
        finally:
            os.unlink(vf_path)
        if not frames:
            logger.warning("Cosmos3 NIM video has no frames")
            return None
        return frames[min(1, len(frames) - 1)]

    def _cosmos_direct_imagine(
        self, image: np.ndarray, action_arr: np.ndarray
    ) -> Optional[np.ndarray]:
        """Try to run Cosmos3 inference in-process via direct import.

        Returns the predicted next frame or ``None`` if the import fails.
        """
        try:
            from cosmos_framework.inference.common.inference import Inference
            from cosmos_framework.inference.args import OmniSetupOverrides
        except Exception as exc:
            logger.debug("cosmos_framework not importable in-process: %s", exc)
            return None

        try:
            import torch as _t
            from PIL import Image as _PILImage

            with tempfile.TemporaryDirectory(prefix="cosmos3_direct_") as tmp:
                tmp_dir = Path(tmp)
                image_path = tmp_dir / "obs.png"
                _save_image(image, image_path)
                action_path = tmp_dir / "action.json"
                self._save_cosmos_action(action_arr, action_path)
                out_dir = tmp_dir / "outputs"
                out_dir.mkdir(parents=True, exist_ok=True)

                prompt = self._build_cosmos_prompt(action_arr)
                cfg = self._build_cosmos_input(image_path, prompt, out_dir, action_path=action_path)
                input_path = tmp_dir / "input.json"
                input_path.write_text(json.dumps(cfg, indent=2))

                checkpoint = str(self.config.get("cosmos_checkpoint", "Cosmos3-Nano"))
                setup = OmniSetupOverrides.model_construct(
                    checkpoint_path=checkpoint,
                    output_dir=out_dir,
                )
                pipe = setup.get_inference_cls().create(setup)
                sample_overrides = setup.get_sample_overrides_cls().from_files(
                    [input_path], overrides=setup.sample_overrides
                )
                for so in sample_overrides:
                    so.output_dir = out_dir / so.name
                    so.download(so.output_dir / "inputs")
                sample_args_list = [so.build_sample(model_config=pipe.model_config) for so in sample_overrides]
                pipe.generate(sample_args_list)

                sample_dir = out_dir / cfg["name"]
                return self._load_cosmos_output(sample_dir, image.shape[:2])
        except Exception as exc:
            logger.warning("Cosmos3 direct import inference failed: %s", exc)
            return None

    def _cosmos_imagine(
        self, observation: Any, action: Any, agent_id: Optional[int] = None
    ) -> Tuple[Any, float, bool]:
        # agent_id is accepted for API symmetry; Cosmos3 is stateless per call.
        image, state, is_dict = _split_observation(observation)
        action_arr = _to_numpy(action).reshape(-1)

        frame: Optional[np.ndarray] = None

        # --- Path 1: NIM HTTP API -----------------------------------------
        if self.config.get("cosmos_nim_url"):
            frame = self._cosmos_nim_imagine(image, action_arr)
            if frame is not None:
                logger.debug("Cosmos3: used NIM HTTP API")

        # --- Path 2: direct Python import ---------------------------------
        if frame is None and self.config.get("cosmos_use_direct", True):
            frame = self._cosmos_direct_imagine(image, action_arr)
            if frame is not None:
                logger.debug("Cosmos3: used direct in-process import")

        # --- Path 3: subprocess -------------------------------------------
        if frame is None:
            with tempfile.TemporaryDirectory(prefix="cosmos3_wm_") as tmp:
                tmp_dir = Path(tmp)
                image_path = tmp_dir / "obs.png"
                _save_image(image, image_path)

                # Write action JSON for forward_dynamics mode
                action_path = tmp_dir / "action.json"
                self._save_cosmos_action(action_arr, action_path)

                input_path = tmp_dir / "input.json"
                out_dir = tmp_dir / "outputs"
                out_dir.mkdir(parents=True, exist_ok=True)

                prompt = self._build_cosmos_prompt(action_arr)
                cfg = self._build_cosmos_input(image_path, prompt, out_dir, action_path=action_path)
                input_path.write_text(json.dumps(cfg, indent=2))

                cmd = self._build_cosmos_command(input_path, out_dir)
                timeout = float(self.config.get("cosmos_timeout", 120.0))
                _run_command(cmd, cwd=self.cosmos_root, timeout=timeout)

                sample_dir = out_dir / cfg["name"]
                frame = self._load_cosmos_output(sample_dir, image.shape[:2])

        return _reconstruct_observation(frame, state, is_dict, action_arr, self.state_integration), 0.0, False

    def _load_cosmos_output(self, sample_dir: Path, target_hw: Tuple[int, int]) -> np.ndarray:
        """Read the second generated frame (or the only image) from Cosmos output."""
        video_path = sample_dir / "vision.mp4"
        if video_path.exists():
            frames = _read_video_frames(video_path)
            if not frames:
                raise RuntimeError(f"Cosmos3 produced empty video: {video_path}")
            frame = frames[min(1, len(frames) - 1)]
        else:
            loaded = False
            for ext in (".png", ".jpg"):
                img_path = sample_dir / f"vision{ext}"
                if img_path.exists():
                    frame = _load_image(img_path)
                    loaded = True
                    break
            if not loaded:
                raise RuntimeError(f"No Cosmos3 output video/image found in {sample_dir}")

        if frame.shape[:2] != target_hw:
            frame = _resize_image(frame, (target_hw[1], target_hw[0]))
        return frame

    def _cosmos_train(
        self,
        observations: List[Any],
        actions: List[Any],
        rewards: List[float],
    ) -> None:
        """Log transition data for Cosmos3 fine-tuning.

        Writes (obs, action, next_obs, reward) tuples to disk in a format
        compatible with Cosmos3 SFT training.  If ``cosmos_train_toml`` is
        configured and the cosmos-framework is available, a fine-tuning job
        can be launched via ``_launch_cosmos_finetune``.
        """
        log_dir = self.output_dir / "cosmos3_train_log"
        log_dir.mkdir(parents=True, exist_ok=True)
        n = min(len(observations) - 1, len(actions), len(rewards))
        for i in range(n):
            img, state, _ = _split_observation(observations[i])
            nxt, _, _ = _split_observation(observations[i + 1])
            act = _to_numpy(actions[i]).reshape(-1)
            rew = float(rewards[i])

            tdir = log_dir / f"{self._log_counter:06d}"
            tdir.mkdir(parents=True, exist_ok=True)
            _save_image(img, tdir / "obs.png")
            _save_image(nxt, tdir / "next_obs.png")
            self._save_cosmos_action(act, tdir / "action.json")
            meta = {
                "action": act.tolist(),
                "reward": rew,
                "state": state.tolist() if state.size else [],
                "timestamp": time.time(),
            }
            (tdir / "meta.json").write_text(json.dumps(meta, default=float))
            self._log_counter += 1
        logger.info("Cosmos3 train: logged %d transitions to %s", n, log_dir)

        # Optionally launch fine-tuning if a TOML recipe is configured
        toml_recipe = self.config.get("cosmos_train_toml")
        if toml_recipe and self._cosmos_available():
            self._launch_cosmos_finetune(toml_recipe, log_dir)

    def _launch_cosmos_finetune(self, toml_recipe: str, data_dir: Path) -> None:
        """Launch a Cosmos3 SFT fine-tuning job (best-effort)."""
        try:
            if _HAS_UV:
                cmd = ["uv", "run", "python", "-m", "cosmos_framework.scripts.train", f"--sft-toml={toml_recipe}"]
            else:
                cmd = [sys.executable, "-m", "cosmos_framework.scripts.train", f"--sft-toml={toml_recipe}"]
            logger.info("Launching Cosmos3 fine-tuning: %s", " ".join(cmd))
            _run_command(cmd, cwd=self.cosmos_root, timeout=float(self.config.get("cosmos_train_timeout", 3600.0)))
        except Exception as exc:
            logger.warning("Cosmos3 fine-tuning launch failed: %s", exc)

    # -----------------------------------------------------------------------
    # DreamDojo backend
    # -----------------------------------------------------------------------
    def _dreamdojo_available(self) -> bool:
        """Return True if any DreamDojo execution path is available."""
        # Direct import
        try:
            import cosmos_predict2  # noqa: F401

            return True
        except Exception:
            pass
        # Subprocess via the DreamDojo root
        if self.dreamdojo_root.is_dir():
            ckpt = self.config.get("dreamdojo_ckpt", "checkpoints/iter_000006000")
            ckpt_path = self.dreamdojo_root / ckpt
            if ckpt_path.exists():
                return True
            # Also accept if the root exists — the subprocess will fail and we
            # fall back to mock, but we want to try.
            return True
        return False

    def _dreamdojo_direct_imagine(
        self, image: np.ndarray, action_arr: np.ndarray
    ) -> Optional[np.ndarray]:
        """Try to run DreamDojo inference in-process via direct import.

        Uses ``ActionStreamingInference`` from the interactive inference module.
        Returns the predicted next frame or ``None`` if the import fails.
        """
        try:
            from cosmos_predict2._src.predict2.interactive.inference.action_video2world import (
                ActionStreamingInference,
            )
        except Exception as exc:
            logger.debug("cosmos_predict2 not importable in-process: %s", exc)
            return None

        try:
            import torch as _t

            experiment = str(
                self.config.get("dreamdojo_experiment", "cosmos_predict2p5_2B_action_gr00t_gr1_self_forcing_no_s3")
            )
            ckpt = str(self.config.get("dreamdojo_ckpt", "checkpoints/iter_000006000"))
            config = str(
                self.config.get(
                    "dreamdojo_config",
                    "cosmos_predict2/_src/predict2/interactive/configs/config_distill.py",
                )
            )
            action_dim = int(self.config.get("dreamdojo_action_dim", 384))
            num_steps = int(self.config.get("dreamdojo_num_steps", 4))
            seed = int(self.config.get("dreamdojo_seed", 1))
            cr1_path = str(
                self.config.get("dreamdojo_cr1_embeddings", "datasets/cr1_dreamdojo_text_embeddings.pt")
            )

            streamer = ActionStreamingInference(
                config_path=config,
                experiment_name=experiment,
                ckpt_path=ckpt,
                s3_credential_path="credentials/s3_checkpoint.secret",
                cr1_embeddings_path=cr1_path,
            )

            # Prepare a single action
            action = np.zeros((1, action_dim), dtype=np.float32)
            action[0, : min(action_arr.shape[0], action_dim)] = action_arr[:action_dim]

            # Generate a single next frame
            with tempfile.TemporaryDirectory(prefix="dreamdojo_direct_") as tmp:
                tmp_dir = Path(tmp)
                image_path = tmp_dir / "obs.png"
                _save_image(image, image_path)
                out_video = tmp_dir / "out.mp4"

                frames = streamer.generate(
                    input_frame=str(image_path),
                    actions=action,
                    num_steps=num_steps,
                    seed=seed,
                    save_output=str(out_video),
                )
                if frames and len(frames) > 0:
                    frame = frames[-1] if isinstance(frames, (list, tuple)) else frames
                    if isinstance(frame, _t.Tensor):
                        frame = frame.detach().cpu().numpy()
                    if frame.ndim == 3 and frame.shape[-1] in (1, 3):
                        return frame.astype(np.uint8)
                    elif frame.ndim == 3 and frame.shape[0] in (1, 3):
                        return frame.transpose(1, 2, 0).astype(np.uint8)
                    return np.asarray(frame, dtype=np.uint8)

                # Fallback: read from saved video
                if out_video.exists():
                    vid_frames = _read_video_frames(out_video)
                    if vid_frames:
                        return vid_frames[-1]
            return None
        except Exception as exc:
            logger.warning("DreamDojo direct import inference failed: %s", exc)
            return None

    def _build_dreamdojo_command(
        self,
        image_path: Union[str, Path],
        action_path: Union[str, Path],
        out_video: Union[str, Path],
    ) -> List[str]:
        experiment = str(
            self.config.get("dreamdojo_experiment", "cosmos_predict2p5_2B_action_gr00t_gr1_self_forcing_no_s3")
        )
        ckpt = str(self.config.get("dreamdojo_ckpt", "checkpoints/iter_000006000"))
        config = str(
            self.config.get(
                "dreamdojo_config",
                "cosmos_predict2/_src/predict2/interactive/configs/config_distill.py",
            )
        )
        action_dim = int(self.config.get("dreamdojo_action_dim", 384))
        max_latent = int(self.config.get("dreamdojo_max_latent_frames", 1))
        fps = float(self.config.get("dreamdojo_fps", 10.0))
        num_steps = int(self.config.get("dreamdojo_num_steps", 4))
        seed = int(self.config.get("dreamdojo_seed", 1))
        cr1_path = str(
            self.config.get("dreamdojo_cr1_embeddings", "datasets/cr1_dreamdojo_text_embeddings.pt")
        )

        if _HAS_UV:
            cmd = [
                "uv",
                "run",
                "--extra",
                "cu128",
                "python",
                "-m",
                "cosmos_predict2._src.predict2.interactive.inference.action_video2world_teleop",
            ]
        else:
            cmd = [
                sys.executable,
                "-m",
                "cosmos_predict2._src.predict2.interactive.inference.action_video2world_teleop",
            ]

        cmd.extend(
            [
                f"--config={config}",
                f"--experiment={experiment}",
                f"--ckpt_path={ckpt}",
                f"--input_frame={image_path}",
                "--action_source=file",
                f"--action_file={action_path}",
                f"--fps={fps}",
                f"--num_steps={num_steps}",
                f"--max_latent_frames={max_latent}",
                "--no_display",
                f"--save_output={out_video}",
                f"--action_dim={action_dim}",
                f"--seed={seed}",
                f"--cr1_embeddings_path={cr1_path}",
            ]
        )
        return cmd

    def _dreamdojo_imagine(
        self, observation: Any, action: Any, agent_id: Optional[int] = None
    ) -> Tuple[Any, float, bool]:
        # agent_id is accepted for API symmetry; DreamDojo is stateless per call.
        image, state, is_dict = _split_observation(observation)
        action_arr = _to_numpy(action).reshape(-1)

        frame: Optional[np.ndarray] = None

        # --- Path 1: direct Python import ---------------------------------
        if self.config.get("dreamdojo_use_direct", True):
            frame = self._dreamdojo_direct_imagine(image, action_arr)
            if frame is not None:
                logger.debug("DreamDojo: used direct in-process import")

        # --- Path 2: subprocess -------------------------------------------
        if frame is None:
            with tempfile.TemporaryDirectory(prefix="dreamdojo_wm_") as tmp:
                tmp_dir = Path(tmp)
                image_path = tmp_dir / "obs.png"
                _save_image(image, image_path)

                # The DreamDojo teleop script reads actions from a .npy file.
                max_latent = int(self.config.get("dreamdojo_max_latent_frames", 1))
                action_file = tmp_dir / "action.npy"
                target_action_dim = int(self.config.get("dreamdojo_action_dim", 384))
                actions = np.zeros((max(max_latent, 1), target_action_dim), dtype=np.float32)
                actions[:, : min(action_arr.shape[0], target_action_dim)] = action_arr[:target_action_dim]
                np.save(action_file, actions)

                out_video = tmp_dir / "out.mp4"
                cmd = self._build_dreamdojo_command(image_path, action_file, out_video)
                timeout = float(self.config.get("dreamdojo_timeout", 120.0))
                _run_command(cmd, cwd=self.dreamdojo_root, timeout=timeout)

                if not out_video.exists():
                    raise RuntimeError(f"DreamDojo did not produce output video: {out_video}")
                frames = _read_video_frames(out_video)
                if not frames:
                    raise RuntimeError("DreamDojo output video is empty")
                frame = frames[-1]

        if frame.shape[:2] != image.shape[:2]:
            frame = _resize_image(frame, (image.shape[1], image.shape[0]))

        return _reconstruct_observation(frame, state, is_dict, action_arr, self.state_integration), 0.0, False

    def _dreamdojo_train(
        self,
        observations: List[Any],
        actions: List[Any],
        rewards: List[float],
    ) -> None:
        """Log transition data for DreamDojo fine-tuning.

        Writes (obs, action, next_obs, reward) tuples to disk in a format
        compatible with DreamDojo/Cosmos-Predict2.5 action-conditioned
        post-training.  If ``dreamdojo_train_config`` is configured and the
        framework is available, a fine-tuning job can be launched.
        """
        log_dir = self.output_dir / "dreamdojo_train_log"
        log_dir.mkdir(parents=True, exist_ok=True)
        n = min(len(observations) - 1, len(actions), len(rewards))
        for i in range(n):
            img, state, _ = _split_observation(observations[i])
            nxt, _, _ = _split_observation(observations[i + 1])
            act = _to_numpy(actions[i]).reshape(-1)
            rew = float(rewards[i])

            tdir = log_dir / f"{self._log_counter:06d}"
            tdir.mkdir(parents=True, exist_ok=True)
            _save_image(img, tdir / "obs.png")
            _save_image(nxt, tdir / "next_obs.png")
            np.save(tdir / "action.npy", act.astype(np.float32))
            meta = {
                "reward": rew,
                "state": state.tolist() if state.size else [],
                "timestamp": time.time(),
            }
            (tdir / "meta.json").write_text(json.dumps(meta, default=float))
            self._log_counter += 1
        logger.info("DreamDojo train: logged %d transitions to %s", n, log_dir)

        # Optionally launch fine-tuning if a config is configured
        train_config = self.config.get("dreamdojo_train_config")
        if train_config and self._dreamdojo_available():
            self._launch_dreamdojo_finetune(train_config, log_dir)

    def _launch_dreamdojo_finetune(self, train_config: str, data_dir: Path) -> None:
        """Launch a DreamDojo action-conditioned post-training job (best-effort)."""
        try:
            experiment = str(
                self.config.get("dreamdojo_train_experiment", "ac_reason_embeddings_rectified_flow_2b_256_320")
            )
            if _HAS_UV:
                cmd = [
                    "uv", "run", "python", "-m", "scripts.train",
                    f"--config={train_config}",
                    f"--experiment={experiment}",
                ]
            else:
                cmd = [
                    sys.executable, "-m", "scripts.train",
                    f"--config={train_config}",
                    f"--experiment={experiment}",
                ]
            logger.info("Launching DreamDojo fine-tuning: %s", " ".join(cmd))
            _run_command(cmd, cwd=self.dreamdojo_root, timeout=float(self.config.get("dreamdojo_train_timeout", 3600.0)))
        except Exception as exc:
            logger.warning("DreamDojo fine-tuning launch failed: %s", exc)


# ---------------------------------------------------------------------------
# Smoke test
# ---------------------------------------------------------------------------
def _smoke_test(backend: str = "mock") -> None:
    """Generate a random observation and call ``imagine()`` three times."""
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

    print(f"\n=== WorldModel smoke test (backend={backend}) ===")
    rng = np.random.default_rng(0)
    obs = rng.integers(0, 255, (64, 64, 3), dtype=np.uint8)
    action = rng.normal(0, 0.1, 4).astype(np.float32)

    wm = WorldModel(backend=backend, fallback_to_mock=True)

    # For mock, train on a tiny random trajectory so the learned heads are not
    # just zero-initialized; for cosmos3/dreamdojo this logs data before trying
    # the heavy inference command.
    obs_seq = [obs, obs, obs, obs]
    act_seq = [action, action, action]
    rew_seq = [-0.1, -0.1, -0.1]
    wm.train(obs_seq, act_seq, rew_seq)

    for step in range(3):
        next_obs, reward, done = wm.imagine(obs, action)
        print(f"step {step}: next_obs shape={next_obs.shape}, reward={reward:.4f}, done={done}")
        obs = next_obs

    print(f"Final backend: {wm.backend}")
    print("=== smoke test complete ===\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="World model smoke test")
    parser.add_argument(
        "--backend",
        choices=["mock", "cosmos3", "dreamdojo"],
        default="mock",
        help="Backend to exercise (cosmos3/dreamdojo will fall back to mock if unavailable)",
    )
    args = parser.parse_args()
    _smoke_test(args.backend)