From a47b75056094ac22c4892d6db090fe147b77e58d Mon Sep 17 00:00:00 2001 From: Krzysztof Rusek Date: Thu, 13 Aug 2026 18:39:53 +0200 Subject: [PATCH 1/3] Add history_reshape function and update usages in SRJaxAgent and history2csv --- ltc/agents/sr_jax.py | 5 +++-- ltc/symbolic/history2csv.py | 3 ++- ltc/symbolic/util.py | 14 ++++++++++++++ 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/ltc/agents/sr_jax.py b/ltc/agents/sr_jax.py index 1d66324..fba423a 100644 --- a/ltc/agents/sr_jax.py +++ b/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 @@ -54,7 +54,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 diff --git a/ltc/symbolic/history2csv.py b/ltc/symbolic/history2csv.py index 325229c..d668c53 100644 --- a/ltc/symbolic/history2csv.py +++ b/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) diff --git a/ltc/symbolic/util.py b/ltc/symbolic/util.py index 8a32fe0..8fa47b9 100644 --- a/ltc/symbolic/util.py +++ b/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) From 0c2bd4fcfedffc290e7a921cbbe68d3621854b79 Mon Sep 17 00:00:00 2001 From: Krzysztof Rusek Date: Thu, 13 Aug 2026 20:50:05 +0200 Subject: [PATCH 2/3] Add support for simulation replay in Makefile with new sr-run target --- Makefile | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 4241221..569d4e6 100644 --- a/Makefile +++ b/Makefile @@ -4,8 +4,15 @@ 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 ?= 5 +SR_RUN_EPOCHS ?= 1 +SR_RUN_STEPS ?= 2000 +SR_RUN_FLAGS ?= + +.PHONY: all distill sr report split report-split sr-run clean all: distill sr report split report-split @@ -19,6 +26,8 @@ report: out/report.html report-split: out/report_split.html +sr-run: $(SR_RUN_STAMPS) + .PRECIOUS: $(CSV_FILES) out: @@ -41,6 +50,16 @@ 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____, 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) $(SR_RUN_FLAGS) --save_plots + touch "$@" + out/report_split.html: $(SR_SPLIT_STAMP) marimo export html ltc/symbolic/report_split.py -o "$@" -f From 8ac8a7647e76b499f626cb4c36e2aa91dd9eb662 Mon Sep 17 00:00:00 2001 From: Krzysztof Rusek Date: Thu, 13 Aug 2026 21:09:12 +0200 Subject: [PATCH 3/3] Guard the symbolic agent against a feature-count mismatch PySR's jax callable indexes X by fit-time column, and JAX clamps out-of-range indices instead of raising. Running a model distilled from a window_size=10 history under the default window_size=5 therefore read the wrong columns, collapsed the expression to a constant, and pinned every station to TX for the whole rollout: constant collisions, zero throughput, flat plots. SRJaxAgent now compares feature_names_in_ against the feature count the simulation supplies and names the window size to use. The sr-run target passes that window size explicitly, and defaults to 10 epochs since the all_* plots draw one point per epoch and a single-epoch run renders as blank axes. Also adds the unit test for history_reshape, covering the row and column ordering the CSV columns and the agent both depend on. Co-Authored-By: Claude Opus 5 --- Makefile | 10 ++++-- ltc/agents/sr_jax.py | 14 +++++++- ltc/run.py | 5 ++- test/symbolic/test_util.py | 71 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 95 insertions(+), 5 deletions(-) create mode 100644 test/symbolic/test_util.py diff --git a/Makefile b/Makefile index 569d4e6..5aa2157 100644 --- a/Makefile +++ b/Makefile @@ -7,9 +7,12 @@ SR_SPLIT_STAMP := $(addprefix out/, $(addsuffix .split.done, $(BASENAMES))) SR_RUN_STAMPS := $(addprefix out/, $(addsuffix .srrun.done, $(BASENAMES))) # Simulation replaying a distilled expression through SRJaxAgent. -SR_EQ ?= 5 -SR_RUN_EPOCHS ?= 1 +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 @@ -57,7 +60,8 @@ 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) $(SR_RUN_FLAGS) --save_plots + --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) diff --git a/ltc/agents/sr_jax.py b/ltc/agents/sr_jax.py index fba423a..418b1d4 100644 --- a/ltc/agents/sr_jax.py +++ b/ltc/agents/sr_jax.py @@ -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"] diff --git a/ltc/run.py b/ltc/run.py index 0192ea1..1087085 100644 --- a/ltc/run.py +++ b/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) diff --git a/test/symbolic/test_util.py b/test/symbolic/test_util.py new file mode 100644 index 0000000..7ed2bbf --- /dev/null +++ b/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()