POMDPPlanners.core.tree.vectorized_belief_tree package

Vectorized belief tree for GPU-batched online POMDP planning.

This package provides VectorizedBeliefTree, a reusable belief-tree data structure stored entirely in flat PyTorch tensors. It mirrors the tensor tree layout used by GPU-vectorized planners (VOPP / PORPP) but contains no planning-algorithm logic — only the batched structural and statistical primitives (keyed child insertion, composite-key lookup, scatter aggregation, group-by-parent, depth-wise traversal) that such algorithms are built from.

See POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree for the design.

class POMDPPlanners.core.tree.vectorized_belief_tree.BeliefStore(*args, **kwargs)[source]

Bases: Protocol

Interface for an external store of per-node belief representations.

Implementations map belief-tree node indices to concrete belief objects. They are entirely optional: the tree is fully functional without one.

create_root(belief)[source]

Register the root belief and return its node index (expected 0).

Parameters:: belief (object) – The concrete root belief object to store.
Return type:: int
Returns:: The integer node index the belief was stored under.

create_successors(parent_belief_indices, action_keys, observation_keys)[source]

Create and store successor beliefs for a batch of transitions.

Parameters:

parent_belief_indices (Tensor) – [batch] parent belief node indices.
action_keys (Tensor) – [batch] integer action keys taken.
observation_keys (Tensor) – [batch] integer observation keys received.

Return type:

Tensor

Returns:

[batch] successor belief node indices, aligned with the inputs.

class POMDPPlanners.core.tree.vectorized_belief_tree.FieldRegistry(device, capacity)[source]

Bases: object

Owns the backing tensors for the registered fields of one node kind.

The registry preallocates one tensor per field with the tree’s current node capacity and keeps every field tensor on the tree’s device. Growing the registry reallocates each field tensor geometrically while preserving existing rows.

Parameters:

device (torch.device)
capacity (int)

property device: torch.device: Device shared by every field tensor in the registry.

grow(new_capacity)[source]

Reallocate every field tensor to new_capacity rows, preserving data.

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

names()[source]

Return the registered field names in insertion order.

Return type:: Tuple[str, ...]

register(spec, reserved_names)[source]

Register a new field, allocating its backing tensor.

Parameters:

spec (FieldSpec) – The field description.
reserved_names (frozenset) – Built-in column names the field may not shadow.

Raises:

ValueError – If the name collides with a reserved name or an already-registered field.

Return type:

None

reset_rows(start, end)[source]

Reset rows [start, end) of every field tensor to its default.

Return type:

None

Parameters:

start (int)
end (int)

specs()[source]

Return the registered FieldSpec objects in insertion order.

Return type:: Tuple[FieldSpec, ...]

tensor(name)[source]

Return the full backing tensor for a registered field.

Raises:: KeyError – If no field with name is registered.
Return type:: Tensor
Parameters:: name (str)

to(device)[source]

Move every field tensor to device in place.

Return type:: None
Parameters:: device (torch.device)

class POMDPPlanners.core.tree.vectorized_belief_tree.FieldSpec(name, shape, dtype, default)[source]

Bases: object

Immutable description of one registered per-node field.

Parameters:

name (str)
shape (Tuple[int, ...])
dtype (torch.dtype)
default (float | int | bool)

name: Field name, unique within its node kind and distinct from the built-in column names.

shape: Trailing shape of the field for a single node. () denotes a scalar field; (num_actions,) denotes a per-action vector, etc.

dtype: The tensor dtype used to store the field.

default: Value used to initialise newly allocated rows.

default: float | int | bool

dtype: torch.dtype

name: str

shape: Tuple[int, ...]

class POMDPPlanners.core.tree.vectorized_belief_tree.VectorizedBeliefTree(device=None, belief_capacity=1024, action_capacity=1024, index_dtype=torch.int64, value_dtype=torch.float32, growth_factor=2.0)[source]

Bases: object

Flat-tensor belief tree with batched insertion, lookup, and reduction.

All persistent tree data is stored in preallocated tensors on self.device. Belief-node and action-node columns are stored separately. Insertion, lookup, and statistic accumulation are batched and free of Python per-row loops so the structure runs on a GPU without per-simulation synchronization.

Parameters:

device (torch.device | None)
belief_capacity (int)
action_capacity (int)
index_dtype (torch.dtype)
value_dtype (torch.dtype)
growth_factor (float)

device: The device every persistent tensor lives on.

index_dtype: Integer dtype for indices and keys (default int64).

value_dtype: Floating dtype for real-valued statistics (default float32).

growth_factor: Geometric factor used when a capacity is exceeded.

Example

See the module-level docstring for a runnable example.

property action_capacity: int: Current preallocated action-node capacity.

action_children(belief_indices)[source]

Return the action children of the queried belief nodes in CSR form.

Parameters:: belief_indices (Tensor) – [M] parent belief node indices.
Return type:: Tuple[Tensor, Tensor]
Returns:: (flat_child_indices, offsets) where the children of belief_indices[m] are flat_child_indices[offsets[m]:offsets[m + 1]].

action_field(name)[source]

Return the active-rows view of a registered action field.

The returned view spans [0, num_action_nodes) and is valid until the next structural mutation (which may reallocate the backing tensor).

Return type:: Tensor
Parameters:: name (str)

action_nodes_at_depth(depth)[source]

Return the indices of every active action node at depth.

Return type:: Tensor
Parameters:: depth (int)

add_action_rewards(action_indices, rewards)[source]

Add rewards to action nodes’ reward sums, summing duplicate indices.

Parameters:

action_indices (Tensor) – [batch] action node indices (may repeat).
rewards (Tensor) – [batch] real-valued rewards.

Return type:

None

property belief_capacity: int: Current preallocated belief-node capacity.

belief_children(action_indices)[source]

Return the belief children of the queried action nodes in CSR form.

Parameters:: action_indices (Tensor) – [M] parent action node indices.
Return type:: Tuple[Tensor, Tensor]
Returns:: (flat_child_indices, offsets) where the children of action_indices[m] are flat_child_indices[offsets[m]:offsets[m + 1]].

belief_field(name)[source]

Return the active-rows view of a registered belief field.

The returned view spans [0, num_belief_nodes) and is valid until the next structural mutation (which may reallocate the backing tensor).

Return type:: Tensor
Parameters:: name (str)

belief_nodes_at_depth(depth)[source]

Return the indices of every active belief node at depth.

Return type:: Tensor
Parameters:: depth (int)

clear()[source]

Remove every node except the root, preserving capacity and fields.

Registered field definitions survive; their rows reset to defaults.

Return type:: None

find_actions(parent_belief_indices, action_keys)[source]

Look up action nodes by (parent_belief, action_key) without mutating.

Parameters:

parent_belief_indices (Tensor) – [batch] parent belief indices.
action_keys (Tensor) – [batch] integer action keys.

Return type:

Tuple[Tensor, Tensor]

Returns:

(action_indices, found_mask). Missing entries hold -1 in action_indices and False in found_mask. Order follows the input batch.

find_beliefs(parent_action_indices, observation_keys)[source]

Look up belief nodes by (parent_action, observation_key).

Parameters:

parent_action_indices (Tensor) – [batch] parent action indices.
observation_keys (Tensor) – [batch] integer observation keys.

Return type:

Tuple[Tensor, Tensor]

Returns:

(belief_indices, found_mask) with -1 for missing entries, in input order.

get_or_create_actions(parent_belief_indices, action_keys)[source]

Return one action-node index per input row, creating missing nodes.

Duplicate (parent_belief, action_key) pairs — whether already in the tree or repeated within the batch — resolve to a single node.

Parameters:

parent_belief_indices (Tensor) – [batch] parent belief indices.
action_keys (Tensor) – [batch] integer action keys.

Return type:

Tuple[Tensor, Tensor]

Returns:

(action_indices, created_mask) in input order. created_mask is True for rows whose node did not exist before this call.

get_or_create_beliefs(parent_action_indices, observation_keys, *, terminal_mask=None)[source]

Return one successor-belief index per input row, creating missing nodes.

Duplicate (parent_action, observation_key) pairs resolve to a single node. A new belief’s depth is its parent action’s depth plus one.

Terminal handling: if terminal_mask is supplied, terminal flags are combined with logical OR — both when several batch rows map to the same new node and when a row targets a belief that already exists (an existing False flag can be promoted to True, never the reverse).

Parameters:

parent_action_indices (Tensor) – [batch] parent action indices.
observation_keys (Tensor) – [batch] integer observation keys.
terminal_mask (Optional[Tensor]) – Optional [batch] boolean terminal flags.

Return type:

Tuple[Tensor, Tensor]

Returns:

(belief_indices, created_mask) in input order.

increment_action_visits(action_indices, counts=None)[source]

Add visit counts to action nodes, summing duplicate indices.

Parameters:

action_indices (Tensor) – [batch] action node indices (may repeat).
counts (Optional[Tensor]) – Optional [batch] integer increments; defaults to one.

Return type:

None

load_state_dict(state)[source]

Restore the tree from a state_dict() payload.

Return type:: None
Parameters:: state (Dict[str, object])

property num_action_nodes: int: Number of active action nodes.

property num_belief_nodes: int: Number of active belief nodes.

parent_actions(belief_indices)[source]

Return the parent action index of each queried belief node.

The root belief maps to -1.

Return type:: Tensor
Parameters:: belief_indices (torch.Tensor)

parent_beliefs(action_indices)[source]

Return the parent belief index of each queried action node.

Return type:: Tensor
Parameters:: action_indices (torch.Tensor)

reduce_belief_children(parent_action_indices, child_values, *, reduction)[source]

Reduce child values grouped by parent action node.

This is a generic segmented reduction — it does not implement any algorithm-specific backup. parent_action_indices[k] is the parent action node of the k-th child and child_values[k] its value.

Parameters:

parent_action_indices (Tensor) – [K] parent action indices per child.
child_values (Tensor) – [K] per-child values.
reduction (str) – One of "sum", "mean", "max", "min".

Return type:

Tensor

Returns:

[num_action_nodes] reduced value per action node. Action nodes with no contributing child hold 0.

register_action_field(name, shape=(), *, dtype=None, default=0)[source]

Register an extra per-action-node field backed by its own tensor.

Parameters:

name (str) – Field name; must not shadow a built-in action column.
shape (Tuple[int, ...]) – Trailing per-node shape (() for scalar).
dtype (Optional[dtype]) – Storage dtype; defaults to value_dtype.
default (float) – Value used to initialise new rows.

Return type:

None

register_belief_field(name, shape=(), *, dtype=None, default=0)[source]

Register an extra per-belief-node field backed by its own tensor.

Parameters:

name (str) – Field name; must not shadow a built-in belief column.
shape (Tuple[int, ...]) – Trailing per-node shape (() for scalar).
dtype (Optional[dtype]) – Storage dtype; defaults to value_dtype.
default (float) – Value used to initialise new rows.

Return type:

None

reset_statistics()[source]

Reset visit counts, reward sums, and registered fields; keep topology.

Return type:: None

property root_index: int: Index of the root belief node (always 0).

state_dict()[source]

Serialise the active tree (topology, statistics, and fields).

Return type:: Dict[str, object]
Returns:: A picklable dict holding the active tensor rows, field definitions, node counters, and reconstruction configuration.

to(device)[source]

Move every built-in and registered tensor to device in place.

Return type:: VectorizedBeliefTree
Returns:: self, for chaining.
Parameters:: device (torch.device)

update_action_statistics(action_indices, rewards, visit_weights=None)[source]

Accumulate reward sums and visit counts for action nodes in one call.

Parameters:

action_indices (Tensor) – [batch] action node indices (may repeat).
rewards (Tensor) – [batch] real-valued rewards.
visit_weights (Optional[Tensor]) – Optional [batch] integer visit increments; defaults to one per row.

Return type:

None

validate()[source]

Assert the structural invariants of the tree.

Raises:: AssertionError – If any invariant is violated.
Return type:: None

Submodules

POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_store module

Protocol separating belief content from belief-tree topology.

VectorizedBeliefTree stores only the topology, integer keys, and statistics of a POMDP belief tree. The mathematical belief attached to each belief node (a particle set, a categorical vector, a Gaussian, an action-observation history, a latent vector, …) lives outside the tree in an object implementing BeliefStore.

The contract is that a belief-tree node index is the identifier the external store uses to look up the belief for that node:

VectorizedBeliefTree stores tree topology and statistics.
BeliefStore          stores or reconstructs the mathematical beliefs.

The tree never calls a BeliefStore; the protocol exists so planners can keep a parallel store keyed by the same integer node indices the tree hands out.

class POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_store.BeliefStore(*args, **kwargs)[source]

Bases: Protocol

Interface for an external store of per-node belief representations.

Implementations map belief-tree node indices to concrete belief objects. They are entirely optional: the tree is fully functional without one.

create_root(belief)[source]

Register the root belief and return its node index (expected 0).

Parameters:: belief (object) – The concrete root belief object to store.
Return type:: int
Returns:: The integer node index the belief was stored under.

create_successors(parent_belief_indices, action_keys, observation_keys)[source]

Create and store successor beliefs for a batch of transitions.

Parameters:

parent_belief_indices (Tensor) – [batch] parent belief node indices.
action_keys (Tensor) – [batch] integer action keys taken.
observation_keys (Tensor) – [batch] integer observation keys received.

Return type:

Tensor

Returns:

[batch] successor belief node indices, aligned with the inputs.

POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree module

Vectorized belief tree stored in flat PyTorch tensors.

This module implements VectorizedBeliefTree, a reusable belief-tree data structure for online POMDP planning. It captures the tensor tree layout of GPU-vectorized planners (e.g. VOPP / PORPP): rather than a graph of Python node objects, the whole tree is a small set of flat, preallocated tensors and a node is an integer index into every column.

The tree alternates two node kinds:

belief node --a--> action node --o--> belief node

An action node is uniquely identified by (parent_belief_index, action_key) and a successor belief node by (parent_action_index, observation_key). Both keys are integers; converting continuous actions/observations into integer keys (binning, hashing, nearest-neighbour, progressive widening) is the caller’s responsibility.

The class stores only topology, integer keys, and statistics. The concrete belief attached to a belief node lives in an external BeliefStore keyed by the integer node index. No planning-policy logic (UCB, softmax action selection, log-sum-exp value backups, PORPP preference updates) lives here — only the batched structural and statistical primitives such algorithms are built from.

Example

Basic usage on CPU or CUDA:

>>> import torch
>>> from POMDPPlanners.core.tree.vectorized_belief_tree import (
...     VectorizedBeliefTree,
... )
>>> tree = VectorizedBeliefTree(device=torch.device("cpu"))
>>> root = tree.root_index
>>> parents = torch.tensor([root, root, root, root])
>>> actions = torch.tensor([0, 1, 1, 2])
>>> action_nodes, created = tree.get_or_create_actions(parents, actions)
>>> bool(action_nodes[1] == action_nodes[2])  # duplicate pair -> one node
True
>>> tree.update_action_statistics(action_nodes, torch.tensor([1.0, 2.0, 3.0, -1.0]))
>>> observations = torch.tensor([4, 2, 2, 7])
>>> beliefs, belief_created = tree.get_or_create_beliefs(action_nodes, observations)
>>> tree.validate()

class POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree.VectorizedBeliefTree(device=None, belief_capacity=1024, action_capacity=1024, index_dtype=torch.int64, value_dtype=torch.float32, growth_factor=2.0)[source]

Bases: object

Flat-tensor belief tree with batched insertion, lookup, and reduction.

All persistent tree data is stored in preallocated tensors on self.device. Belief-node and action-node columns are stored separately. Insertion, lookup, and statistic accumulation are batched and free of Python per-row loops so the structure runs on a GPU without per-simulation synchronization.

Parameters:

device (torch.device | None)
belief_capacity (int)
action_capacity (int)
index_dtype (torch.dtype)
value_dtype (torch.dtype)
growth_factor (float)

device: The device every persistent tensor lives on.

index_dtype: Integer dtype for indices and keys (default int64).

value_dtype: Floating dtype for real-valued statistics (default float32).

growth_factor: Geometric factor used when a capacity is exceeded.

Example

See the module-level docstring for a runnable example.

property action_capacity: int: Current preallocated action-node capacity.

action_children(belief_indices)[source]

Return the action children of the queried belief nodes in CSR form.

Parameters:: belief_indices (Tensor) – [M] parent belief node indices.
Return type:: Tuple[Tensor, Tensor]
Returns:: (flat_child_indices, offsets) where the children of belief_indices[m] are flat_child_indices[offsets[m]:offsets[m + 1]].

action_field(name)[source]

Return the active-rows view of a registered action field.

The returned view spans [0, num_action_nodes) and is valid until the next structural mutation (which may reallocate the backing tensor).

Return type:: Tensor
Parameters:: name (str)

action_nodes_at_depth(depth)[source]

Return the indices of every active action node at depth.

Return type:: Tensor
Parameters:: depth (int)

add_action_rewards(action_indices, rewards)[source]

Add rewards to action nodes’ reward sums, summing duplicate indices.

Parameters:

action_indices (Tensor) – [batch] action node indices (may repeat).
rewards (Tensor) – [batch] real-valued rewards.

Return type:

None

property belief_capacity: int: Current preallocated belief-node capacity.

belief_children(action_indices)[source]

Return the belief children of the queried action nodes in CSR form.

Parameters:: action_indices (Tensor) – [M] parent action node indices.
Return type:: Tuple[Tensor, Tensor]
Returns:: (flat_child_indices, offsets) where the children of action_indices[m] are flat_child_indices[offsets[m]:offsets[m + 1]].

belief_field(name)[source]

Return the active-rows view of a registered belief field.

The returned view spans [0, num_belief_nodes) and is valid until the next structural mutation (which may reallocate the backing tensor).

Return type:: Tensor
Parameters:: name (str)

belief_nodes_at_depth(depth)[source]

Return the indices of every active belief node at depth.

Return type:: Tensor
Parameters:: depth (int)

clear()[source]

Remove every node except the root, preserving capacity and fields.

Registered field definitions survive; their rows reset to defaults.

Return type:: None

find_actions(parent_belief_indices, action_keys)[source]

Look up action nodes by (parent_belief, action_key) without mutating.

Parameters:

parent_belief_indices (Tensor) – [batch] parent belief indices.
action_keys (Tensor) – [batch] integer action keys.

Return type:

Tuple[Tensor, Tensor]

Returns:

(action_indices, found_mask). Missing entries hold -1 in action_indices and False in found_mask. Order follows the input batch.

find_beliefs(parent_action_indices, observation_keys)[source]

Look up belief nodes by (parent_action, observation_key).

Parameters:

parent_action_indices (Tensor) – [batch] parent action indices.
observation_keys (Tensor) – [batch] integer observation keys.

Return type:

Tuple[Tensor, Tensor]

Returns:

(belief_indices, found_mask) with -1 for missing entries, in input order.

get_or_create_actions(parent_belief_indices, action_keys)[source]

Return one action-node index per input row, creating missing nodes.

Duplicate (parent_belief, action_key) pairs — whether already in the tree or repeated within the batch — resolve to a single node.

Parameters:

parent_belief_indices (Tensor) – [batch] parent belief indices.
action_keys (Tensor) – [batch] integer action keys.

Return type:

Tuple[Tensor, Tensor]

Returns:

(action_indices, created_mask) in input order. created_mask is True for rows whose node did not exist before this call.

get_or_create_beliefs(parent_action_indices, observation_keys, *, terminal_mask=None)[source]

Return one successor-belief index per input row, creating missing nodes.

Duplicate (parent_action, observation_key) pairs resolve to a single node. A new belief’s depth is its parent action’s depth plus one.

Terminal handling: if terminal_mask is supplied, terminal flags are combined with logical OR — both when several batch rows map to the same new node and when a row targets a belief that already exists (an existing False flag can be promoted to True, never the reverse).

Parameters:

parent_action_indices (Tensor) – [batch] parent action indices.
observation_keys (Tensor) – [batch] integer observation keys.
terminal_mask (Optional[Tensor]) – Optional [batch] boolean terminal flags.

Return type:

Tuple[Tensor, Tensor]

Returns:

(belief_indices, created_mask) in input order.

increment_action_visits(action_indices, counts=None)[source]

Add visit counts to action nodes, summing duplicate indices.

Parameters:

action_indices (Tensor) – [batch] action node indices (may repeat).
counts (Optional[Tensor]) – Optional [batch] integer increments; defaults to one.

Return type:

None

load_state_dict(state)[source]

Restore the tree from a state_dict() payload.

Return type:: None
Parameters:: state (Dict[str, object])

property num_action_nodes: int: Number of active action nodes.

property num_belief_nodes: int: Number of active belief nodes.

parent_actions(belief_indices)[source]

Return the parent action index of each queried belief node.

The root belief maps to -1.

Return type:: Tensor
Parameters:: belief_indices (torch.Tensor)

parent_beliefs(action_indices)[source]

Return the parent belief index of each queried action node.

Return type:: Tensor
Parameters:: action_indices (torch.Tensor)

reduce_belief_children(parent_action_indices, child_values, *, reduction)[source]

Reduce child values grouped by parent action node.

This is a generic segmented reduction — it does not implement any algorithm-specific backup. parent_action_indices[k] is the parent action node of the k-th child and child_values[k] its value.

Parameters:

parent_action_indices (Tensor) – [K] parent action indices per child.
child_values (Tensor) – [K] per-child values.
reduction (str) – One of "sum", "mean", "max", "min".

Return type:

Tensor

Returns:

[num_action_nodes] reduced value per action node. Action nodes with no contributing child hold 0.

register_action_field(name, shape=(), *, dtype=None, default=0)[source]

Register an extra per-action-node field backed by its own tensor.

Parameters:

name (str) – Field name; must not shadow a built-in action column.
shape (Tuple[int, ...]) – Trailing per-node shape (() for scalar).
dtype (Optional[dtype]) – Storage dtype; defaults to value_dtype.
default (float) – Value used to initialise new rows.

Return type:

None

register_belief_field(name, shape=(), *, dtype=None, default=0)[source]

Register an extra per-belief-node field backed by its own tensor.

Parameters:

name (str) – Field name; must not shadow a built-in belief column.
shape (Tuple[int, ...]) – Trailing per-node shape (() for scalar).
dtype (Optional[dtype]) – Storage dtype; defaults to value_dtype.
default (float) – Value used to initialise new rows.

Return type:

None

reset_statistics()[source]

Reset visit counts, reward sums, and registered fields; keep topology.

Return type:: None

property root_index: int: Index of the root belief node (always 0).

state_dict()[source]

Serialise the active tree (topology, statistics, and fields).

Return type:: Dict[str, object]
Returns:: A picklable dict holding the active tensor rows, field definitions, node counters, and reconstruction configuration.

to(device)[source]

Move every built-in and registered tensor to device in place.

Return type:: VectorizedBeliefTree
Returns:: self, for chaining.
Parameters:: device (torch.device)

update_action_statistics(action_indices, rewards, visit_weights=None)[source]

Accumulate reward sums and visit counts for action nodes in one call.

Parameters:

action_indices (Tensor) – [batch] action node indices (may repeat).
rewards (Tensor) – [batch] real-valued rewards.
visit_weights (Optional[Tensor]) – Optional [batch] integer visit increments; defaults to one per row.

Return type:

None

validate()[source]

Assert the structural invariants of the tree.

Raises:: AssertionError – If any invariant is violated.
Return type:: None

POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree_fields module

Registered-field bookkeeping for the vectorized belief tree.

A registered field is an extra per-node column (belief-side or action-side) that a planning algorithm attaches to the tree without modifying the core storage schema. For example, a PORPP/VOPP-style planner registers a preferences belief field of shape (num_actions,) and a value belief field; a POMCP-style planner registers a q_value action field.

Each field is described by an immutable FieldSpec and backed by a preallocated tensor that a FieldRegistry grows in lockstep with the node capacity of the owning tree.

class POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree_fields.FieldRegistry(device, capacity)[source]

Bases: object

Owns the backing tensors for the registered fields of one node kind.

The registry preallocates one tensor per field with the tree’s current node capacity and keeps every field tensor on the tree’s device. Growing the registry reallocates each field tensor geometrically while preserving existing rows.

Parameters:

device (torch.device)
capacity (int)

property device: torch.device: Device shared by every field tensor in the registry.

grow(new_capacity)[source]

Reallocate every field tensor to new_capacity rows, preserving data.

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

names()[source]

Return the registered field names in insertion order.

Return type:: Tuple[str, ...]

register(spec, reserved_names)[source]

Register a new field, allocating its backing tensor.

Parameters:

spec (FieldSpec) – The field description.
reserved_names (frozenset) – Built-in column names the field may not shadow.

Raises:

ValueError – If the name collides with a reserved name or an already-registered field.

Return type:

None

reset_rows(start, end)[source]

Reset rows [start, end) of every field tensor to its default.

Return type:

None

Parameters:

start (int)
end (int)

specs()[source]

Return the registered FieldSpec objects in insertion order.

Return type:: Tuple[FieldSpec, ...]

tensor(name)[source]

Return the full backing tensor for a registered field.

Raises:: KeyError – If no field with name is registered.
Return type:: Tensor
Parameters:: name (str)

to(device)[source]

Move every field tensor to device in place.

Return type:: None
Parameters:: device (torch.device)

class POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree_fields.FieldSpec(name, shape, dtype, default)[source]

Bases: object

Immutable description of one registered per-node field.

Parameters:

name (str)
shape (Tuple[int, ...])
dtype (torch.dtype)
default (float | int | bool)

name: Field name, unique within its node kind and distinct from the built-in column names.

shape: Trailing shape of the field for a single node. () denotes a scalar field; (num_actions,) denotes a per-action vector, etc.

dtype: The tensor dtype used to store the field.

default: Value used to initialise newly allocated rows.

default: float | int | bool

dtype: torch.dtype

name: str

shape: Tuple[int, ...]

POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree_lookup module

Vectorized composite-key lookup and segmented-grouping helpers.

These are pure tensor functions used by VectorizedBeliefTree. None of them mutate persistent tree state; each operates on the flat column tensors passed in and returns index or value tensors. Every function is batched — there are no Python loops over individual rows.

The two structural keys of a belief tree are composite integer pairs:

(parent_belief_index, action_key) identifying an action node;
(parent_action_index, observation_key) identifying a belief node.

Matching them is done with match_pairs(), which packs a pair-match into a single torch.unique() call rather than an unsafe integer packing that could overflow int64 when either component is large.

Complexity notes use E for the number of existing rows, Q for the query batch size, A for the number of active nodes, and M for the number of queried parents.

POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree_lookup.csr_children(parent_column, query_parents)[source]

Group node indices by parent into a CSR (segmented) representation.

Parameters:

parent_column (Tensor) – [A] parent id of every active node of one kind. Position a holds the parent id of node a.
query_parents (Tensor) – [M] parent ids whose children are requested.

Return type:

Tuple[Tensor, Tensor]

Returns:

A tuple (flat_children, offsets) where offsets has length M + 1 and the children of query_parents[m] are flat_children[offsets[m]:offsets[m + 1]].

Complexity:: O(A log A + M log A).

POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree_lookup.match_pairs(existing_first, existing_second, query_first, query_second)[source]

Match each query pair against the existing pairs.

Parameters:

existing_first (Tensor) – [E] first component of the existing pairs.
existing_second (Tensor) – [E] second component of the existing pairs.
query_first (Tensor) – [Q] first component of the query pairs.
query_second (Tensor) – [Q] second component of the query pairs.

Return type:

Tensor

Returns:

[Q] tensor whose i-th entry is the index into the existing pairs that matches query i, or -1 if no existing pair matches. Order follows the query batch.

In belief-tree terms each pair is a composite node key: for action-node lookup the components are (parent_belief_index, action_key); for belief-node lookup they are (parent_action_index, observation_key). The existing_* tensors are the corresponding column tensors of the nodes already in the tree and the query_* tensors are the keys being searched; a returned index is the matching child node, -1 meaning that edge does not exist yet.

Complexity:: O((E + Q) log (E + Q)) from a single lexicographic torch.unique.

Example

Three existing action nodes keyed by (parent_belief, action_key), looked up by three query pairs (matching on the pair, not either component alone):

>>> import torch
>>> existing_first = torch.tensor([0, 0, 1])   # parent belief index
>>> existing_second = torch.tensor([2, 4, 2])  # action key
>>> query_first = torch.tensor([0, 1, 0])
>>> query_second = torch.tensor([4, 2, 9])
>>> match_pairs(existing_first, existing_second, query_first, query_second)
tensor([ 1,  2, -1])

Query (0, 4) matches existing row 1, (1, 2) matches row 2, and (0, 9) has no match so it returns -1.

POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree_lookup.scatter_or(indices, values, size)[source]

Reduce boolean values by logical OR, grouped by indices.

Parameters:

indices (Tensor) – [N] target index of every value (may contain duplicates).
values (Tensor) – [N] boolean values to OR into their target positions.
size (int) – Length of the output tensor.

Return type:

Tensor

Returns:

[size] boolean tensor; position p is True iff any value routed to p is True. Untargeted positions are False.

POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree_lookup.segment_reduce(segment_ids, values, num_segments, reduction)[source]

Reduce values grouped by segment_ids into a dense tensor.

Parameters:

segment_ids (Tensor) – [N] segment index of every value.
values (Tensor) – [N] values to reduce.
num_segments (int) – Length of the output tensor.
reduction (str) – One of "sum", "mean", "max", "min".

Return type:

Tensor

Returns:

[num_segments] reduced values. Empty segments are 0 for every reduction (documented default so the result is always well defined).

Raises:

ValueError – If reduction is not a supported name.

POMDPPlanners.core.tree.vectorized_belief_tree.vectorized_belief_tree_lookup.unique_pairs(first, second)[source]

Deduplicate a batch of integer pairs.

Parameters:

first (Tensor) – [Q] first component of the pairs.
second (Tensor) – [Q] second component of the pairs.

Return type:

Tuple[Tensor, Tensor, Tensor]

Returns:

A tuple (unique_first, unique_second, inverse) where the unique arrays hold each distinct pair once and inverse maps every input row to its unique-pair index (inverse has length Q).

Complexity:: O(Q log Q).