arena/environments/tera_spatial/tests/test_replay.py

"""The bridge's correctness gate.

The whole claim this package makes is: *a trace produced by driving Tera from
Python replays, inside Tera, to a bit-identical FNV-1a-64 checksum — and a trace
that has been edited by so much as one float does not.* If that holds, a reward
that reaches Python is a reward that TypeScript actually produced, and Arena can
grade a spatial rollout without a second implementation of the simulator to
disagree with the first.

So this file is not a smoke test. It exercises **every** public scenario at two
seeds — 4 environments x 4 scenarios x 2 seeds = 32 episodes — and for each one:

1. drives the episode a step at a time from Python, using actions that made the
   round trip through `json`, and checks every per-step checksum against the
   trace Tera produced under its own baseline;
2. replays the Python-driven trace in a *fresh* environment;
3. replays it again after a `json.dumps`/`json.loads` cycle, because the wire
   format is the thing under test and a float that does not survive Python's
   repr is a silent divergence;
4. tampers with it seven ways and requires every one to be rejected.

Run:

    uv run python -m unittest discover -s environments/tera_spatial/tests -v
"""

from __future__ import annotations

import copy
import json
import time
import unittest

from tera_spatial import ENVIRONMENT_IDS, TeraError, TeraWorker
from tera_spatial import hashes
from tera_spatial.closure import bare_specifiers

SEEDS = (115, 2718)


def scenario_seed_matrix(worker: TeraWorker) -> list[tuple[str, str, str, int]]:
    """(env id, scenario id, split, seed) for every public scenario at every seed."""
    matrix: list[tuple[str, str, str, int]] = []
    for manifest in worker.manifests():
        for split in ("train", "dev"):
            for scenario_id in manifest["scenarioIds"][split]:
                for seed in SEEDS:
                    matrix.append((manifest["id"], scenario_id, split, seed))
    return matrix


class VendoredSourceTests(unittest.TestCase):
    """The checksums in a trace are only worth something if the code is the pinned code."""

    def test_vendored_tree_matches_recorded_hashes(self) -> None:
        self.assertEqual(hashes.verify(), [])

    def test_closure_is_the_expected_shape(self) -> None:
        # 23 files and 283 KB is the measured closure of `src/arena/index.ts`. A
        # change here is not necessarily wrong, but it is never incidental.
        self.assertEqual(len(hashes.VENDORED_FILES), 23)
        self.assertEqual(hashes.CLOSURE_BYTES, 282_684)

    def test_closure_has_no_bare_specifiers(self) -> None:
        """No `node_modules`: the wheel must run the simulator with nothing installed."""
        self.assertEqual(bare_specifiers(), [])


class ReplayGateTests(unittest.TestCase):
    """Python drives, TypeScript simulates, and the two agree to the last bit."""

    worker: TeraWorker

    @classmethod
    def setUpClass(cls) -> None:
        started = time.perf_counter()
        cls.worker = TeraWorker()
        cls.startup_seconds = time.perf_counter() - started

    @classmethod
    def tearDownClass(cls) -> None:
        cls.worker.close()

    def test_worker_exposes_the_four_environments(self) -> None:
        ids = [manifest["id"] for manifest in self.worker.manifests()]
        self.assertEqual(sorted(ids), sorted(ENVIRONMENT_IDS))

    def test_every_scenario_and_seed_round_trips_and_replays(self) -> None:
        matrix = scenario_seed_matrix(self.worker)
        self.assertEqual(len(matrix), 32, "4 environments x 4 scenarios x 2 seeds")
        for env_id, scenario_id, split, seed in matrix:
            with self.subTest(env=env_id, scenario=scenario_id, seed=seed):
                self._assert_round_trip(env_id, scenario_id, split, seed)

    def _assert_round_trip(self, env_id: str, scenario_id: str, split: str, seed: int) -> None:
        # The reference episode runs entirely inside TypeScript, under the
        # baseline `src/arena/index.ts` exports. Its trace is the ground truth.
        reference = self.worker.oracle(env_id, seed, scenario_id, policy="scripted")
        oracle_trace = reference["trace"]
        self.assertGreater(len(oracle_trace["steps"]), 0)

        # Now drive the same episode from Python, one step at a time, with every
        # action having crossed the pipe as JSON. This is the direction that
        # matters: it is how a rollout will actually run.
        with self.worker.open(env_id, seed, scenario_id) as episode:
            self.assertEqual(
                episode.info["stateChecksum"], oracle_trace["initialStateChecksum"]
            )
            for frame in oracle_trace["steps"]:
                result = episode.step(json.loads(json.dumps(frame["action"])))
                self.assertEqual(result["info"]["stateChecksum"], frame["stateChecksum"])
                self.assertEqual(result["reward"], frame["reward"])
                self.assertEqual(result["terminated"], frame["terminated"])
                self.assertEqual(result["truncated"], frame["truncated"])
            driven = episode.trace()

        # Python drove it; the envelope checksum says TypeScript agrees it is the
        # same episode.
        self.assertEqual(driven["checksum"], oracle_trace["checksum"])
        self.assertEqual(driven["finalStateChecksum"], oracle_trace["finalStateChecksum"])
        self.assertEqual(driven["cumulativeReward"], oracle_trace["cumulativeReward"])

        # And the wire format survives Python's float repr. `json` is the only
        # thing between the two runtimes, so it is part of the claim.
        wired = json.loads(json.dumps(driven))
        self.assertEqual(wired["checksum"], oracle_trace["checksum"])

        replayed = self.worker.replay(env_id, wired)
        self.assertEqual(replayed["finalStateChecksum"], oracle_trace["finalStateChecksum"])
        self.assertEqual(replayed["cumulativeReward"], oracle_trace["cumulativeReward"])
        self.assertEqual(replayed["steps"], len(oracle_trace["steps"]))

    def test_snapshot_restore_continues_an_episode_exactly(self) -> None:
        for env_id in ENVIRONMENT_IDS:
            with self.subTest(env=env_id):
                reference = self.worker.oracle(env_id, 115, {"split": "train"})
                frames = reference["trace"]["steps"]
                cut = max(1, len(frames) // 2)
                with self.worker.open(env_id, 115, {"split": "train"}) as episode:
                    for frame in frames[:cut]:
                        episode.step(frame["action"])
                    snapshot = json.loads(json.dumps(episode.snapshot()))
                with self.worker.open(env_id, 115, {"split": "train"}) as resumed:
                    restored = resumed.restore(snapshot)
                    self.assertEqual(
                        restored["info"]["stateChecksum"], frames[cut - 1]["stateChecksum"]
                    )
                    for frame in frames[cut:]:
                        result = resumed.step(frame["action"])
                        self.assertEqual(
                            result["info"]["stateChecksum"], frame["stateChecksum"]
                        )

    def test_a_broken_seal_is_rejected(self) -> None:
        """Seven edits, none of them re-sealed. The envelope checksum catches all seven."""
        for env_id in ENVIRONMENT_IDS:
            clean = self.worker.oracle(env_id, 115, {"split": "train"})["trace"]
            self.worker.replay(env_id, copy.deepcopy(clean))  # the control: it replays
            for name, tamper in _TAMPERS.items():
                with self.subTest(env=env_id, tamper=name):
                    forged = copy.deepcopy(clean)
                    tamper(forged)
                    with self.assertRaises(TeraError) as caught:
                        self.worker.replay(env_id, forged)
                    self.assertIn("arena", str(caught.exception))

    def test_a_resealed_forgery_is_rejected(self) -> None:
        """The edit that a checksum alone cannot catch.

        Every tamper above breaks the envelope seal, so on its own the previous
        test only proves the seal works. The interesting adversary edits the
        trace *and recomputes the checksum* — a rollout reporting a reward it
        did not earn would look exactly like this. Catching it is `replay()`
        re-running the simulator and comparing, which is the actual claim.
        """
        for env_id in ENVIRONMENT_IDS:
            clean = self.worker.oracle(env_id, 115, {"split": "train"})["trace"]
            for name, tamper in _TAMPERS.items():
                if name == "envelope-checksum":
                    continue  # re-sealing it is just the clean trace again
                with self.subTest(env=env_id, tamper=name):
                    forged = self._reseal(copy.deepcopy(clean), tamper)
                    # The seal is now valid — prove it, or the test is trivially
                    # passing for the wrong reason.
                    core = {k: v for k, v in forged.items() if k != "checksum"}
                    self.assertEqual(forged["checksum"], self.worker.checksum(core))
                    with self.assertRaises(TeraError) as caught:
                        self.worker.replay(env_id, forged)
                    self.assertIn("arena", str(caught.exception))

    def _reseal(self, trace: dict, tamper) -> dict:
        tamper(trace)
        core = {key: value for key, value in trace.items() if key != "checksum"}
        trace["checksum"] = self.worker.checksum(core)
        return trace

    def test_stepping_a_finished_episode_is_an_error(self) -> None:
        reference = self.worker.oracle("crow-nav-v1", 115, {"split": "train"})
        with self.worker.open("crow-nav-v1", 115, {"split": "train"}) as episode:
            for frame in reference["trace"]["steps"]:
                episode.step(frame["action"])
            with self.assertRaises(TeraError):
                episode.step(reference["trace"]["steps"][-1]["action"])

    def test_an_unknown_scenario_is_an_error(self) -> None:
        with self.assertRaises(TeraError):
            self.worker.open("crow-nav-v1", 115, "dev-there-is-no-such-place")

    def test_baselines_come_from_typescript_and_separate(self) -> None:
        """The oracle op exists so the reward denominator is never a Python port.

        Inaction below zero, scripted above it, is Tera's own documented claim;
        checking it here is checking that the bridge is calling the real
        baselines and not, say, handing both policies the same action.
        """
        for env_id in ENVIRONMENT_IDS:
            with self.subTest(env=env_id):
                scripted = self.worker.oracle(env_id, 115, {"split": "train"}, "scripted")
                inaction = self.worker.oracle(env_id, 115, {"split": "train"}, "inaction")
                self.assertGreater(scripted["cumulativeReward"], inaction["cumulativeReward"])
                self.assertLess(inaction["cumulativeReward"], 0.0)
                self.assertGreater(scripted["cumulativeReward"], 0.0)


def _bump_reward(trace: dict) -> None:
    trace["steps"][0]["reward"] += 1.0


def _bump_cumulative(trace: dict) -> None:
    trace["cumulativeReward"] += 1.0


def _swap_action(trace: dict) -> None:
    action = trace["steps"][0]["action"]
    key = sorted(k for k, v in action.items() if isinstance(v, (int, float)))[0]
    action[key] = float(action[key]) + 0.5


def _forge_state_checksum(trace: dict) -> None:
    trace["steps"][0]["stateChecksum"] = "fnv1a64:0000000000000000"


def _forge_envelope_checksum(trace: dict) -> None:
    trace["checksum"] = "fnv1a64:0000000000000000"


def _drop_last_step(trace: dict) -> None:
    # Truncation alone keeps every surviving frame honest, so only the final
    # state and the cumulative reward can catch it.
    trace["steps"] = trace["steps"][:-1]


def _forge_source_hashes(trace: dict) -> None:
    trace["sourceHashes"]["simulator"] = "sha256:" + "0" * 64


_TAMPERS = {
    "reward": _bump_reward,
    "cumulative-reward": _bump_cumulative,
    "action": _swap_action,
    "state-checksum": _forge_state_checksum,
    "envelope-checksum": _forge_envelope_checksum,
    "dropped-step": _drop_last_step,
    "source-hashes": _forge_source_hashes,
}


if __name__ == "__main__":
    unittest.main()