PR-A4: behavior_3 trial-mode test suite (+ cleanup)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: m2kulkarni

tests/test_cache_freeze.py

@@ -0,0 +1,167 @@
+"""Contract: KV cache slots written while an agent is off-map (removed=1) are
+excluded from attention via per-agent garbage_mask.
+
+The attention math is verified directly: we drive the transformer's
+forward_eval with a synthetic state dict, mark some slots as garbage,
+and assert that the attention weights at those slots are zero (after
+softmax).
+
+We use the _probe_attention path in models.py which captures attention
+weights per layer in state["_attn_weights"].
+"""
+
+import os
+
+# Without this, _USE_LEGACY_EVAL defaults True and the streaming KV path
+# (which owns garbage_mask) is bypassed — the test would silently no-op.
+os.environ.setdefault("PUFFER_TRANSFORMER_LEGACY_EVAL", "0")
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+
+class _MinimalPolicy(nn.Module):
+    """Stub that satisfies TransformerWrapper's policy contract:
+    encode_observations(obs, state) -> (B, hidden) and
+    decode_actions(hidden) -> (logits, values)."""
+
+    def __init__(self, obs_dim, hidden, n_actions):
+        super().__init__()
+        self.encoder = nn.Linear(obs_dim, hidden)
+        self.decoder_a = nn.Linear(hidden, n_actions)
+        self.decoder_v = nn.Linear(hidden, 1)
+        self.is_continuous = False
+
+    def encode_observations(self, obs, state=None):
+        return self.encoder(obs)
+
+    def decode_actions(self, hidden):
+        return self.decoder_a(hidden), self.decoder_v(hidden).squeeze(-1)
+
+
+class _StubEnv:
+    """Just exposes single_observation_space.shape — the only thing
+    TransformerWrapper.__init__ reads from env."""
+
+    def __init__(self, obs_dim):
+        from gymnasium import spaces
+
+        self.single_observation_space = spaces.Box(low=-1, high=1, shape=(obs_dim,))
+
+
+def _make_wrapper(batch_size=4, horizon=8, hidden=16, n_heads=2):
+    from pufferlib.models import TransformerWrapper
+
+    env = _StubEnv(obs_dim=hidden)
+    policy = _MinimalPolicy(obs_dim=hidden, hidden=hidden, n_actions=3)
+    return TransformerWrapper(
+        env=env,
+        policy=policy,
+        horizon=horizon,
+        num_layers=2,
+        num_heads=n_heads,
+        input_size=hidden,
+        hidden_size=hidden,
+    )
+
+
+def test_garbage_mask_excludes_slots_from_attention():
+    """If garbage_mask[a, k] = True, the softmax weight at slot k for agent a
+    must be 0 after the next forward."""
+    torch.manual_seed(0)
+    B, T, H = 4, 8, 16
+    wrapper = _make_wrapper(batch_size=B, horizon=T, hidden=H, n_heads=2)
+    wrapper.eval()
+
+    state = wrapper.init_eval_state(batch_size=B, device="cpu", dtype=torch.float32)
+    state["_probe_attention"] = True
+
+    # Step 5 times so cache fills slots 0..4 for all agents. Agent 0 has
+    # `removed=True` at steps 2 and 3 — slots 2, 3 should be marked garbage.
+    for step in range(5):
+        obs = torch.randn(B, H)
+        removed = torch.zeros(B, dtype=torch.bool)
+        if step in (2, 3):
+            removed[0] = True
+        state["removed"] = removed
+        state["_attn_weights"] = []  # reset per step
+        with torch.no_grad():
+            wrapper.forward_eval(obs, state)
+
+    # After step 4: garbage_mask[0, 2] and [0, 3] should be True
+    gm = state["garbage_mask"]
+    assert gm[0, 2].item() and gm[0, 3].item(), f"garbage slots not marked for agent 0: {gm[0]}"
+    # Other agents: nothing marked
+    assert not gm[1:].any().item(), f"non-removed agents should have empty garbage_mask: {gm[1:]}"
+
+    # Now step once more (no removed) and inspect attention weights for agent 0
+    state["removed"] = torch.zeros(B, dtype=torch.bool)
+    state["_attn_weights"] = []
+    obs = torch.randn(B, H)
+    with torch.no_grad():
+        wrapper.forward_eval(obs, state)
+
+    # Per-layer attention weights are (B, H, 1, horizon).
+    # Agent 0's slots 2 and 3 must have zero weight (masked out by garbage_mask).
+    for layer_rec in state["_attn_weights"]:
+        w = layer_rec["weights"]  # (B, H, 1, horizon)
+        assert w[0, :, 0, 2].abs().max().item() < 1e-6, (
+            f"layer {layer_rec['layer']}: slot 2 weight nonzero for agent 0: {w[0, :, 0, 2]}"
+        )
+        assert w[0, :, 0, 3].abs().max().item() < 1e-6, f"layer {layer_rec['layer']}: slot 3 weight nonzero for agent 0"
+        # Sanity: other agents' slots 2, 3 should still get nonzero weight
+        assert w[1, :, 0, 2].abs().max().item() > 1e-6, "agent 1 slot 2 should NOT be masked"
+
+
+def test_garbage_mask_clears_on_full_reset():
+    """reset_eval_state(state, done_indices=None) must zero garbage_mask."""
+    B, T, H = 4, 8, 16
+    wrapper = _make_wrapper(batch_size=B, horizon=T, hidden=H, n_heads=2)
+    state = wrapper.init_eval_state(batch_size=B, device="cpu", dtype=torch.float32)
+    state["garbage_mask"][:] = True
+    wrapper.reset_eval_state(state, done_indices=None)
+    assert not state["garbage_mask"].any().item(), "garbage_mask should be zeroed after full reset"
+
+
+def test_garbage_mask_clears_per_agent_on_partial_reset():
+    """reset_eval_state(state, done_indices=[a]) must zero garbage_mask[a]
+    but leave other agents untouched."""
+    B, T, H = 4, 8, 16
+    wrapper = _make_wrapper(batch_size=B, horizon=T, hidden=H, n_heads=2)
+    state = wrapper.init_eval_state(batch_size=B, device="cpu", dtype=torch.float32)
+    state["garbage_mask"][:] = True
+    wrapper.reset_eval_state(state, done_indices=torch.tensor([1, 3]))
+    assert state["garbage_mask"][0].all().item(), "agent 0 garbage_mask should be unchanged"
+    assert not state["garbage_mask"][1].any().item(), "agent 1 garbage_mask should be cleared"
+    assert state["garbage_mask"][2].all().item(), "agent 2 garbage_mask should be unchanged"
+    assert not state["garbage_mask"][3].any().item(), "agent 3 garbage_mask should be cleared"
+
+
+def test_no_garbage_mask_no_regression():
+    """Without `removed` in state, forward_eval must still work — model
+    creates a fresh garbage_mask (all False), so attention is unchanged
+    from the pre-fix behavior."""
+    torch.manual_seed(0)
+    B, T, H = 4, 8, 16
+    wrapper = _make_wrapper(batch_size=B, horizon=T, hidden=H, n_heads=2)
+    wrapper.eval()
+    state = wrapper.init_eval_state(batch_size=B, device="cpu", dtype=torch.float32)
+    # No state["removed"] key
+    for _ in range(3):
+        obs = torch.randn(B, H)
+        with torch.no_grad():
+            wrapper.forward_eval(obs, state)
+    # garbage_mask stays all-False
+    assert not state["garbage_mask"].any().item(), "without removed signal, no slots should be marked garbage"
+
+
+if __name__ == "__main__":
+    test_garbage_mask_excludes_slots_from_attention()
+    test_garbage_mask_clears_on_full_reset()
+    test_garbage_mask_clears_per_agent_on_partial_reset()
+    test_no_garbage_mask_no_regression()
+    print("test_cache_freeze: PASS")

tests/test_env_level_trial.py

@@ -0,0 +1,178 @@
+"""Contract tests for B'' env-level trial semantic.
+
+Design (see docs/src/trial_mode.md, "Env-level trials"):
+  - Each env has ONE trial clock (env->trial_count, env->trial_start_timestep),
+    not per-agent.
+  - On ego goal-reach mid-trial: ego goes off-map (removed=1, INVALID_POSITION,
+    vx=vy=0). No truncations / terminals yet — wait for trial-end.
+  - env trial-end fires when ALL active egos in env have removed=1 OR env's
+    per_trial_timeout elapses since trial start. At env trial-end:
+      * truncations[i] = 1 for every active ego in env
+      * trial_ended_this_step[i] = 1 for every active ego in env
+      * All entities (egos + co-players) reset to init position; removed=0
+      * env->trial_count++, env->trial_start_timestep = env->timestep
+  - At env episode-end (env->trial_count == max_trials):
+      * terminals[i] = 1 for every active ego in env
+      * Option D: all egos removed=1 + off-map until c_reset
+
+These tests run on a tiny env (per_trial_timeout=5, k=2) for determinism.
+"""
+
+import os
+import sys
+import numpy as np
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+MAP_DIR = "resources/drive/binaries/nuplan_201"
+INI = "pufferlib/config/ocean/drive.ini"
+
+
+def _make_env(k=2, scenario_length=5, num_agents=4, goal_radius=2.0):
+    from pufferlib.ocean.drive import Drive
+
+    return Drive(
+        num_agents=num_agents,
+        map_dir=MAP_DIR,
+        num_maps=10,
+        scenario_length=scenario_length,
+        ini_file=INI,
+        goal_behavior=3,
+        k_scenarios=k,
+        max_trials_per_episode=k,
+        per_trial_timeout=scenario_length,
+        goal_radius=goal_radius,
+        report_interval=10000,
+    )
+
+
+def _zero_actions(env):
+    return np.zeros(env.action_space.shape, dtype=env.actions.dtype)
+
+
+def test_removed_buffer_exists_and_is_zero_at_reset():
+    """Python-side `removed` SHM buffer exists and starts all-zero."""
+    env = _make_env()
+    env.reset(seed=42)
+    assert hasattr(env, "removed"), "env must expose a `removed` SHM buffer"
+    assert np.asarray(env.removed, dtype=bool).shape == (env.num_agents,)
+    assert not np.asarray(env.removed, dtype=bool).any(), "removed must be all-zero after reset"
+    env.close()
+
+
+def test_env_trial_end_fires_on_timeout_only():
+    """Tight goal_radius so no ego reaches. Trial-end MUST fire at
+    per_trial_timeout for the env, with truncations=1 on every ego."""
+    env = _make_env(k=2, scenario_length=5, num_agents=4, goal_radius=2.0)
+    env.reset(seed=42)
+    actions = _zero_actions(env)
+    truncations_at = None
+    for t in range(1, 10):
+        env.step(actions)
+        if np.asarray(env.truncations, dtype=bool).any():
+            truncations_at = t
+            break
+    assert truncations_at == 5, f"trial-end (timeout) should fire at tick=5, got {truncations_at}"
+    # Trial-end fires for ALL active agents simultaneously
+    tr = np.asarray(env.truncations, dtype=bool)
+    te = np.asarray(env.trial_ended_this_step, dtype=bool)
+    assert tr.all() or tr.sum() >= 1, f"truncations should fire env-wide: {tr}"
+    assert (tr == te).all(), f"truncations and trial_ended_this_step must align: tr={tr}, te={te}"
+    env.close()
+
+
+def test_ego_goes_off_map_on_reach():
+    """Wide goal_radius so all egos reach quickly. Each ego should become
+    removed=1 the step after it reaches goal."""
+    env = _make_env(k=2, scenario_length=20, num_agents=4, goal_radius=200.0)
+    env.reset(seed=42)
+    actions = _zero_actions(env)
+    if env.action_space.shape[-1] == 2:
+        actions[:, 0] = 0.1  # gentle accel — within speed limit
+    saw_removed = False
+    for _ in range(20):
+        env.step(actions)
+        if np.asarray(env.removed, dtype=bool).any():
+            saw_removed = True
+            break
+    assert saw_removed, "At least one ego should have removed=1 mid-trial after reaching goal"
+    env.close()
+
+
+def test_env_trial_end_resets_all_entities_to_init():
+    """After a NON-terminal env trial-end (trial < max_trials), all egos
+    must be back on-map (removed=0). k must be >= 3 so trial 1 end isn't
+    the same step as episode-end."""
+    env = _make_env(k=3, scenario_length=5, num_agents=4, goal_radius=200.0)
+    env.reset(seed=42)
+    actions = _zero_actions(env)
+    if env.action_space.shape[-1] == 2:
+        actions[:, 0] = 0.1
+    # Run until first env trial-end
+    for t in range(1, 20):
+        env.step(actions)
+        if np.asarray(env.truncations, dtype=bool).any():
+            # At trial-end step itself, the reset has already fired in C —
+            # removed should already be 0 (entities back at init).
+            removed_after = np.asarray(env.removed, dtype=bool)
+            term = np.asarray(env.terminals, dtype=bool)
+            assert not term.any(), f"trial 1 must not also be episode-end (k={3}); got terminals={term}"
+            assert not removed_after.any(), (
+                f"After env trial-end (mid-episode), all egos must be back on-map. Got: {removed_after}"
+            )
+            env.close()
+            return
+    raise AssertionError("env trial-end never fired in 20 steps")
+
+
+def test_episode_end_fires_after_max_trials():
+    """After max_trials env trial-ends, terminals must fire for all egos.
+    Option D semantic: removed=1 stays until c_reset."""
+    env = _make_env(k=2, scenario_length=3, num_agents=4, goal_radius=2.0)
+    env.reset(seed=42)
+    actions = _zero_actions(env)
+    trial_ends = 0
+    term_at = None
+    for t in range(1, 20):
+        env.step(actions)
+        if np.asarray(env.truncations, dtype=bool).any():
+            trial_ends += 1
+        if np.asarray(env.terminals, dtype=bool).any():
+            term_at = t
+            break
+    assert trial_ends >= 1, f"expected ≥1 trial-end before episode end, got {trial_ends}"
+    assert term_at is not None, "terminals never fired within 20 steps"
+    # At terminals, all egos should be removed (Option D)
+    assert np.asarray(env.removed, dtype=bool).all(), (
+        f"after episode-end terminals, all egos should be removed: {env.removed}"
+    )
+    env.close()
+
+
+def test_truncations_not_fired_on_individual_reach():
+    """Before env trial-end, individual reaches must NOT fire truncations.
+    Only the env-level trial-end does."""
+    env = _make_env(k=2, scenario_length=30, num_agents=4, goal_radius=200.0)
+    env.reset(seed=42)
+    actions = _zero_actions(env)
+    if env.action_space.shape[-1] == 2:
+        actions[:, 0] = 0.1
+    for t in range(1, 6):
+        env.step(actions)
+        rem = np.asarray(env.removed, dtype=bool)
+        tr = np.asarray(env.truncations, dtype=bool)
+        # If any ego has removed=1 but not all of them, truncations must NOT
+        # fire yet — we're mid-trial waiting for stragglers.
+        if rem.any() and not rem.all():
+            assert not tr.any(), f"step={t}: removed={rem} but truncations={tr} — env trial-end fired prematurely"
+    env.close()
+
+
+if __name__ == "__main__":
+    test_removed_buffer_exists_and_is_zero_at_reset()
+    test_env_trial_end_fires_on_timeout_only()
+    test_ego_goes_off_map_on_reach()
+    test_env_trial_end_resets_all_entities_to_init()
+    test_episode_end_fires_after_max_trials()
+    test_truncations_not_fired_on_individual_reach()
+    print("test_env_level_trial: PASS")