diff --git a/Makefile b/Makefile index 4241221..5aa2157 100644 --- a/Makefile +++ b/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____, 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 diff --git a/ltc/agents/sr_jax.py b/ltc/agents/sr_jax.py index 1d66324..418b1d4 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 @@ -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 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/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) 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()