fix(agent): back Policies with a hashable tuple

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

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(