"""RosDiff's intelligence & routing layer: which model, which simulator, and what it will cost. Every generation goes through `plan_scene` / `plan_code`, which look at the request (and what retrieval found for it), rate its complexity, and choose the cheapest model tier that should handle it reliably. When a draft fails validation, `model_for_attempt` escalates to the next tier for the repair, so hard tasks end up on the strongest model without anyone picking it. Users never see a model picker; the plan's tier and reason are shown as information. There are two registries, one for scenes and one for code, of capable low-cost Chinese open-weight models on OpenRouter chosen by benchmark and price rather than by release date (see README for sources). Either can be replaced with SCENE_MODELS / CODE_MODELS (JSON) without code changes. Debugging: the first repair after a failed check goes to a free model through OpenCode Zen (DEBUG_MODEL, default opencode/deepseek-v4-flash-free), never through OpenRouter. If that call fails (rate limit, model withdrawn, OpenCode missing) the same attempt falls back to the next paid tier, so a free model can never cost a repair. Simulators: RosDiff scenes are MuJoCo (MJCF) scenes, so the routing chooses between what can run them today: MuJoCo (one run with video), MuJoCo Warp (many randomized copies on the GPU) and Isaac Sim (a photoreal RTX render of the same motion). Other simulators are listed as not available yet rather than pretended. """ from __future__ import annotations import json import os import re from dataclasses import asdict, dataclass, field TIERS = ("low", "mid", "high") @dataclass(frozen=True) class ModelSpec: id: str # OpenRouter model id tier: str # low | mid | high provider: str input_usd_per_m: float output_usd_per_m: float note: str = "" free: bool = False # served free through OpenCode Zen: recorded, never billed # OpenRouter list prices (USD per million tokens, Sept 2026) and the benchmark each was chosen on. CODE_REGISTRY = ( ModelSpec("deepseek/deepseek-v4-flash", "low", "DeepSeek", 0.089, 0.177, "79.0% SWE-bench Verified; very cheap"), ModelSpec("deepseek/deepseek-v4-pro", "mid", "DeepSeek", 0.435, 0.87, "80.6% SWE-bench Verified"), ModelSpec( "z-ai/glm-5.2", "high", "Zhipu AI (Z.ai)", 0.49, 1.54, "62.1% SWE-bench Pro, best open model on hard code" ), ) SCENE_REGISTRY = ( ModelSpec( "deepseek/deepseek-v4-flash", "low", "DeepSeek", 0.089, 0.177, "listed for completeness; scenes start at mid" ), ModelSpec("z-ai/glm-5.2", "mid", "Zhipu AI (Z.ai)", 0.49, 1.54, "within a point of Claude Opus 4.8 on FrontierSWE"), ModelSpec( "moonshotai/kimi-k2.6", "high", "Moonshot AI", 0.74, 3.49, "80.2% SWE-bench Verified; top open-weight AA index" ), ) DEFAULT_REGISTRY = CODE_REGISTRY # kept for callers that don't say which task def load_registry(task: str = "code") -> tuple[ModelSpec, ...]: """SCENE_MODELS / CODE_MODELS (or MODEL_REGISTRY for both) = '[{"id": ..., "tier": ..., ...}, ...]'.""" raw = os.environ.get("SCENE_MODELS" if task == "scene" else "CODE_MODELS", "").strip() raw = raw or os.environ.get("MODEL_REGISTRY", "").strip() if not raw: return SCENE_REGISTRY if task == "scene" else CODE_REGISTRY specs = tuple(ModelSpec(**m) for m in json.loads(raw)) if {s.tier for s in specs} != set(TIERS): raise ValueError("a model registry needs one model for each tier: low, mid, high") return specs def all_models() -> dict[str, ModelSpec]: return {m.id: m for task in ("code", "scene") for m in load_registry(task)} def model_for_tier(tier: str, registry: tuple[ModelSpec, ...] | None = None, task: str = "code") -> ModelSpec: registry = registry or load_registry(task) return next(m for m in registry if m.tier == tier) def free_debugger(settings) -> ModelSpec | None: if not getattr(settings, "free_debug", False) or not settings.debug_model: return None return ModelSpec(settings.debug_model, "free", "OpenCode Zen", 0.0, 0.0, "free debugging model", free=True) def is_free(model_id: str | None) -> bool: return bool(model_id) and model_id.startswith("opencode/") def escalate(tier: str, steps: int = 1) -> str: return TIERS[min(TIERS.index(tier) + steps, len(TIERS) - 1)] # --------------------------------------------------------------------------- complexity _CODE_COMPLEX = ( "nav2", "navigation stack", "moveit", "slam", "costmap", "planner", "behavior tree", "behaviour tree", "state machine", "lifecycle", "ros2_control", "hardware interface", "controller manager", "action server", "perception", "point cloud", "pointcloud", "kalman", "ekf", "localization", "localisation", "multi-robot", "fleet", "autonomous", "inverse kinematics", "trajectory", "mpc", "reinforcement", ) # fmt: skip _CODE_MEDIUM = ( "service", "action", "parameter", "launch", "timer", "tf", "transform", "filter", "odometry", "odom", "lidar", "laser", "imu", "camera", "image", "twist", "subscriber", "publisher", "client", "path", "c++", ) # fmt: skip _SCENE_COMPLEX = ( "maze", "warehouse", "rough terrain", "heightfield", "uneven", "hinge", "door that opens", "drawer", "joint", "articulated", "pendulum", "tendon", "cloth", "rope", "multiple robots", "two robots", "conveyor", "stairs with", "spiral", "ramp and stairs", ) # fmt: skip _SCENE_MEDIUM = ( "stairs", "steps", "ramp", "slope", "table", "shelf", "wall", "doorway", "obstacle", "boxes", "cubes", "ball", "friction", "terrain", "tray", ) # fmt: skip def _count(text: str, words: tuple[str, ...]) -> list[str]: t = text.lower() return [w for w in words if w in t] @dataclass class Plan: task: str # scene | code complexity: str # simple | medium | complex tier: str model: str reason: str category: str # simulation_ai | coding (repairs are booked as debugging) estimate_inr: tuple[float, float] = (0.0, 0.0) simulator: str | None = None simulator_reason: str | None = None signals: list[str] = field(default_factory=list) debugger: str | None = None # the free model tried first on the first repair def to_dict(self) -> dict: d = asdict(self) d["estimate_inr"] = list(self.estimate_inr) return d def _level(score: int) -> str: return "complex" if score >= 4 else "medium" if score >= 2 else "simple" def _tier_for(level: str) -> str: return {"simple": "low", "medium": "mid", "complex": "high"}[level] def plan_scene(description: str, robot: str | None, settings, n_numbers: int | None = None) -> Plan: complex_hits = _count(description, _SCENE_COMPLEX) medium_hits = _count(description, _SCENE_MEDIUM) objects = n_numbers if n_numbers is not None else len(re.findall(r"\d+(?:\.\d+)?", description)) # everyday parts add up to 'medium' at most; only genuinely hard ones push a task to the top tier score = 3 * len(complex_hits) + min(len(medium_hits), 3) + (objects >= 6) + (len(description) > 400) level = _level(score) # Scenes must be spatially exact; the cheapest tier is only used for very small ones. tier = max(_tier_for(level), "mid", key=TIERS.index) signals = complex_hits + medium_hits reason = { "simple": "a small scene", "medium": "a scene with several objects to place exactly", "complex": "a scene with articulated or intricate parts", }[level] sim, sim_reason = choose_simulator(description) plan = Plan("scene", level, tier, model_for_tier(tier, task="scene").id, reason, "simulation_ai", signals=signals) plan.simulator, plan.simulator_reason = sim, sim_reason plan.debugger = (d := free_debugger(settings)) and d.id plan.estimate_inr = estimate_inr(plan, settings) return plan def plan_code(request: str, distro: str, language: str, settings, n_interfaces: int = 0) -> Plan: complex_hits = _count(request, _CODE_COMPLEX) medium_hits = _count(request, _CODE_MEDIUM) score = 3 * len(complex_hits) + min(len(medium_hits), 3) + (1 if language == "cpp" else 0) score += 1 if n_interfaces >= 3 else 0 score += 1 if len(request) > 500 else 0 level = _level(score) tier = _tier_for(level) reason = { "simple": "a small, self-contained ROS 2 node", "medium": "a ROS 2 package with several parts to connect", "complex": "a ROS 2 system with navigation, control or planning", }[level] if language == "cpp" and level != "complex": reason += " in C++" plan = Plan("code", level, tier, model_for_tier(tier).id, reason, "coding", signals=complex_hits + medium_hits) plan.debugger = (d := free_debugger(settings)) and d.id plan.estimate_inr = estimate_inr(plan, settings) return plan def candidates_for_attempt(plan: Plan, attempt: int) -> list[ModelSpec]: """Models to try for one attempt, in order (the next is only used if a call fails outright). Attempt 0: the planned tier. Attempt 1: the free debugger, falling back to one tier up. Later repairs keep going up one tier each. Without a free debugger every repair simply goes one tier up.""" registry = load_registry(plan.task) if attempt == 0: return [model_for_tier(plan.tier, registry)] if plan.debugger: paid = model_for_tier(escalate(plan.tier, attempt), registry) if attempt == 1: return [ModelSpec(plan.debugger, "free", "OpenCode Zen", 0.0, 0.0, free=True), paid] return [paid] return [model_for_tier(escalate(plan.tier, attempt), registry)] def model_for_attempt(plan: Plan, attempt: int) -> ModelSpec: """The paid model an attempt ends on in the worst case.""" return candidates_for_attempt(plan, attempt)[-1] # --------------------------------------------------------------------------- simulators SIMULATORS = { "mujoco": {"label": "MuJoCo", "what": "one run with video and 3D replay", "available": True}, "mujoco_warp": {"label": "MuJoCo Warp", "what": "64 randomized copies at once on the GPU", "available": True}, "isaac_sim": {"label": "Isaac Sim", "what": "photoreal RTX render of the same motion", "available": True}, "gazebo": {"label": "Gazebo", "what": "not available in RosDiff yet", "available": False}, "webots": {"label": "Webots", "what": "not available in RosDiff yet", "available": False}, "pybullet": {"label": "PyBullet", "what": "not available in RosDiff yet", "available": False}, } _WANTS_RENDER = ("photoreal", "realistic render", "camera image", "rgb", "lighting", "vision", "synthetic image", "rendered image", "ray tracing", "ray-traced") # fmt: skip _WANTS_MANY = ("robust", "random", "variation", "many times", "stress", "domain randomization", "monte carlo", "reinforcement", "train a policy", "statistics", "how often") # fmt: skip def choose_simulator(text: str) -> tuple[str, str]: if hits := _count(text, _WANTS_MANY): return "mujoco_warp", f"you asked about {hits[0]}: many randomized copies show how often it works" if hits := _count(text, _WANTS_RENDER): return "isaac_sim", f"you mentioned {hits[0]}: MuJoCo simulates, Isaac Sim renders it photoreal" return "mujoco", "contact-rich rigid-body physics with the real robot model, fast and precise" # --------------------------------------------------------------------------- cost # Typical tokens per model call, measured from RosDiff's prompts (context + robot facts / interface definitions). TOKENS = { ("scene", "direct"): (16_000, 3_500), ("code", "direct"): (20_000, 6_000), ("code", "opencode"): (60_000, 9_000), # OpenCode's agent loop re-reads context across several steps } def call_cost_usd(model: ModelSpec, input_tokens: int, output_tokens: int) -> float: return (input_tokens * model.input_usd_per_m + output_tokens * model.output_usd_per_m) / 1e6 def estimate_inr(plan: Plan, settings) -> tuple[float, float]: """(likely, worst case): one call on the planned tier vs. every attempt, escalating.""" mode = "opencode" if plan.task == "code" and settings.codegen_backend == "opencode" else "direct" tin, tout = TOKENS[(plan.task, mode)] attempts = settings.mjcf_max_attempts if plan.task == "scene" else settings.code_max_attempts first = call_cost_usd(model_for_attempt(plan, 0), tin, tout) # worst case: every attempt runs, and the free debugger fails so its paid fallback is billed worst = sum(call_cost_usd(model_for_attempt(plan, i), tin * (1 + i // 2), tout) for i in range(attempts)) return round(first * settings.usd_inr, 2), round(worst * settings.usd_inr, 2)