react-manip 0.1.0__py3-none-any.whl

react_manip/core/metrics.py

@@ -0,0 +1,196 @@
+"""Step-level tracking for fixed-horizon, sustained-success evaluation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import math
+from typing import Any, Iterable, Mapping
+
+from .specs import SuccessRule
+
+
+def _failure_flags(failure: Any) -> tuple[bool, bool]:
+    if failure is None or failure is False:
+        return False, False
+    if isinstance(failure, Mapping):
+        return True, bool(failure.get("recoverable", False))
+    return True, False
+
+
+@dataclass
+class SustainedSuccessTracker:
+    """Collect continuous-state and discrete-event success signals."""
+
+    control_hz: float = 20.0
+    steps: int = 0
+    state_steps: int = 0
+    success_steps: int = 0
+    longest_streak_steps: int = 0
+    current_streak_steps: int = 0
+    completed_events: int = 0
+    failure_count: int = 0
+    recoverable_failure_count: int = 0
+    recovery_count: int = 0
+    time_to_first_failure_sec: float | None = None
+    recovery_latency_sec: list[float] = field(default_factory=list)
+    max_scalars: dict[str, float] = field(default_factory=dict)
+    sum_scalars: dict[str, float] = field(default_factory=dict)
+    count_scalars: dict[str, int] = field(default_factory=dict)
+    last_scalars: dict[str, float] = field(default_factory=dict)
+    _failure_started_at: float | None = field(default=None, init=False, repr=False)
+    _last_failure: bool = field(default=False, init=False, repr=False)
+    _last_info: Mapping[str, Any] = field(default_factory=dict, init=False, repr=False)
+
+    def update(
+        self,
+        success_state: bool | None,
+        *,
+        events: Iterable[Any] = (),
+        failure: Any = None,
+        dt: float | None = None,
+        info: Mapping[str, Any] | None = None,
+    ) -> None:
+        """Add one control tick to the trace.
+
+        ``success_state=None`` means the task did not expose a continuous
+        predicate for this tick. It is conservatively counted as a non-success
+        tick, preventing missing signals from inflating occupancy.
+        """
+
+        dt = float(dt if dt is not None else 1.0 / self.control_hz)
+        if dt <= 0:
+            raise ValueError("dt must be positive")
+        self.steps += 1
+        self.state_steps += 1
+        if success_state is True:
+            self.success_steps += 1
+            self.current_streak_steps += 1
+            self.longest_streak_steps = max(self.longest_streak_steps, self.current_streak_steps)
+        elif success_state is False:
+            self.current_streak_steps = 0
+        else:
+            # An unknown predicate cannot be part of a continuous success
+            # streak and is conservatively treated as a failed tick.
+            self.current_streak_steps = 0
+
+        event_list = list(events)
+        self.completed_events += len(event_list)
+        failed, recoverable = _failure_flags(failure)
+        now = (self.steps - 1) * dt
+        if failed and not self._last_failure:
+            self.failure_count += 1
+            if recoverable:
+                self.recoverable_failure_count += 1
+            self._failure_started_at = now
+            if self.time_to_first_failure_sec is None:
+                self.time_to_first_failure_sec = now
+        elif self._last_failure and self._failure_started_at is not None:
+            self.recovery_count += 1
+            self.recovery_latency_sec.append(max(0.0, now - self._failure_started_at))
+            self._failure_started_at = None
+        self._last_failure = failed
+        self._last_info = info or {}
+        # Optional, task-independent safety and real-time signals.  Tasks may
+        # emit any of these fields without changing the tracker schema.
+        for key in (
+            "peak_force_N",
+            "impulse_Ns",
+            "joint_limit_violation",
+            "inference_latency_ms",
+            "action_age_ms",
+            # Optional task scoring signals. Tasks may expose one or more of
+            # these without creating a task-specific runner branch.
+            "total_score",
+            "circle_hold_score",
+            "hole_drop_score",
+            "time_score",
+            "distance_score",
+            "attempt_score",
+            "distance_travelled_m",
+            "distance_fraction",
+            "lane_score",
+            "max_lane_deviation_m",
+            "speed_m_s",
+            "insertion_depth_m",
+            "max_insertion_depth_m",
+            "current_circle_hold_sec",
+            "max_circle_hold_sec",
+        ):
+            value = (info or {}).get(key)
+            if isinstance(value, (int, float)) and math.isfinite(float(value)):
+                scalar = float(value)
+                self.max_scalars[key] = max(self.max_scalars.get(key, float("-inf")), scalar)
+                self.sum_scalars[key] = self.sum_scalars.get(key, 0.0) + scalar
+                self.count_scalars[key] = self.count_scalars.get(key, 0) + 1
+                self.last_scalars[key] = scalar
+
+    def record_failure(self, *, recoverable: bool = False, at_sec: float | None = None) -> None:
+        """Record an out-of-band failure without adding a control tick."""
+
+        if not self._last_failure:
+            self.failure_count += 1
+            if recoverable:
+                self.recoverable_failure_count += 1
+            if self.time_to_first_failure_sec is None:
+                self.time_to_first_failure_sec = (
+                    float(at_sec) if at_sec is not None else self.steps / self.control_hz
+                )
+        self._last_failure = True
+
+    def summarize(self, rule: SuccessRule, *, terminal_success: bool = False) -> dict[str, Any]:
+        occupancy = self.success_steps / self.state_steps if self.state_steps else 0.0
+        longest = self.longest_streak_steps / self.control_hz
+        passed = (
+            occupancy >= rule.min_occupancy
+            and longest >= rule.min_longest_streak_sec
+            and self.completed_events >= rule.min_events
+            and (rule.max_failures is None or self.failure_count <= rule.max_failures)
+            and (not rule.require_terminal_success or terminal_success)
+        )
+        metrics = {
+            "steps": self.steps,
+            "state_steps": self.state_steps,
+            "sim_time_sec": self.steps / self.control_hz,
+            "success_steps": self.success_steps,
+            "success_occupancy": round(occupancy, 6),
+            "longest_success_streak_sec": round(longest, 6),
+            "completed_events": self.completed_events,
+            "event_rate_hz": round(self.completed_events / (self.steps / self.control_hz), 6)
+            if self.steps
+            else 0.0,
+            "failure_count": self.failure_count,
+            "recoverable_failure_count": self.recoverable_failure_count,
+            "recovery_count": self.recovery_count,
+            "time_to_first_failure_sec": self.time_to_first_failure_sec,
+            "mean_recovery_latency_sec": round(sum(self.recovery_latency_sec) / len(self.recovery_latency_sec), 6)
+            if self.recovery_latency_sec
+            else None,
+            "terminal_success": bool(terminal_success),
+            "passed": bool(passed),
+        }
+        for key, value in self.max_scalars.items():
+            metrics[f"max_{key}"] = round(value, 6)
+        if "max_peak_force_N" in metrics:
+            metrics["peak_force_N"] = metrics["max_peak_force_N"]
+        for key in (
+            "circle_hold_score", "hole_drop_score", "attempt_score",
+            "distance_score", "distance_travelled_m", "distance_fraction",
+            "max_lane_deviation_m", "speed_m_s", "insertion_depth_m",
+            "max_insertion_depth_m", "max_circle_hold_sec",
+        ):
+            max_key = f"max_{key}"
+            if max_key in metrics:
+                # Scores and a longest-duration measure are monotonic within
+                # a task episode, so their maximum is the final audit value.
+                metrics[key] = metrics[max_key]
+        for key in ("time_score",):
+            if key in self.last_scalars:
+                metrics[key] = round(self.last_scalars[key], 6)
+        # Lane and total scores may decrease after a vehicle leaves its lane.
+        for key in ("lane_score", "total_score", "attempt_score"):
+            if key in self.last_scalars:
+                metrics[key] = round(self.last_scalars[key], 6)
+        for key, value in self.sum_scalars.items():
+            count = self.count_scalars[key]
+            metrics[f"mean_{key}"] = round(value / count, 6) if count else None
+        return metrics

react_manip/core/observation.py

@@ -0,0 +1,59 @@
+"""Policy-observation filtering for privileged task state.
+
+Geometry truth is exposed through :class:`GeometryFeatureAdapter`, not through
+the observation passed to a normal policy.  Legacy robosuite tasks can still
+produce named object sensors, so the benchmark runner enforces the task's
+visibility contract at its final policy boundary.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Mapping
+
+
+def policy_observation(observation: Any, task_spec: Any, env: Any = None) -> Any:
+    """Return the non-privileged view that may be passed to a policy.
+
+    Tasks list any legacy raw keys in ``privileged_observation_keys``.  The
+    robosuite aggregate ``object-state`` is always privileged when geometry is
+    declared internal.  Non-mapping observations are returned unchanged so
+    native state-vector tasks keep their existing API.
+    """
+
+    if not isinstance(observation, Mapping):
+        return observation
+    contract = getattr(task_spec, "observation", None)
+    if contract is None or getattr(contract, "geometry_view", "internal") != "internal":
+        return observation
+    if "geometry" in tuple(getattr(contract, "default_policy_view", ())):
+        return observation
+
+    task_env = getattr(env, "env", env)
+    hidden = set(getattr(task_env, "privileged_observation_keys", ()))
+    hidden.add("object-state")
+    if not hidden.intersection(observation):
+        return observation
+
+    filtered = observation.copy()
+    for key in hidden:
+        filtered.pop(key, None)
+    return filtered
+
+
+def policy_info(info: Mapping[str, Any], policy: Any, env: Any) -> Mapping[str, Any]:
+    """Attach geometry truth only for an explicitly privileged oracle.
+
+    A regular policy sees the original public ``info`` mapping.  Oracle code
+    must opt in with ``requires_geometry_truth = True``; this makes privileged
+    use visible in code review and prevents accidental model leakage.
+    """
+
+    if not bool(getattr(policy, "requires_geometry_truth", False)):
+        raw = getattr(env, 'env', env)
+        hidden = set(getattr(raw, 'privileged_info_keys', ()))
+        return {k:v for k,v in info.items() if k not in hidden} if hidden else info
+    from .geometry import GeometryFeatureAdapter
+
+    enriched = dict(info)
+    enriched["geometry_ground_truth"] = GeometryFeatureAdapter.for_task(env).extract()
+    return enriched

react_manip/core/protocol.py

@@ -0,0 +1,91 @@
+"""Protocols shared by tasks, policies, and the fixed-horizon runner."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+from typing import Any, Mapping, Protocol, Sequence, runtime_checkable
+
+
+@dataclass(frozen=True)
+class StepResult:
+    """Normalized result of either Gymnasium (5-tuple) or robosuite (4-tuple)."""
+
+    observation: Any
+    reward: float
+    terminated: bool
+    truncated: bool
+    info: Mapping[str, Any]
+
+    @property
+    def done(self) -> bool:
+        return self.terminated or self.truncated
+
+
+@runtime_checkable
+class BenchmarkTask(Protocol):
+    """Minimum environment interface expected by :class:`BenchmarkRunner`."""
+
+    action_dim: int
+
+    def reset(self, *args: Any, **kwargs: Any) -> Any:
+        ...
+
+    def step(self, action: Any) -> Any:
+        ...
+
+    def close(self) -> None:
+        ...
+
+
+@runtime_checkable
+class Policy(Protocol):
+    """Policy interface; model-specific dependencies stay outside the runner."""
+
+    def reset(self, *, seed: int | None = None, task_spec: Any = None) -> None:
+        ...
+
+    def act(self, observation: Any, info: Mapping[str, Any] | None = None) -> Any:
+        ...
+
+
+def normalize_reset(value: Any) -> tuple[Any, Mapping[str, Any]]:
+    """Normalize old robosuite and Gymnasium reset return conventions."""
+
+    if isinstance(value, tuple) and len(value) == 2 and isinstance(value[1], Mapping):
+        return value[0], value[1]
+    return value, {}
+
+
+def normalize_step(value: Sequence[Any]) -> StepResult:
+    """Normalize a four- or five-element environment step result."""
+
+    if len(value) == 5:
+        obs, reward, terminated, truncated, info = value
+        reward = float(reward)
+        if not math.isfinite(reward):
+            raise ValueError(f"environment returned non-finite reward: {reward}")
+        return StepResult(obs, reward, bool(terminated), bool(truncated), info or {})
+    if len(value) == 4:
+        obs, reward, done, info = value
+        reward = float(reward)
+        if not math.isfinite(reward):
+            raise ValueError(f"environment returned non-finite reward: {reward}")
+        truncated = bool(done and (info or {}).get('TimeLimit.truncated', False))
+        return StepResult(obs, reward, bool(done) and not truncated, truncated, info or {})
+    raise ValueError("env.step(action) must return a 4- or 5-element tuple")
+
+
+class AdaptivePolicy(Policy, Protocol):
+    """Optional contract for task-selected practice/exam protocols.
+
+    begin_episode clears transient episode state and sets trial RNG, preserving
+    learned adaptation throughout practice AND exams. Snapshot hooks are only
+    required by the historical frozen protocol, not online evaluation.
+    end_episode receives filtered final feedback, including horizon timeouts.
+    """
+
+    def begin_episode(self, *, phase: str, episode_index: int, seed: int, task_spec: Any) -> None: ...
+    def end_episode(self, **feedback: Any) -> None: ...
+    def get_adaptation_state(self) -> Any: ...
+    def set_adaptation_state(self, state: Any) -> None: ...

react_manip/core/registry.py

@@ -0,0 +1,94 @@
+"""Task and robot registries used for additive extension."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from .specs import TaskSpec
+
+Factory = Callable[..., Any]
+
+
+@dataclass(frozen=True)
+class TaskDefinition:
+    spec: TaskSpec
+    factory: Factory
+
+
+class TaskRegistry:
+    def __init__(self) -> None:
+        self._items: dict[str, TaskDefinition] = {}
+
+    def register(self, spec: TaskSpec | dict[str, Any], factory: Factory, *, replace: bool = False) -> TaskDefinition:
+        spec = spec if isinstance(spec, TaskSpec) else TaskSpec.from_mapping(spec)
+        if spec.task_id in self._items and not replace:
+            raise KeyError(f"task already registered: {spec.task_id}")
+        definition = TaskDefinition(spec, factory)
+        self._items[spec.task_id] = definition
+        return definition
+
+    def get(self, task_id: str) -> TaskDefinition:
+        try:
+            return self._items[task_id]
+        except KeyError as exc:
+            known = ", ".join(sorted(self._items)) or "<empty>"
+            raise KeyError(f"unknown task {task_id!r}; registered tasks: {known}") from exc
+
+    def list(self) -> tuple[TaskDefinition, ...]:
+        return tuple(self._items[key] for key in sorted(self._items))
+
+
+class RobotRegistry:
+    def __init__(self) -> None:
+        self._items: dict[str, Any] = {}
+
+    def register(self, robot_id: str, adapter_or_factory: Any, *, replace: bool = False) -> Any:
+        robot_id = str(robot_id).lower()
+        if robot_id in self._items and not replace:
+            raise KeyError(f"robot already registered: {robot_id}")
+        self._items[robot_id] = adapter_or_factory
+        return adapter_or_factory
+
+    def get(self, robot_id: str) -> Any:
+        robot_id = str(robot_id).lower()
+        if robot_id in ("aloha", "aloha_mobile"):
+            from ..robots.catalog import normalize_robot_id
+            robot_id = normalize_robot_id(robot_id)
+        try:
+            return self._items[robot_id]
+        except KeyError as exc:
+            known = ", ".join(sorted(self._items)) or "<empty>"
+            raise KeyError(f"unknown robot {robot_id!r}; registered robots: {known}") from exc
+
+    def list(self) -> tuple[tuple[str, Any], ...]:
+        return tuple((key, self._items[key]) for key in sorted(self._items))
+
+
+_TASKS = TaskRegistry()
+_ROBOTS = RobotRegistry()
+
+
+def get_task_registry() -> TaskRegistry:
+    return _TASKS
+
+
+def get_robot_registry() -> RobotRegistry:
+    return _ROBOTS
+
+
+def register_task(spec: TaskSpec | dict[str, Any], factory: Factory | None = None, **kwargs: Any):
+    """Register directly or use as ``@register_task(spec)`` decorator."""
+
+    if factory is not None:
+        return _TASKS.register(spec, factory, **kwargs)
+
+    def decorator(fn: Factory) -> Factory:
+        _TASKS.register(spec, fn, **kwargs)
+        return fn
+
+    return decorator
+
+
+def register_robot(robot_id: str, adapter_or_factory: Any, **kwargs: Any) -> Any:
+    return _ROBOTS.register(robot_id, adapter_or_factory, **kwargs)