arena/environments/tera_spatial/tera_crow_nav/taskset.py

"""tera-crow-nav: fly a crow to a waypoint on sixteen replies.

Tera's `crow-nav-v1` is a three-dimensional waypoint problem over the same flight
controller the renderer flies — momentum, a bank-coupled turn rate, a finite
flight-energy budget, thermals, drag that rises with the square of airspeed, and
an envelope at 2 m and 40 m that ends the flight if you touch it. The simulator
runs at 10 Hz for up to 300 steps. A language model does not fly at 10 Hz.

So a reply is a decision, not a frame: one control vector plus the number of
steps to hold it for. Sixteen replies, twelve steps each at most.

**Those two integers are the environment**, and they are measured. `measure_hold.py`
flies Tera's own scripted baseline at every refresh rate over sixteen scenario x
seed episodes:

    hold    return    goal          turns   best hold   return    goal
       1     5.484    100%              1          12     0.031      0%
       5     5.479    100%              4          12     1.039      0%
       8     4.906     94%              8          11     3.420     44%
      12     4.401     88%             12           7     4.974     88%
      24     1.392     44%             14           6     5.409    100%
      32    -2.807      0%             16           5     5.479    100%
     300    -3.764      0%             20           4     5.483    100%

Read the left table for house rule 4's premise and the right one for its
conclusion. Open-loop fails outright — one decision for the whole flight returns
-3.764 against an inaction floor of -4.161, which is to say a bird that sets a
course and leaves does no better than a bird that does nothing. And the budget
BINDS all the way to sixteen: at eight turns the best possible constant hold
reaches the waypoint 44% of the time, at twelve 88%, and only at fourteen does it
reach every one. Sixteen is the first budget at which the ceiling is reachable and
fifteen turns of slack is not on offer.

Two rewards, both products:

  flight    the episode's return against nine tenths of what Tera's scripted
            controller returned on the same scenario at the same seed. Smooth:
            reaching the waypoint is already the biggest term inside it — the
            simulator pays +3 of a ~5.5 return for arriving — so a near-miss
            scores about four tenths and a wander scores zero.
  economy   that same quality, multiplied by having arrived and by how directly.
            Goal attainment MULTIPLIES and is never a term of its own; standing
            beside the sum it would double-count a bonus the return already
            contains. Efficiency multiplies too, because a flight that never
            arrived cannot be quick about arriving.

Minimum-power inaction returns -4.161 and scores exactly 0.000 on both. The
scripted controller flown inside this budget scores 1.000 on both. `gate`,
`quality`, `step_ratio` and `replay_ok` are recorded and never rewarded.
"""

from __future__ import annotations

import math
from typing import Any, ClassVar

import verifiers.v1 as vf

from tera_spatial.spatial import (
    SpatialConfig,
    SpatialData,
    SpatialTask,
    TeraSpatialTaskset,
    TurnView,
)

MISSION = """You are flying a crow to a waypoint.

The simulator advances in fixed steps of {step:g} s. You do not fly it step by step:
each reply is ONE set of control positions plus a `hold`, and the crow holds those
positions for that many steps before you are shown where it ended up.

  You get {turns} replies. A reply may be held for 1 to {max_hold} steps.
  That is {budget} steps of flight in total if you use every one of them.

The controls, each a number from -1 to 1:

  forward   airspeed demand. -1 is minimum power (about 4 m/s), +1 is full
            (about 16 m/s). Flapping hard drains flight energy, and a tired crow
            has less speed and climb authority than a rested one.
  turn      rate of turn, up to about 1.75 rad/s at full deflection and full
            speed. The crow banks into it, which adds a little more.
  pitch     nose attitude, up to about 0.72 rad. Pitch trades airspeed for
            height and back; it is not the climb control.
  climb     powered vertical demand, up to about 5 m/s at full deflection.
  glide     true or false. Gliding recovers flight energy and holds about
            7.5 m/s, but you give up powered climb while you do it.

The world:

  The waypoint counts as reached within {radius:g} m in three dimensions, and pays a
  large bonus when you get there.
  You must stay between {floor:g} m and {ceiling:g} m altitude and inside +/-{bound:g} m in x and z.
  Touching any of those ends the flight immediately and costs more than the
  waypoint was worth.
  Every step costs a little time, and power, turn, pitch and climb demand each
  cost a little more. Pointing away from the waypoint costs a little.

Reply with ONE ```json code block and nothing else that matters:

```json
{{"forward": 0.5, "turn": 0.0, "pitch": 0.0, "climb": 0.0, "glide": false, "hold": 6}}
```

Any control you leave out keeps the value it already had. A reply with no readable
JSON object holds your last controls for {default_hold} steps and costs you the turn — there
is no way to pause."""


def _degrees(radians: float) -> float:
    return math.degrees(radians)


def _wrap(value: float) -> float:
    return (value + math.pi) % (2 * math.pi) - math.pi


class CrowNavConfig(SpatialConfig):
    """Crow-nav's knobs are the shared ones. The defaults are the measured ones and
    a config that changes `max_turns` or `max_hold` is changing the environment, not
    tuning it — see the two tables in this module's docstring."""


class CrowNavTaskset(TeraSpatialTaskset, vf.Taskset[SpatialTask, CrowNavConfig]):
    """Sixteen waypoint flights over the four public crow scenarios.

    The generic form is not decoration: verifiers v1 resolves a taskset's config
    through `__orig_bases__`, so `vf.Taskset[SpatialTask, CrowNavConfig]` in the
    bases is the whole of the wiring. There is no `CONFIG` classvar to set, and a
    taskset that omits this reads every per-environment setting off the base config
    and silently ignores the TOML.
    """

    ENV_ID: ClassVar[str] = "crow-nav-v1"

    def mission(self, manifest: dict[str, Any], data: SpatialData) -> str:
        return MISSION.format(
            step=manifest["fixedStepSeconds"],
            turns=data.max_turns,
            max_hold=data.max_hold,
            default_hold=data.default_hold,
            budget=data.max_turns * data.max_hold,
            radius=2.5,
            floor=2,
            ceiling=40,
            bound=120,
        )

    def board(self, view: TurnView) -> str:
        """The state vector, rendered as a panel a reader could fly on.

        Every number here is the environment's own published observation or plain
        arithmetic over it: `distanceToGoalM` and the deltas are fields on the
        manifest, and the bearing is `atan2` of two of them against a third. That
        trigonometry is given away deliberately. What crow-nav measures is whether a
        model can fly a body with momentum to a point inside a turn budget, not
        whether it can do `atan2` in its head with thinking switched off — and the
        scripted controller the reward is normalised against reads exactly these
        fields, so nothing here moves the denominator.

        What it does NOT contain is anything the simulator has not already
        published: this is a dict that crossed a JSON pipe, with no handle in it to
        walk back to the model.
        """
        o = view.observation
        bearing = _degrees(_wrap(math.atan2(-o["deltaX"], -o["deltaZ"]) - o["yaw"]))
        contact = o["altitudeBoundContact"]
        seconds = view.steps * 0.1
        lines = [
            f"turn {view.turn} of {view.turn + view.turns_left - 1}"
            f"  ·  step {view.steps} of {view.budget}"
            f"  ·  {seconds:.1f} s flown",
            "",
            f"position    x {o['x']:>8.1f}   y {o['y']:>6.1f}   z {o['z']:>8.1f}"
            f"      altitude {o['y']:.1f} m of 2–40",
            f"attitude    yaw {_degrees(o['yaw']):>6.0f}°   pitch {_degrees(o['pitch']):>4.0f}°"
            f"        speed {o['speedMps']:.1f} m/s   vertical {o['verticalSpeedMps']:+.1f} m/s",
            f"waypoint    x {o['goalX']:>8.1f}   y {o['goalY']:>6.1f}   z {o['goalZ']:>8.1f}",
            f"            {o['distanceToGoalM']:.1f} m away  ·  {o['deltaY']:+.1f} m of height to make up"
            f"  ·  bearing {bearing:+.0f}° (a positive turn closes it)",
            f"envelope    {'clear' if contact == 'none' else contact + ' altitude bound — you are on it'}",
        ]
        if view.last is None:
            lines.append("controls    at rest; you have not flown yet")
        else:
            action = view.last.action
            lines.append(
                "you flew    "
                + "  ".join(
                    f"{name} {_show(action.get(name, 0))}"
                    for name in ("forward", "turn", "pitch", "climb", "glide")
                )
                + f"   held {view.last.hold} step{'s' if view.last.hold != 1 else ''}"
            )
            if view.last.problem:
                lines.append(f"⚠  {view.last.problem}, so those were held again")
        lines.append("")
        lines.append(
            f"{view.turns_left} repl{'y' if view.turns_left == 1 else 'ies'} left, "
            f"{view.budget - view.steps} steps of flight. Your controls?"
        )
        return "\n".join(lines)


def _show(value: Any) -> str:
    if isinstance(value, bool):
        return "yes" if value else "no"
    return f"{float(value):+.2f}"


__all__ = ["CrowNavConfig", "CrowNavTaskset"]