Debug and test SRJaxAgent

Makefile

@@ -4,8 +4,18 @@ CSV_FILES      := $(addprefix out/, $(addsuffix .csv, $(BASENAMES)))
 FOREST_FILES   := $(addprefix out/, $(addsuffix .forest.pkl, $(BASENAMES)))
 SR_STAMPS      := $(addprefix out/, $(addsuffix .sr.done, $(BASENAMES)))
 SR_SPLIT_STAMP := $(addprefix out/, $(addsuffix .split.done, $(BASENAMES)))
+SR_RUN_STAMPS  := $(addprefix out/, $(addsuffix .srrun.done, $(BASENAMES)))
 
-.PHONY: all distill sr report split report-split clean
+# Simulation replaying a distilled expression through SRJaxAgent.
+SR_EQ          ?= 2
+# The `all_*` plots draw one point per epoch, so a single epoch renders as blank axes.
+SR_RUN_EPOCHS  ?= 10
+SR_RUN_STEPS   ?= 2000
+# Must match the window the model was distilled from, else the expression reads wrong columns.
+SR_WINDOW      ?= 10
+SR_RUN_FLAGS   ?=
+
+.PHONY: all distill sr report split report-split sr-run clean
 
 all: distill sr report split report-split
 
@@ -19,6 +29,8 @@ report: out/report.html
 
 report-split: out/report_split.html
 
+sr-run: $(SR_RUN_STAMPS)
+
 .PRECIOUS: $(CSV_FILES)
 
 out:
@@ -41,6 +53,17 @@ out/%.split.done: out/%.csv ltc/symbolic/sr_split.py
 	python -m ltc.symbolic.sr_split --file "$<" --output "out/$*" --pysr_output_dir out/output_split
 	touch "$@"
 
+# Run the simulator with the distilled expression as the agent policy.
+# The stem is history_<n>_<n_final>_<seed>_<commit>, so n and seed come back out of it.
+# ltc.run writes its own history_*.pkl.lz4 in the repo root, hence the stamp file.
+out/%.srrun.done: out/%.split.done
+	python -m ltc.run --agent_type sr-jax \
+		--sr_pkl "out/$*.split_sr.pkl" --sr_eq $(SR_EQ) \
+		--n $(word 2,$(subst _, ,$*)) --seed $(word 4,$(subst _, ,$*)) \
+		--n_epochs $(SR_RUN_EPOCHS) --n_steps $(SR_RUN_STEPS) --window_size $(SR_WINDOW) \
+		$(SR_RUN_FLAGS) --save_plots
+	touch "$@"
+
 out/report_split.html: $(SR_SPLIT_STAMP)
 	marimo export html ltc/symbolic/report_split.py -o "$@" -f
 

ltc/agents/sr_jax.py

@@ -8,7 +8,7 @@
 
 from ltc.sim.constants import Features
 from ltc.sim.features import select_features
-from ltc.symbolic.util import SimplexCode
+from ltc.symbolic.util import SimplexCode, history_reshape
 
 
 @dataclass
@@ -19,7 +19,19 @@ class SRJaxState(AgentState):
 class SRJaxAgent(BaseAgent):
     FEATURES = tuple(Features)
 
-    def __init__(self, sr_model, equation_index: int, n_actions: int = 2):
+    def __init__(self, sr_model, equation_index: int, n_actions: int = 2, n_features: int | None = None):
+        feature_names = getattr(sr_model, 'feature_names_in_', None)
+        expected = 0 if feature_names is None else len(feature_names)
+        if n_features is not None and expected and n_features != expected:
+            # PySR's jax callable indexes X by fit-time column, and JAX clamps out-of-range
+            # indices instead of raising, so a mismatch silently evaluates the expression on
+            # the wrong columns and the policy degenerates to a constant action.
+            raise ValueError(
+                f'The symbolic model was fitted on {expected} features but the simulation supplies '
+                f'{n_features}. Set --window_size {expected // len(Features)} to match the history '
+                f'the model was distilled from.'
+            )
+
         jaxeq = sr_model.jax(equation_index)
         callable_fn = jax.jit(jaxeq["callable"])
         parameters = jaxeq["parameters"]
@@ -54,7 +66,8 @@ def sample(
     ) -> Array:
         # env_state: [window_size, n_features] raw int obs
         env_state = select_features(env_state, SRJaxAgent.FEATURES)
-        x = env_state.reshape(-1).astype(jnp.float32)[jnp.newaxis]  # [1, w*f]
+        # a single obs is a one-step, one-agent history: [1, 1, w, f] -> [1, w*f]
+        x = history_reshape(env_state[jnp.newaxis, jnp.newaxis]).astype(jnp.float32)
         yhat = callable_fn(x, state.parameters)                     # [T-1] or [1*(T-1)]
         codes = yhat.reshape(1, simplex.T - 1)                      # [1, T-1]
         return simplex.decode(codes)[0]                             # scalar

ltc/run.py

@@ -386,7 +386,10 @@ def setup_args():
             raise ValueError('--sr_pkl is required when --agent_type is sr-jax.')
         with open(args.sr_pkl, 'rb') as f:
             sr_model = pickle.load(f)
-        drl = SRJaxAgent(sr_model, equation_index=args.sr_eq, n_actions=num_actions)
+        drl = SRJaxAgent(
+            sr_model, equation_index=args.sr_eq, n_actions=num_actions,
+            n_features=window_size * len(Features)
+        )
     elif agent_type == 'aloha-qtf':
         drl = ALOHAQTF()
         node_ids = jnp.arange(n, dtype=jnp.int32)

ltc/symbolic/history2csv.py

@@ -10,6 +10,7 @@
 import pandas as pd
 
 from ltc.agents import QNetwork
+from ltc.symbolic.util import history_reshape
 from ltc.utils.history import resolve_history_file, unpack_history
 
 FEATURE_NAMES = ["buffer", "channel", "ret_c", "no_tx", "action_tx", "action_cs"]
@@ -80,7 +81,7 @@ def compute_qvals(params, state, observations, key):
 
     # Build flat feature matrix: [n_agents * n_steps, window_size * n_features]
     # transpose [n_steps, n_agents, w, f] -> [n_agents, n_steps, w, f] then flatten last two dims
-    XX = np.asarray(observations).transpose(1, 0, 2, 3).reshape(n_agents * n_steps, -1)
+    XX = history_reshape(observations)
     actions = np.asarray(qvals.argmax(axis=-1)).flatten()  # [n_agents * n_steps]
     agent_ids = np.repeat(np.arange(n_agents), n_steps)
 

ltc/symbolic/util.py

@@ -65,6 +65,20 @@ def decode(self, codes: jax.Array) -> jax.Array:
         return jnp.argmax(inner_products, axis=1)
 
 
+@jax.jit
+def history_reshape(observations: jax.Array) -> jax.Array:
+    """
+    Flatten a history of observations into a PySR-style feature matrix.
+
+    [n_steps, n_agents, window_size, n_features] -> [n_agents * n_steps, window_size * n_features]
+
+    Rows are grouped by agent (agent-major), and within a row the features are
+    window-step-major, matching ``build_column_names`` in ``history2csv``.
+    """
+    n_steps, n_agents = observations.shape[0], observations.shape[1]
+    return observations.transpose(1, 0, 2, 3).reshape(n_agents * n_steps, -1)
+
+
 if __name__ == "__main__":
     T = 5
     sc = SimplexCode(T=T)

test/symbolic/test_util.py

@@ -0,0 +1,71 @@
+import unittest
+
+import jax.numpy as jnp
+import numpy as np
+
+from ltc.symbolic.util import history_reshape
+
+
+class TestHistoryReshape(unittest.TestCase):
+    N_STEPS = 7
+    N_AGENTS = 3
+    WINDOW = 4
+    N_FEATURES = 6
+
+    def setUp(self):
+        shape = (self.N_STEPS, self.N_AGENTS, self.WINDOW, self.N_FEATURES)
+        # distinct values so any mis-ordering is detectable
+        self.obs = jnp.arange(np.prod(shape), dtype=jnp.int32).reshape(shape)
+
+    def test_shape(self):
+        XX = history_reshape(self.obs)
+        self.assertEqual(
+            XX.shape,
+            (self.N_AGENTS * self.N_STEPS, self.WINDOW * self.N_FEATURES),
+        )
+
+    def test_rows_are_agent_major(self):
+        """Row a*n_steps + t is agent a at step t, flattened."""
+        XX = np.asarray(history_reshape(self.obs))
+        obs = np.asarray(self.obs)
+        for a in range(self.N_AGENTS):
+            for t in range(self.N_STEPS):
+                np.testing.assert_array_equal(
+                    XX[a * self.N_STEPS + t], obs[t, a].reshape(-1)
+                )
+
+    def test_matches_per_agent_concatenation(self):
+        obs = np.asarray(self.obs)
+        expected = np.concatenate(
+            [obs[:, a].reshape(self.N_STEPS, -1) for a in range(self.N_AGENTS)], axis=0
+        )
+        np.testing.assert_array_equal(np.asarray(history_reshape(self.obs)), expected)
+
+    def test_columns_are_window_step_major(self):
+        """Within a row, features are grouped by window step (matches build_column_names)."""
+        XX = np.asarray(history_reshape(self.obs))
+        obs = np.asarray(self.obs)
+        row = XX[0]  # agent 0, step 0
+        for w in range(self.WINDOW):
+            lo, hi = w * self.N_FEATURES, (w + 1) * self.N_FEATURES
+            np.testing.assert_array_equal(row[lo:hi], obs[0, 0, w])
+
+    def test_single_observation_via_newaxis(self):
+        """The SRJaxAgent path: one [w, f] obs promoted to a [1, 1, w, f] history."""
+        single = self.obs[0, 0]
+        x = history_reshape(single[jnp.newaxis, jnp.newaxis])
+        self.assertEqual(x.shape, (1, self.WINDOW * self.N_FEATURES))
+        np.testing.assert_array_equal(np.asarray(x)[0], np.asarray(single).reshape(-1))
+
+    def test_rejects_non_history_rank(self):
+        with self.assertRaises(Exception):
+            history_reshape(jnp.zeros((self.WINDOW, self.N_FEATURES)))
+
+    def test_preserves_dtype(self):
+        self.assertEqual(history_reshape(self.obs).dtype, self.obs.dtype)
+        f = self.obs.astype(jnp.float32)
+        self.assertEqual(history_reshape(f).dtype, jnp.float32)
+
+
+if __name__ == "__main__":
+    unittest.main()