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(