experiments/shared/src/methods.py

"""Method registry for paper-facing experiments."""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class MethodSpec:
    name: str
    display_name: str
    category: str
    kind: str
    role: str
    notes: str


METHODS: dict[str, MethodSpec] = {
    "flowmo": MethodSpec(
        name="flowmo",
        display_name="FlowMo",
        category="A_learned_world_model",
        kind="learned_world_model",
        role="proposed",
        notes="Proposed flow-momentum world model with short state history, long drift context, and zero-context residual transition.",
    ),
    "leworldmodel": MethodSpec(
        name="leworldmodel",
        display_name="LeWorldModel",
        category="A_learned_world_model",
        kind="learned_world_model",
        role="external_style_baseline",
        notes="JEPA-like latent predictor under the shared clean-image rollout protocol.",
    ),
    "planet": MethodSpec(
        name="planet",
        display_name="PlaNet RSSM",
        category="A_learned_world_model",
        kind="learned_world_model",
        role="external_style_baseline",
        notes="Recurrent state-space dynamics world model under the shared clean-image rollout protocol.",
    ),
    "tdmpc2": MethodSpec(
        name="tdmpc2",
        display_name="TD-MPC2 Dynamics",
        category="A_learned_world_model",
        kind="learned_world_model",
        role="external_style_baseline",
        notes="Latent dynamics world model under the shared clean-image rollout protocol.",
    ),
    "pid_los_controller": MethodSpec(
        name="pid_los_controller",
        display_name="PID/LOS controller",
        category="B_traditional_controller",
        kind="traditional_controller",
        role="planning_baseline",
        notes="Classical image-pose-based line-of-sight controller for downstream planning comparison.",
    ),
    "physics_mpc_no_flow": MethodSpec(
        name="physics_mpc_no_flow",
        display_name="Physics MPC No-Flow",
        category="B_traditional_controller",
        kind="traditional_controller",
        role="planning_baseline",
        notes="Classical controller without external-flow compensation.",
    ),
    "current_estimator_mpc": MethodSpec(
        name="current_estimator_mpc",
        display_name="Current-Estimator MPC",
        category="B_traditional_controller",
        kind="traditional_controller",
        role="planning_baseline",
        notes="Classical controller using recent image-pose drift estimate.",
    ),
    "oracle_flow_mpc": MethodSpec(
        name="oracle_flow_mpc",
        display_name="Oracle-Flow MPC",
        category="B_traditional_controller",
        kind="traditional_controller",
        role="oracle_planning_upper_bound",
        notes="Oracle controller with true local flow; used only as a planning upper/reference bound.",
    ),
}


PAPER_LEARNED_METHODS = [name for name, spec in METHODS.items() if spec.category == "A_learned_world_model"]
TRADITIONAL_METHODS = [name for name, spec in METHODS.items() if spec.kind == "traditional_controller"]


def display_name(method: str) -> str:
    return METHODS[method].display_name