From c150977c6208d7609cfc829068c6a4e4620b866d Mon Sep 17 00:00:00 2001 From: gauraansh Date: Mon, 6 Jul 2026 03:30:54 +0530 Subject: [PATCH] fix(agent): back Policies with a hashable tuple Agent.policies is marked static=True to support vmap/scan over batched Agents (#227), but Policies currently stores policy_arr as a jnp.ndarray. As described in #346, constructing equivalent Agents repeatedly causes eqx.filter_jit to miss its compilation cache. Store policies internally as a hashable nested tuple instead, exposing policy_arr via a property that materializes a JAX array on access. This preserves the static Agent.policies partition while allowing equivalent policy tables to hash by value. Constructing an Agent inside an active jax.jit trace required an additional change, since traced arrays cannot be converted back to concrete Python values. To avoid this, add pure-Python tuple builders for policy construction, and only convert to a JAX array at the property/return boundary. Fixes #346 --- pymdp/agent.py | 60 ++++++++++++++++++++++++++++++++++----- pymdp/control.py | 74 ++++++++++++++++++++++++++++++++++++++---------- 2 files changed, 112 insertions(+), 22 deletions(-) diff --git a/pymdp/agent.py b/pymdp/agent.py index f9df47cf..2ca2700e 100644 --- a/pymdp/agent.py +++ b/pymdp/agent.py @@ -211,14 +211,14 @@ def __init__( if ( policies is None and B_action_dependencies is not None ): # note, this only works when B_action_dependencies is not the trivial case of [[0], [1], ...., [num_factors-1]] - policies_multi = control.construct_policies( + policies_multi_tup = control._construct_policies_tuple( self.num_controls_multi, self.num_controls_multi, policy_len, control_fac_idx, ) B, pB, self.action_maps = self._flatten_B_action_dims(B, pB, self.B_action_dependencies) - policies = self._construct_flattend_policies(policies_multi, self.action_maps) + policies = self._construct_flattend_policies_tuple(policies_multi_tup, self.action_maps) self.sampling_mode = "full" # extract shapes from A and B @@ -307,18 +307,23 @@ def __init__( # construct policies if policies is None: - policies_array = control.construct_policies( + policies_tup = control._construct_policies_tuple( self.num_states, self.num_controls, self.policy_len, self.control_fac_idx, ) - self.policies = control.Policies(policies_array) + self.policies = control.Policies(policies_tup) else: - if not isinstance(policies, control.Policies): - self.policies = control.Policies(jnp.array(policies)) - else: + if isinstance(policies, control.Policies): self.policies = policies + elif isinstance(policies, tuple): + # already a concrete nested tuple (e.g. from `_construct_flattend_policies_tuple`); + # pass it straight through so `Policies.__init__` doesn't need to round-trip it + # through `jnp.array` first, which would defeat the point under an active jit trace + self.policies = control.Policies(policies) + else: + self.policies = control.Policies(jnp.array(policies)) if C is None: C = [jnp.ones((self.batch_size, self.num_obs[m])) / self.num_obs[m] for m in range(self.num_modalities)] @@ -1136,6 +1141,47 @@ def _construct_flattend_policies(self, policies: Array, action_maps: list[dict[s policies_flat = jnp.stack(policies_flat, axis=-1) return policies_flat + def _construct_flattend_policies_tuple( + self, policies_tup: tuple, action_maps: list[dict[str, Any]] + ) -> tuple: + """ + Pure-Python equivalent of `_construct_flattend_policies`, operating on and + returning nested tuples instead of `jnp.ndarray`. Mirrors + `utils.get_combination_index`'s mixed-radix encoding exactly, but never invokes + a JAX primitive, so it's safe to call while `Agent.__init__` is itself being + traced under `jax.jit` (see `control._construct_policies_tuple`'s docstring for + why that matters). + """ + num_policies = len(policies_tup) + horizon = len(policies_tup[0]) + + columns = [] + for action_map in action_maps: + multi_dependency = action_map["multi_dependency"] + if multi_dependency == []: + columns.append(tuple(tuple(0 for _ in range(horizon)) for _ in range(num_policies))) + continue + + dims = action_map["multi_dims"] + column = [] + for policy in policies_tup: + row = [] + for t in range(horizon): + xs = [policy[t][d] for d in multi_dependency] + index = 0 + product = 1 + for i in reversed(range(len(dims))): + index += xs[i] * product + product *= dims[i] + row.append(index) + column.append(tuple(row)) + columns.append(tuple(column)) + + return tuple( + tuple(tuple(columns[a][pol][t] for a in range(len(action_maps))) for t in range(horizon)) + for pol in range(num_policies) + ) + def _get_default_params(self) -> dict[str, Any] | None: method = self.inference_algo default_params = None diff --git a/pymdp/control.py b/pymdp/control.py index 2df5d038..20b3e546 100644 --- a/pymdp/control.py +++ b/pymdp/control.py @@ -18,23 +18,46 @@ from pymdp.maths import factor_dot, log_stable, stable_entropy, stable_xlogx, spm_wnorm from pymdp import utils +def _to_nested_tuple(x): + """Recursively convert nested lists/tuples to nested tuples, for hashable-by-value storage.""" + if isinstance(x, (list, tuple)): + return tuple(_to_nested_tuple(v) for v in x) + return x + + class Policies(eqx.Module): - """ + """ A class for storing an array of policies and its properties - + """ - policy_arr: Array + _policy_tup: tuple = eqx.field(static=True) + _dtype: jnp.dtype = eqx.field(static=True) horizon: int = eqx.field(static=True) num_policies: int = eqx.field(static=True) - def __init__(self, policy_arr: Array) -> None: - self.num_policies = policy_arr.shape[0] - self.horizon = policy_arr.shape[1] - self.policy_arr = policy_arr - + def __init__(self, policy_arr: Array | tuple) -> None: + if isinstance(policy_arr, tuple): + # Already a concrete nested tuple (e.g. from `_construct_policies_tuple`). + # Building it never invoked a JAX primitive, so it's safe to use directly + # even if this constructor is itself being called while `Agent.__init__` is + # traced under `jax.jit` -- there's no array here to concretize. + self._policy_tup = policy_arr + self.num_policies = len(policy_arr) + self.horizon = len(policy_arr[0]) + self._dtype = jnp.int32 + else: + self.num_policies = policy_arr.shape[0] + self.horizon = policy_arr.shape[1] + self._dtype = policy_arr.dtype + self._policy_tup = _to_nested_tuple(policy_arr.tolist()) + + @property + def policy_arr(self) -> Array: + return jnp.array(self._policy_tup, dtype=self._dtype) + def __getitem__(self, idx: int) -> Array: return self.policy_arr[idx] - + def __len__(self) -> int: return self.num_policies @@ -175,6 +198,27 @@ def construct_policies( Policy matrix with shape `(num_policies, policy_len, num_factors)`. """ + policies_tup = _construct_policies_tuple(num_states, num_controls, policy_len, control_fac_idx) + return jnp.array(policies_tup, dtype=jnp.int32) + + +def _construct_policies_tuple( + num_states: Sequence[int], + num_controls: Sequence[int] | None = None, + policy_len: int = 1, + control_fac_idx: Sequence[int] | None = None, +) -> tuple: + """ + Pure-Python (no JAX operations) combinatorial construction of the policy table, as a + nested tuple of shape `(num_policies, policy_len, num_factors)`. + + `num_states`/`num_controls`/`policy_len`/`control_fac_idx` are always plain Python + values (never traced), so this never invokes a JAX primitive. That makes it safe to + call even while `Agent.__init__` is itself being traced under `jax.jit`, unlike + building the result via `jnp.array`/`jnp.stack` first: converting a traced array back + with `.tolist()` raises `ConcretizationTypeError`, since `.tolist()` requires a + concrete value that doesn't exist until the trace finishes. + """ num_factors = len(num_states) if control_fac_idx is None: if num_controls is not None: @@ -184,14 +228,14 @@ def construct_policies( if num_controls is None: num_controls = [num_states[c_idx] if c_idx in control_fac_idx else 1 for c_idx in range(num_factors)] - + x = num_controls * policy_len - policies = list(itertools.product(*[list(range(i)) for i in x])) - - for pol_i in range(len(policies)): - policies[pol_i] = jnp.array(policies[pol_i]).reshape(policy_len, num_factors) + flat_policies = itertools.product(*[range(i) for i in x]) - return jnp.stack(policies) + return tuple( + tuple(flat[t * num_factors : (t + 1) * num_factors] for t in range(policy_len)) + for flat in flat_policies + ) def update_posterior_policies(