POMDPPlanners.core.tree package

MCTS tree data structures.

Two implementations live here:

POMDPPlanners.core.tree.anytree_based — node-object tree (each node is a Python object subclassing anytree.NodeMixin). The original implementation, preserved for backward compatibility.
POMDPPlanners.core.tree.arena — column-store SoA tree (Tree holds one list per node attribute, nodes are integer IDs). Faster on every measured tree-side operation.

The legacy classes (ActionNode, BeliefNode, BaseNode, print_tree, get_optimal_action_*) are re-exported from this package so existing from POMDPPlanners.core.tree import ActionNode imports keep working.

class POMDPPlanners.core.tree.ActionNode(action, parent=None, children=(), data=None)[source]

Bases: BaseNode

Parameters:: data (Any)

extend_child_cdf(weight)[source]

Append a newly-attached child’s weight to the cached CDF.

If no CDF is cached yet, this is a no-op — the next sample will rebuild the full CDF from current children.

Return type:: None
Parameters:: weight (float | int)

get_belief_node_child(observation, environment)[source]

Return type:

Optional[BeliefNode]

Parameters:

observation (Any)
environment (Environment)

invalidate_child_cdf()[source]

Clear the cached child sampling CDF.

Called by BeliefNode when a child’s weight mutates or it detaches, forcing sample_child_node() to rebuild from scratch.

Return type:: None

property name

print()[source]

sample_child_node()[source]

Return type:: BeliefNode

property spec

class POMDPPlanners.core.tree.BaseNode(parent=None, children=(), data=None)[source]

Bases: NodeMixin

Parameters:: data (Any)

property immediate_cost: float | None: Get the immediate cost value.

property immediate_reward: float | None: Get the immediate reward value.

class POMDPPlanners.core.tree.BeliefNode(belief, observation=None, weight=1.0, parent=None, children=(), data=None)[source]

Bases: BaseNode

Parameters:

belief (Belief)
observation (Any)
weight (float | int)
data (Any)

get_child(action)[source]

Return type:: Optional[ActionNode]
Parameters:: action (Any)

property name

print()[source]

property spec

update_belief(action, observation, pomdp, **kwargs)[source]

Parameters:

action (Any)
observation (Any)
pomdp (Environment)

property weight: float | int

POMDPPlanners.core.tree.get_optimal_action_cost_setting(belief_node)[source]

Return type:: Any
Parameters:: belief_node (BeliefNode)

POMDPPlanners.core.tree.get_optimal_action_reward_setting(belief_node)[source]

Return type:: Any
Parameters:: belief_node (BeliefNode)

POMDPPlanners.core.tree.print_tree(tree)[source]

Parameters:: tree (BeliefNode | ActionNode)

Submodules

POMDPPlanners.core.tree.anytree_based module

Anytree-based implementation of the MCTS tree.

Each node is a Python object (subclass of anytree.NodeMixin) with parent/children references. This is the original tree implementation, preserved alongside the column-store arena implementation in POMDPPlanners.core.tree.arena.

Use:

from POMDPPlanners.core.tree.anytree_based import (
    ActionNode, BeliefNode, get_optimal_action_cost_setting,
)

class POMDPPlanners.core.tree.anytree_based.ActionNode(action, parent=None, children=(), data=None)[source]

Bases: BaseNode

Parameters:: data (Any)

extend_child_cdf(weight)[source]

Append a newly-attached child’s weight to the cached CDF.

If no CDF is cached yet, this is a no-op — the next sample will rebuild the full CDF from current children.

Return type:: None
Parameters:: weight (float | int)

get_belief_node_child(observation, environment)[source]

Return type:

Optional[BeliefNode]

Parameters:

observation (Any)
environment (Environment)

invalidate_child_cdf()[source]

Clear the cached child sampling CDF.

Called by BeliefNode when a child’s weight mutates or it detaches, forcing sample_child_node() to rebuild from scratch.

Return type:: None

property name

print()[source]

sample_child_node()[source]

Return type:: BeliefNode

property spec

class POMDPPlanners.core.tree.anytree_based.BaseNode(parent=None, children=(), data=None)[source]

Bases: NodeMixin

Parameters:: data (Any)

property immediate_cost: float | None: Get the immediate cost value.

property immediate_reward: float | None: Get the immediate reward value.

class POMDPPlanners.core.tree.anytree_based.BeliefNode(belief, observation=None, weight=1.0, parent=None, children=(), data=None)[source]

Bases: BaseNode

Parameters:

belief (Belief)
observation (Any)
weight (float | int)
data (Any)

get_child(action)[source]

Return type:: Optional[ActionNode]
Parameters:: action (Any)

property name

print()[source]

property spec

update_belief(action, observation, pomdp, **kwargs)[source]

Parameters:

action (Any)
observation (Any)
pomdp (Environment)

property weight: float | int

POMDPPlanners.core.tree.anytree_based.get_optimal_action_cost_setting(belief_node)[source]

Return type:: Any
Parameters:: belief_node (BeliefNode)

POMDPPlanners.core.tree.anytree_based.get_optimal_action_reward_setting(belief_node)[source]

Return type:: Any
Parameters:: belief_node (BeliefNode)

POMDPPlanners.core.tree.anytree_based.print_tree(tree)[source]

Parameters:: tree (BeliefNode | ActionNode)

POMDPPlanners.core.tree.arena module

Column-store implementation of the MCTS tree.

A single Tree instance holds one Python list per node attribute. A node is identified by an integer ID — its index into every column. There are no node objects; operations on the tree are methods of this class that take (node_id, ...) and read or write the columns directly.

Use:

from POMDPPlanners.core.tree.arena import Tree, BELIEF, ACTION

tree = Tree()
root_id = tree.add_belief_node(some_belief)
action_id = tree.add_action_node("up", parent_id=root_id)
child_id = tree.add_belief_node(next_belief, observation="o", parent_id=action_id)
tree.visit_count[root_id] += 1

Schema is declared once in __init__. Adding a new per-node attribute means adding a line to __init__ (the empty column) and one matching line to _allocate (the default value at allocation). Direct attribute access is used in _allocate rather than a generic setattr loop because the per-allocation cost is on the hot path.

This layout uses contiguous typed vectors instead of a graph of Python objects, integer IDs instead of object references, and a single allocation point so per-node overhead is one append per column rather than a Python class instantiation plus GC bookkeeping.

In addition to the columns, the tree maintains two per-parent indexes:

children_cdf[parent] — cumulative distribution function over child weights, in the same order as children_ids[parent]. Maintained incrementally on every add_*_node. Enables O(log K) weighted sampling via sample_belief_child.
obs_child_lookup and action_child_lookup — (parent_id, key) → child_id dicts populated when the observation/action is hashable. Enable O(1) child lookup via get_belief_child_indexed and get_action_child_indexed.

If the user mutates weight[id] after add_belief_node the CDF for the parent becomes stale; call recompute_children_cdf(parent_id) to rebuild. Indexed lookup is unavailable for unhashable observations/actions (e.g. np.ndarray); use the linear-scan get_belief_child / get_action_child for those.

class POMDPPlanners.core.tree.arena.Tree[source]

Bases: object

Column-store search tree.

add_action_node(action, parent_id, data=None, action_key=None)[source]

Allocate an action node under parent_id and return its ID.

Also: (a) extend the parent’s CDF by 1.0 (action nodes default to unit weight); (b) register (parent_id, key) → node_id in action_child_lookup where key is the explicit action_key if provided (caller guarantees hashable), or the raw action (silently dropped if unhashable).

Return type:

int

Parameters:

action (Any)
parent_id (int)
data (Any)
action_key (Hashable | None)

add_belief_node(belief, observation=None, weight=1.0, parent_id=None, data=None, obs_key=None)[source]

Allocate a belief node and return its ID.

If parent_id is provided, also: (a) extend the parent’s CDF by weight; (b) register (parent_id, key) → node_id in obs_child_lookup where key is the explicit obs_key if provided (caller guarantees hashable), or the raw observation (silently dropped if unhashable).

Return type:

int

Parameters:

belief (Belief)
observation (Any)
weight (float)
parent_id (int | None)
data (Any)
obs_key (Hashable | None)

best_action_by_cost(belief_id)[source]

Return the action of the lowest-q_value action child of belief_id.

Arena equivalent of get_optimal_action_cost_setting() from anytree_based.

Return type:: Any
Parameters:: belief_id (int)

best_action_by_reward(belief_id)[source]

Return the action of the highest-q_value action child of belief_id.

Arena equivalent of get_optimal_action_reward_setting() from anytree_based.

Return type:: Any
Parameters:: belief_id (int)

depth(node_id)[source]

Number of edges from this node to the root.

Return type:: int
Parameters:: node_id (int)

get_action(action_id)[source]

Return type:: Any
Parameters:: action_id (int)

get_action_child(belief_id, action)[source]

Return the ID of the action child of belief_id matching action.

Linear scan; compares numpy arrays by content, everything else by ==. Use when action may not be hashable; otherwise prefer get_action_child_indexed() for O(1) lookup.

Return type:

Optional[int]

Parameters:

belief_id (int)
action (Any)

get_action_child_indexed(belief_id, action=None, action_key=None)[source]

O(1) lookup of an action child of belief_id by hashable key.

If action_key is provided the caller guarantees it is hashable; the lookup uses it directly. Otherwise, the raw action is tried and None is returned for unhashable values.

Return type:

Optional[int]

Parameters:

belief_id (int)
action (Any)
action_key (Hashable | None)

get_belief(belief_id)[source]

Return type:: Belief
Parameters:: belief_id (int)

get_belief_child(action_id, observation, env)[source]

Return the ID of the belief child of action_id matching observation.

Linear scan using env.is_equal_observation. Use this when the observation may not be hashable (e.g. np.ndarray) or when equality is custom; otherwise prefer get_belief_child_indexed() for O(1) lookup.

Return type:

Optional[int]

Parameters:

action_id (int)
observation (Any)
env (Environment)

get_belief_child_indexed(action_id, observation=None, obs_key=None)[source]

O(1) lookup of a belief child of action_id by hashable key.

If obs_key is provided the caller guarantees it is hashable; the lookup uses it directly. Otherwise, the raw observation is tried and None is returned for unhashable values.

Return type:

Optional[int]

Parameters:

action_id (int)
observation (Any)
obs_key (Hashable | None)

get_children_ids(node_id)[source]

Return type:: List[int]
Parameters:: node_id (int)

get_immediate_cost(action_id)[source]

Return type:: Optional[float]
Parameters:: action_id (int)

get_immediate_reward(action_id)[source]

Return type:: Optional[float]
Parameters:: action_id (int)

get_observation(belief_id)[source]

Return type:: Any
Parameters:: belief_id (int)

get_parent_id(node_id)[source]

Return type:: Optional[int]
Parameters:: node_id (int)

get_q_value(action_id)[source]

Return type:: float
Parameters:: action_id (int)

get_v_value(belief_id)[source]

Return type:: float
Parameters:: belief_id (int)

get_visit_count(node_id)[source]

Return type:: int
Parameters:: node_id (int)

get_weight(belief_id)[source]

Return type:: float
Parameters:: belief_id (int)

increment_visit_count(node_id)[source]

Increment visit_count[node_id] by one. Works for either node kind.

Return type:: None
Parameters:: node_id (int)

increment_weight(child_id, delta)[source]

Increment weight[child_id] by delta and patch the parent’s CDF.

Cheaper than recompute_children_cdf() for the common case of a single weight bump (e.g. POMCPOW’s observation widening): walks the parent’s CDF only from this child’s position onward, adding delta to each remaining entry. O(K - position) where K is the number of siblings.

If child_id is the root (no parent), only the weight is updated; there is no parent CDF to patch.

Return type:

None

Parameters:

child_id (int)
delta (float)

print(node_id=0)[source]

Print the subtree rooted at node_id as an indented tree.

Return type:: None
Parameters:: node_id (int)

recompute_children_cdf(parent_id)[source]

Rebuild the parent’s CDF from current weight values.

Use this when a caller has mutated weight[id] for one or more children after they were added.

Return type:: None
Parameters:: parent_id (int)

reserve(capacity)[source]

Pre-allocate capacity slots in every per-node column.

After this call len(tree) is unchanged, but capacity slots are physically resident in each column so the first capacity subsequent allocations write at a cursor instead of triggering list.append (and the periodic O(N) realloc bursts that come with it). Useful when the maximum tree size is known in advance (e.g., 2 * n_simulations * depth for POMCPOW-style planners).

Allocations beyond capacity fall back to append and grow the columns normally. Calling reserve() more than once is idempotent in the sense that columns are re-padded up to capacity if more headroom is needed; existing entries are preserved.

Return type:: None
Parameters:: capacity (int)

sample_belief_child(action_id)[source]

Sample one belief child of action_id proportional to its weight.

Uses the maintained CDF: O(log K) per sample via bisect. If the caller has mutated weight[id] after the children were added, call recompute_children_cdf() first to refresh the CDF.

Return type:: int
Parameters:: action_id (int)

set_immediate_cost(node_id, value)[source]

Set immediate_cost and mirror to immediate_reward = -value.

Return type:

None

Parameters:

node_id (int)
value (float | None)

set_immediate_reward(node_id, value)[source]

Set immediate_reward and mirror to immediate_cost = -value.

Return type:

None

Parameters:

node_id (int)
value (float | None)

update_action_q_with_return(action_id, return_sample)[source]

Standard MCTS Q-update: bump visits and incrementally average return_sample.

Return type:

None

Parameters:

action_id (int)
return_sample (float)

update_belief(belief_id, action, observation, pomdp, **kwargs)[source]

Replace the belief at belief_id with the result of belief.update.

Return type:

None

Parameters:

belief_id (int)
action (Any)
observation (Any)
pomdp (Environment)
kwargs (Any)