`si_policy_search` crashes when `num_policies` exceeds max_branching's default (10)

[gauraaansh]

max_branching=10 causes a hard crash for any model with more than 10 policies.

Repro, using the repo's GraphEnv from si_graph_world.ipynb, with the unmodified si_policy_search call from that notebook:

import jax.numpy as jnp
import jax.random as jr
import jax.tree_util as jtu
from pymdp.envs import GraphEnv
from pymdp.envs.graph_worlds import generate_connected_clusters
from pymdp.utils import list_array_zeros, list_array_norm_dist
from pymdp.agent import Agent
from pymdp.planning.si import si_policy_search

# same setup as si_graph_world.ipynb, just cluster_size bumped from 3 to 5
graph, _ = generate_connected_clusters(cluster_size=5, connections=2)  # 11 nodes
env = GraphEnv(graph, object_location=4, agent_location=0)
A, B = env.A, env.B
A_dependencies, B_dependencies = env.A_dependencies, env.B_dependencies
C = list_array_zeros([a.shape[0] for a in A])
C[1] = C[1].at[1].set(10.0)
D = [jnp.ones(d.shape[0]) for d in env.D]
D[0] = D[0].at[0].set(100.0)
D[1] = D[1].at[4].set(100.0)
D = list_array_norm_dist(D)
agent = Agent(A, B, C, D, A_dependencies=A_dependencies, B_dependencies=B_dependencies, policy_len=1)

key = jr.PRNGKey(0)
qs = agent.infer_states(
    observations=jtu.tree_map(lambda x: jnp.broadcast_to(x, (agent.batch_size,) + x.shape), [jnp.array([0]), jnp.array([0])]),
    past_actions=None, empirical_prior=agent.D, qs_hist=None,
)
search_fn = si_policy_search(horizon=3)  # exact call from the notebook, no overrides
q_pi, extra = search_fn(agent, qs, key)

This raises
ValueError: index can't contain negative values from deep inside jnp.pad (pymdp/planning/si.py:402)

and nothing in the traceback hints that the real cause is max_branching being too small.

I tried to isolate the cause with a single variable test on the same graph and agent (num_policieshing. With max_branching=10 (one less than num_policies) it crashes. With max_branching=11
(equal to num_policies) it succeeds. With max_branching=12 it also succeeds. That confirms max_branching < num_policies.

cause:

Tree.children_indices and Tree.children_probs are pre-allocated at shape (batc When a node's real number of surviving children (up to num_policies for a policy node) exceeds max_branching, _update_node's padding step computes a negative pad width, which JAX rejects. Nothing validates max_branching >= num_policies before this point. Since max_branching defaults to 10, any
model with more than 10 policies crashes using pure defaults.

something that can work:

derive max_branching's default from num_policies (policy-node fan-out) and n combinations after topk_obsspace/pruning (observation-node fan-out), instead of a fixed
constant.
Or raise a clear error at construction time if max_branching < num_poli inside jnp.pad with an unrelated looking message.

[conorheins]

@tverbele @riddhipits any insight to share here?

[tverbele]

Yes, max_branching basically allows you to constrain the tree memory. The reason it's not tied to num_policies is that I can have models with a huge number of policies but only a few ones remain after pruning, so I can constrain it with max_branching. You have a valid point about the error message though. One other way could be to set topk to max_branching if its < than num_policies

[gauraaansh]

Makes sense. I tried, capping to the top max_branching policies by probability when num_policies > max_branching , and left max_branching alone as a memory constraint independent of the model,

One thing that happened when doing so: doing the top-k selection breaks the correspondence between q_pi's columns and agent.policies, since the surviving policies get reindexed by probability than be in their original order.

Downstream, sample_policy does policies[argmax(q_pi)], so it silently picks the wrong policy whenever the branching cap actually kicks in, and the wrong action selected. adding a small bit of bookkeeping (tracking which original policy each surviving slot maps to) should help, so the returned q_pi stays correctly aligned with agent.policies regardless of numberr of policies that get capped.