#!/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)