Spaces:
Paused
Paused
| """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") | |
| 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] | |
| 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) | |