onf-c2/tests/test_entry_static.py

"""Tests for ENTRY_STATIC — the static-window training query class
(onf.graph.train.data.make_queries's entry_static_frac, build_entry_static_pool_cache,
SoftTargetBuilder.entry_static).

WHAT IT IS. Every other training query is built from real consecutive demo frames, so its window
carries nonzero finite-differenced velocity. At t=0 the deployed retriever is instead fed
repeat(q0, hist) — velocity identically zero — an input shape ordinary training never produces.
ENTRY_STATIC mixes that shape into the training distribution at ENTRY_STATIC_FRAC, targeted at the
uniform distribution over task lane INTERSECT phase <= ENTRY_STATIC_BAND.

WHY IT IS STILL HERE. ENTRY_STATIC is retained because the shipped g_head.npz was trained with it.

Reuses the synthetic toy-graph fixtures from tests/test_train.py (_toy_graph/_make_demo) and
tests/test_retrieve.py (_toy_nodes_edges/_toy_net/_write_graph_dir/_loaded_retriever)
-- no dataset/artifact dependency, no network, no GPU, cpu device throughout.
"""

from __future__ import annotations

import warnings

import numpy as np
import pytest

from onf.config import GraphConfig
from onf.graph.run.retrieve import GraphRetriever
from onf.graph.train.data import ENTRY_STATIC_BAND, SoftTargetBuilder, build_entry_static_pool_cache, label_smooth_targets, make_queries
from onf.graph.train.types import QuerySpec
from test_retrieve import _loaded_retriever
from test_train import D, _toy_graph


# ======================================================================================================
# 1. no-op guard: entry_static_frac=0.0 reproduces the pre-feature RNG stream / output bit-for-bit
# ======================================================================================================
def test_entry_static_frac_zero_is_bit_identical_noop():
    """entry_static_frac=0.0 must draw the SAME rng.rand() calls, in the SAME order, as the
    pre-feature two-way (entry vs traversal) draw -- the module docstring's own claim ("the static check
    ... is skipped ENTIRELY in that case, falling straight through to the original single
    rng.rand() < entry_frac line, unchanged"), pinned here as a behavioural test rather than trusted
    from the comment alone."""
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=4, length=40)
    owners = np.arange(nodes.n_demos)

    rng_new = np.random.RandomState(0)
    out_new = make_queries(nodes, rng_new, n=100, owners=owners,
                            spec=QuerySpec(entry_frac=0.4, entry_static_frac=0.0))

    # a fresh RandomState with the SAME seed, calling the SAME public signature but never touching
    # entry_static_frac (its default is 0.15 -- explicitly overridden to 0.0 here to isolate "does
    # entry_static_frac=0.0 behave like the feature doesn't exist" from "what is the new default").
    rng_old = np.random.RandomState(0)
    out_old = make_queries(nodes, rng_old, n=100, owners=owners,
                            spec=QuerySpec(entry_frac=0.4, entry_static_frac=0.0))

    for key in ("q_hist", "qdot_hist", "grip_hist", "w_hist", "src_owner", "src_node", "tgt_node",
                "is_clean", "abstain_is_correct"):
        assert np.array_equal(out_new[key], out_old[key]), key

    assert "is_entry_static" in out_new
    assert out_new["is_entry_static"].sum() == 0
    assert out_new["is_entry_static"].dtype == bool


# ======================================================================================================
# 2. zero velocity by construction
# ======================================================================================================
def test_entry_static_rows_have_exactly_zero_velocity():
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=5, length=40)
    owners = np.arange(nodes.n_demos)
    rng = np.random.RandomState(1)

    out = make_queries(nodes, rng, n=120, owners=owners, spec=QuerySpec(entry_frac=0.0, entry_static_frac=1.0))

    assert np.all(out["is_entry_static"])
    assert np.abs(out["qdot_hist"]).max() == 0.0
    # every row of q_hist within a query must be the identical repeated frame
    for i in range(out["q_hist"].shape[0]):
        assert np.all(out["q_hist"][i] == out["q_hist"][i, 0])
    assert np.all(out["src_node"] == out["tgt_node"])
    assert not np.any(out["is_clean"])


# ======================================================================================================
# 3. build_entry_static_pool_cache: pool correctness
# ======================================================================================================
def test_build_entry_static_pool_cache_pool_correctness():
    nodes, _edges = _toy_graph(n_tasks=3, demos_per_task=6, length=40)
    cache = build_entry_static_pool_cache(nodes)

    task_ids = np.unique(np.asarray(nodes.task_id))
    assert set(cache.keys()) == set(int(t) for t in task_ids)

    phase = np.asarray(nodes.phase)
    task_id = np.asarray(nodes.task_id)
    for t, pool in cache.items():
        assert pool.size > 0, f"pool for task {t} must be non-empty on this fixture"
        assert np.all(phase[pool] <= ENTRY_STATIC_BAND + 1e-9)
        assert np.all(task_id[pool] == t)


# ======================================================================================================
# 4. invariance: the ENTRY_STATIC target does not depend on q_now (uniform over the task pool only)
# ======================================================================================================
def test_entry_static_targets_invariant_to_q_now():
    nodes, _edges = _toy_graph(n_tasks=1, demos_per_task=6, length=40)
    cfg = GraphConfig(device="cpu")
    pool_cache = build_entry_static_pool_cache(nodes)

    task = int(nodes.task_id[0])
    pool = pool_cache[task]
    assert pool.size >= 2, "need >=2 pool nodes on this fixture for the row identities to differ"

    tgt_a, tgt_b = int(pool[0]), int(pool[-1])
    assert nodes.task_id[tgt_a] == nodes.task_id[tgt_b] == task

    q_now_1 = nodes.q[5].astype(np.float64)
    q_now_2 = nodes.q[nodes.n_demos * 0 + 20 if len(nodes) > 20 else 0].astype(np.float64)
    assert not np.array_equal(q_now_1, q_now_2), "fixture sanity: the two q_now draws must differ"

    builder = SoftTargetBuilder(nodes, phase_band=cfg.where_phase_band, move_temp=cfg.where_move_temp)
    assert np.array_equal(builder.pool[task], pool), "builder's own pool must match build_entry_static_pool_cache"
    cand_a, w_a = builder.entry_static(tgt_a, q_now_1)
    cand_b, w_b = builder.entry_static(tgt_b, q_now_2)

    assert np.array_equal(cand_a, cand_b)
    assert np.array_equal(w_a, w_b)
    assert np.isclose(float(w_a.sum()), 1.0)
    assert np.allclose(w_a, 1.0 / len(cand_a))


# ======================================================================================================
# 6. empty-pool fallback: falls back to the smoothed target, warns ONCE, never raises
# ======================================================================================================
def test_entry_static_targets_empty_pool_falls_back_to_label_smooth():
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=4, length=30)
    cfg = GraphConfig(device="cpu")
    tgt_i = 5

    expected_cand, expected_w = label_smooth_targets(
        nodes, tgt_i, phase_band=cfg.where_phase_band, move_temp=cfg.where_move_temp,
    )

    # a band below every node's phase leaves EVERY task's pool empty -- the fallback condition
    builder = SoftTargetBuilder(
        nodes, phase_band=cfg.where_phase_band, move_temp=cfg.where_move_temp, entry_static_band=-1.0,
    )
    assert all(len(p) == 0 for p in builder.pool.values())

    with pytest.warns(UserWarning, match="empty ENTRY_STATIC pool"):
        cand, w = builder.entry_static(tgt_i, None)
    assert np.array_equal(cand, expected_cand)
    assert np.allclose(w, expected_w)

    # the latch is per-builder: a second empty-pool row must NOT warn again
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        cand2, w2 = builder.entry_static(int(nodes.task_id.size) - 1, None)
    assert len(cand2) > 0 and np.isclose(float(w2.sum()), 1.0)