POMDPPlanners.planners.mcts_planners.constrained_zero package

ConstrainedZero: Neural MCTS for Chance-Constrained POMDPs.

This package implements the ConstrainedZero algorithm (Moss et al., IJCAI 2024), which extends BetaZero to solve CC-POMDPs by adding a failure probability head, safety-constrained PUCT, and adaptive failure threshold calibration.

References

Moss, R. J., Jamgochian, A., Fischer, J., Corso, A., & Kochenderfer, M. J. (2024). ConstrainedZero: Chance-Constrained POMDP Planning Using Learned Probabilistic Failure Surrogates and Adaptive Safety Constraints. Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence (IJCAI), 6752-6760. https://www.ijcai.org/proceedings/2024/746

Classes:

ConstrainedZero: Main planner extending BetaZero for CC-POMDPs ConstrainedZeroNetwork: Three-head network with policy, value, and failure heads ConstrainedTrainingBuffer: Replay buffer with failure targets ConstrainedTrainingExample: Training datum with failure target

class POMDPPlanners.planners.mcts_planners.constrained_zero.ConstrainedTrainingBuffer(n_buffer=1)[source]

Bases: TrainingBuffer

Iteration-slot replay buffer for ConstrainedZero training.

Extends TrainingBuffer to store and sample ConstrainedTrainingExample instances, returning a 4-tuple from sample_batch() that includes failure targets.

Parameters:

n_buffer (int) – Number of policy-iteration slots to retain. With the default n_buffer=1 only the current iteration’s data is used for training (on-policy).

Example

>>> import numpy as np
>>> from POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_training_buffer import (
...     ConstrainedTrainingBuffer, ConstrainedTrainingExample,
... )
>>> buf = ConstrainedTrainingBuffer(n_buffer=1)
>>> buf.begin_iteration()
>>> buf.add(ConstrainedTrainingExample(np.zeros(4), np.array([0.5, 0.5]), 1.0, 0.0))
>>> len(buf)
1
>>> batch = buf.sample_batch(1)
>>> len(batch)
4
>>> batch[0].shape
(1, 4)
add(example)[source]

Append a constrained example to the current iteration slot.

Return type:

None

Parameters:

example (ConstrainedTrainingExample)

sample_batch(batch_size)[source]

Sample a random mini-batch including failure targets.

Parameters:

batch_size (int) – Number of examples to sample (with replacement if batch_size > len(buffer)).

Return type:

Tuple[ndarray, ndarray, ndarray, ndarray]

Returns:

Tuple of (belief_features, policy_targets, value_targets, failure_targets). - belief_features: shape (batch_size, belief_dim) - policy_targets: shape (batch_size, policy_dim) - value_targets: shape (batch_size,) - failure_targets: shape (batch_size,)

class POMDPPlanners.planners.mcts_planners.constrained_zero.ConstrainedTrainingExample(belief_features, policy_target, value_target, failure_target)[source]

Bases: object

Single training datum for ConstrainedZero network training.

Parameters:
belief_features

Belief feature vector phi(b), shape (belief_dim,).

policy_target

Q-weighted policy target pi_qw, shape (n_actions,) (discrete).

value_target

Discounted return g_t.

failure_target

Binary episode-level failure indicator (1.0 if failure occurred).

belief_features: ndarray
failure_target: float
policy_target: ndarray
value_target: float
class POMDPPlanners.planners.mcts_planners.constrained_zero.ConstrainedZero(environment, discount_factor, depth, name, action_sampler, failure_fn, delta_0=0.01, eta=1e-05, delta_compounding=1.0, k_a=1.0, alpha_a=0.5, k_o=1.0, alpha_o=0.5, exploration_constant=1.0, time_out_in_seconds=None, n_simulations=None, min_visit_count_per_action=1, network=None, belief_representation=None, state_dim=None, z_q=1.0, z_n=1.0, temperature=1.0, n_buffer=1, training_batch_size=256, training_epochs=10, learning_rate=0.001, weight_decay=0.0001, hidden_sizes=(128, 128), use_dropout=True, p_dropout=0.2, track_gradients=False, normalize_inputs=True, normalize_values=True, log_path=None, debug=False, use_queue_logger=False)[source]

Bases: BetaZero

ConstrainedZero: Neural MCTS for Chance-Constrained POMDPs.

Extends BetaZero with:

  1. 3-head network: Adds a failure probability head alongside policy and value.

  2. SPUCT selection: Safety-constrained PUCT that masks unsafe actions.

  3. Adaptive Delta (conformal inference): Calibrates the failure threshold during tree search using online conformal inference.

  4. Failure propagation: Tracks failure probability per action node using p = p_immediate + delta_compounding * (1 - p_immediate) * p_next.

  5. Constrained policy targets: Applies safety mask during target computation.

Parameters:
failure_fn

User-provided function state -> bool defining failure.

delta_0

Nominal failure probability threshold.

eta

Learning rate for adaptive Delta calibration.

delta_compounding

Discount factor for failure propagation.

Example

>>> import numpy as np
>>> np.random.seed(42)
>>> from POMDPPlanners.environments.tiger_pomdp import TigerPOMDP
>>> from POMDPPlanners.core.belief import get_initial_belief
>>> from POMDPPlanners.utils.action_samplers import DiscreteActionSampler
>>> from POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_zero import ConstrainedZero
>>>
>>> env = TigerPOMDP(discount_factor=0.95)
>>> sampler = DiscreteActionSampler(env.get_actions())
>>> planner = ConstrainedZero(
...     environment=env,
...     discount_factor=0.95,
...     depth=3,
...     name="CZ_Tiger",
...     action_sampler=sampler,
...     n_simulations=20,
...     state_dim=1,
...     failure_fn=lambda s: False,
... )
>>> belief = get_initial_belief(env, n_particles=10)
>>> actions, run_data = planner.action(belief)
>>> actions[0] in env.get_actions()
True
get_metric_keys()[source]

Return the loss-metric key names produced by train_step().

Return type:

List[str]

network: ConstrainedZeroNetwork
class POMDPPlanners.planners.mcts_planners.constrained_zero.ConstrainedZeroNetwork(belief_dim, action_space_type, n_actions=None, action_dim=None, hidden_sizes=(128, 128), use_dropout=True, p_dropout=0.2)[source]

Bases: BetaZeroNetwork

Three-head neural network for ConstrainedZero.

Architecture:

  • Shared trunk: Linear(belief_dim, h) → ReLU → [Dropout] → ... → Linear(h, h) → ReLU → [Dropout]

  • Policy head: inherited from BetaZeroNetwork

  • Value head: inherited from BetaZeroNetwork

  • Failure head: Linear(h, h//2) -> ReLU -> Linear(h//2, 1)

The failure head outputs a raw logit. During predict(), sigmoid is applied to produce a failure probability in [0, 1].

Parameters:

  • belief_dim (int) – Dimensionality of the belief feature vector phi(b).

  • action_space_type (str) – "discrete" or "continuous".

  • n_actions (Optional[int]) – Number of discrete actions (required when action_space_type="discrete").

  • action_dim (Optional[int]) – Dimensionality of continuous actions (required when action_space_type="continuous").

  • hidden_sizes (Sequence[int]) – Tuple of hidden layer widths for the shared trunk.

  • use_dropout (bool) – If True, apply dropout after each ReLU in the shared trunk (default True).

  • p_dropout (float) – Dropout probability for trunk layers (default 0.2).

Example

>>> import numpy as np
>>> from POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_zero_network import ConstrainedZeroNetwork
>>> net = ConstrainedZeroNetwork(belief_dim=4, action_space_type="discrete", n_actions=3, use_dropout=False)
>>> policy, value, failure_prob = net.predict(np.zeros(4, dtype=np.float32))
>>> policy.shape
(3,)
>>> isinstance(value, float)
True
>>> 0.0 <= failure_prob <= 1.0
True
forward(belief_features)[source]

Forward pass returning raw policy, value, and failure logit.

Parameters:

belief_features (Tensor) – Tensor of shape (batch, belief_dim) or (belief_dim,).

Return type:

Tuple[Tensor, Tensor, Tensor]

Returns:

Tuple of (policy_output, value, failure_logit) tensors.

predict(belief_features)[source]

Single-sample inference returning numpy policy, value, and failure probability.

Switches to eval mode before inference to disable dropout, then restores the original training mode.

Parameters:

belief_features (ndarray) – 1-D array of shape (belief_dim,).

Return type:

Tuple[ndarray, float, float]

Returns:

Tuple of (policy, value, failure_prob). - Discrete: policy is a probability vector summing to 1. - Continuous: policy is [mean, log_std]. - value is a Python float. - failure_prob is a Python float in [0, 1].

predict_batch(belief_features_batch)[source]

Batched inference returning numpy policy, value, and failure probability arrays.

Switches to eval mode before inference to disable dropout, then restores the original training mode.

Parameters:

belief_features_batch (ndarray) – 2-D array of shape (N, belief_dim).

Return type:

Tuple[ndarray, ndarray, ndarray]

Returns:

Tuple of (policies, values, failure_probs). - Discrete: policies is (N, n_actions) probability matrix. - Continuous: policies is (N, 2*action_dim) with [mean, log_std]. - values is (N,) array of floats. - failure_probs is (N,) array of floats in [0, 1].

Submodules

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_puct module

Safety-constrained PUCT (SPUCT) action selection for ConstrainedZero.

Implements the SPUCT selection rule that masks unsafe actions based on their estimated failure probability relative to an adaptive threshold Delta’.

References

Moss, R. J., Jamgochian, A., Fischer, J., Corso, A., & Kochenderfer, M. J. (2024). ConstrainedZero: Chance-Constrained POMDP Planning Using Learned Probabilistic Failure Surrogates and Adaptive Safety Constraints. Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence (IJCAI), 6752-6760. https://www.ijcai.org/proceedings/2024/746

Functions:

spuct_selection: Select among existing children using safety-masked PUCT. spuct_action_progressive_widening: Progressive widening with SPUCT selection. spuct_selection_arena: Arena-tree variant of spuct_selection. spuct_action_progressive_widening_arena: Arena variant of widening+SPUCT.

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_puct.spuct_action_progressive_widening(belief_node, alpha_a, action_sampler, exploration_constant, k_a, failure_dict, delta_prime, action_priors=None, min_visit_count_per_action=1)[source]

Progressive widening with SPUCT selection instead of PUCT.

Parameters:

  • belief_node (BeliefNode) – Current belief node.

  • alpha_a (float) – Progressive widening exponent (0 < alpha_a <= 1).

  • action_sampler (ActionSampler) – Sampler for generating new candidate actions.

  • exploration_constant (float) – PUCT exploration constant c.

  • k_a (float) – Progressive widening coefficient.

  • failure_dict (Dict[int, float]) – Maps id(action_node) to estimated failure probability.

  • delta_prime (float) – Adaptive failure threshold.

  • action_priors (Optional[ndarray]) – Prior probabilities for existing children.

  • min_visit_count_per_action (int) – At the root, ensure every child has been visited at least this many times before selecting via SPUCT.

Return type:

ActionNode

Returns:

Selected or newly created action node.

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_puct.spuct_action_progressive_widening_arena(tree, belief_id, alpha_a, action_sampler, exploration_constant, k_a, failure_dict, delta_prime, action_priors=None, min_visit_count_per_action=1)[source]

Arena variant of spuct_action_progressive_widening().

Returns the action-node ID selected by progressive widening + SPUCT.

Return type:

int

Parameters:
POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_puct.spuct_selection(belief_node, exploration_constant, failure_dict, delta_prime, action_priors=None)[source]

Select an action child using the safety-constrained PUCT criterion.

The selection rule is:

a* = argmax subject_to(a) * [Q_norm(b,a) + c * P(a|b) * sqrt(N(b)) / (1 + N(b,a))]

where subject_to(a) = I(f(a) <= Delta') masks unsafe actions. If ALL actions are unsafe, falls back to unconstrained selection.

Parameters:

  • belief_node (BeliefNode) – Current belief node with at least one action child.

  • exploration_constant (float) – Exploration constant c.

  • failure_dict (Dict[int, float]) – Maps id(action_node) to estimated failure probability.

  • delta_prime (float) – Adaptive failure threshold.

  • action_priors (Optional[ndarray]) – Prior probabilities P(a|b) aligned with belief_node.children. If None, uniform priors are used.

Return type:

ActionNode

Returns:

The action node with the highest safety-masked PUCT score.

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_puct.spuct_selection_arena(tree, belief_id, exploration_constant, failure_dict, delta_prime, action_priors=None)[source]

Arena variant of spuct_selection(). Returns the action-node ID.

failure_dict is keyed by action-node integer ID (not id(node)).

Return type:

int

Parameters:

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_training module

Training utilities for ConstrainedZero network.

This module provides the loss function and training loop for the 3-head ConstrainedZero network. It extends the BetaZero training with an additional binary cross-entropy loss for the failure head.

References

Moss, R. J., Jamgochian, A., Fischer, J., Corso, A., & Kochenderfer, M. J. (2024). ConstrainedZero: Chance-Constrained POMDP Planning Using Learned Probabilistic Failure Surrogates and Adaptive Safety Constraints. Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence (IJCAI), 6752-6760. https://www.ijcai.org/proceedings/2024/746

Functions:

compute_constrained_zero_loss: Combined value + policy + failure loss. train_constrained_network: Multi-epoch training on a replay buffer.

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_training.compute_constrained_zero_loss(network, belief_features, policy_targets, value_targets, failure_targets)[source]

Compute the ConstrainedZero combined loss.

L = MSE(v, g_t) + CrossEntropy(p, pi_t) + BCE(failure_logit, f_t)

Parameters:

  • network (ConstrainedZeroNetwork) – The ConstrainedZero 3-head network.

  • belief_features (Tensor) – Batch of belief feature vectors, shape (B, belief_dim).

  • policy_targets (Tensor) – Batch of policy targets, shape (B, policy_dim).

  • value_targets (Tensor) – Batch of scalar value targets, shape (B,).

  • failure_targets (Tensor) – Batch of binary failure targets, shape (B,).

Return type:

Tuple[Tensor, Dict[str, float]]

Returns:

Tuple of (total_loss, component_dict) where component_dict contains "value_loss", "policy_loss", and "failure_loss" as Python floats.

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_training.train_constrained_network(network, buffer, n_epochs=10, batch_size=256, learning_rate=0.001, weight_decay=0.0001, track_gradients=False, input_mean=None, input_std=None, value_mean=None, value_std=None)[source]

Train the 3-head network for multiple epochs on buffered data.

Parameters:

  • network (ConstrainedZeroNetwork) – Network to train (modified in-place).

  • buffer (ConstrainedTrainingBuffer) – Replay buffer with constrained training examples.

  • n_epochs (int) – Number of full passes over the buffer.

  • batch_size (int) – Mini-batch size.

  • learning_rate (float) – Adam learning rate.

  • weight_decay (float) – L2 regularisation coefficient.

  • track_gradients (bool) – When True, gradient and weight norms are computed per-batch/epoch and included in the returned metrics.

  • input_mean (Optional[ndarray]) – Per-feature mean for input normalisation (None = disabled).

  • input_std (Optional[ndarray]) – Per-feature std for input normalisation (None = disabled).

  • value_mean (Optional[float]) – Scalar mean for value normalisation (None = disabled).

  • value_std (Optional[float]) – Scalar std for value normalisation (None = disabled).

Returns:

"total_loss", "value_loss", "policy_loss", "failure_loss". When track_gradients is True, also includes "grad_norm/global", "grad_norm/trunk", "grad_norm/policy_head", "grad_norm/value_head", "grad_norm/failure_head", and "weight_norm/global".

Return type:

Dict[str, List[float]]

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_training_buffer module

Iteration-slot replay buffer for ConstrainedZero training examples.

This module extends the BetaZero training buffer with an additional failure target, used for training the 3-head ConstrainedZero network.

References

Moss, R. J., Jamgochian, A., Fischer, J., Corso, A., & Kochenderfer, M. J. (2024). ConstrainedZero: Chance-Constrained POMDP Planning Using Learned Probabilistic Failure Surrogates and Adaptive Safety Constraints. Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence (IJCAI), 6752-6760. https://www.ijcai.org/proceedings/2024/746

Classes:

ConstrainedTrainingExample: Training datum with failure target. ConstrainedTrainingBuffer: Buffer returning 4-tuple batches.

class POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_training_buffer.ConstrainedTrainingBuffer(n_buffer=1)[source]

Bases: TrainingBuffer

Iteration-slot replay buffer for ConstrainedZero training.

Extends TrainingBuffer to store and sample ConstrainedTrainingExample instances, returning a 4-tuple from sample_batch() that includes failure targets.

Parameters:

n_buffer (int) – Number of policy-iteration slots to retain. With the default n_buffer=1 only the current iteration’s data is used for training (on-policy).

Example

>>> import numpy as np
>>> from POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_training_buffer import (
...     ConstrainedTrainingBuffer, ConstrainedTrainingExample,
... )
>>> buf = ConstrainedTrainingBuffer(n_buffer=1)
>>> buf.begin_iteration()
>>> buf.add(ConstrainedTrainingExample(np.zeros(4), np.array([0.5, 0.5]), 1.0, 0.0))
>>> len(buf)
1
>>> batch = buf.sample_batch(1)
>>> len(batch)
4
>>> batch[0].shape
(1, 4)
add(example)[source]

Append a constrained example to the current iteration slot.

Return type:

None

Parameters:

example (ConstrainedTrainingExample)

sample_batch(batch_size)[source]

Sample a random mini-batch including failure targets.

Parameters:

batch_size (int) – Number of examples to sample (with replacement if batch_size > len(buffer)).

Return type:

Tuple[ndarray, ndarray, ndarray, ndarray]

Returns:

Tuple of (belief_features, policy_targets, value_targets, failure_targets). - belief_features: shape (batch_size, belief_dim) - policy_targets: shape (batch_size, policy_dim) - value_targets: shape (batch_size,) - failure_targets: shape (batch_size,)

class POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_training_buffer.ConstrainedTrainingExample(belief_features, policy_target, value_target, failure_target)[source]

Bases: object

Single training datum for ConstrainedZero network training.

Parameters:
belief_features

Belief feature vector phi(b), shape (belief_dim,).

policy_target

Q-weighted policy target pi_qw, shape (n_actions,) (discrete).

value_target

Discounted return g_t.

failure_target

Binary episode-level failure indicator (1.0 if failure occurred).

belief_features: ndarray
failure_target: float
policy_target: ndarray
value_target: float

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_zero module

ConstrainedZero planner: neural MCTS for Chance-Constrained POMDPs.

This module implements the ConstrainedZero algorithm, which extends BetaZero to solve CC-POMDPs. It adds a 3-head network with a failure probability head, safety-constrained PUCT (SPUCT), adaptive failure threshold calibration via conformal inference, and constrained policy targets for training.

References

Moss, R. J., Jamgochian, A., Fischer, J., Corso, A., & Kochenderfer, M. J. (2024). ConstrainedZero: Chance-Constrained POMDP Planning Using Learned Probabilistic Failure Surrogates and Adaptive Safety Constraints. Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence (IJCAI), 6752-6760. https://www.ijcai.org/proceedings/2024/746

Implementation note:

Operates on the column-store arena POMDPPlanners.core.tree.arena.Tree via its BetaZero parent. _failure_dict and _delta_dict are keyed by integer node IDs (instead of id(node) as in the legacy anytree version).

Classes:

ConstrainedZero: Main planner extending BetaZero.

class POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_zero.ConstrainedZero(environment, discount_factor, depth, name, action_sampler, failure_fn, delta_0=0.01, eta=1e-05, delta_compounding=1.0, k_a=1.0, alpha_a=0.5, k_o=1.0, alpha_o=0.5, exploration_constant=1.0, time_out_in_seconds=None, n_simulations=None, min_visit_count_per_action=1, network=None, belief_representation=None, state_dim=None, z_q=1.0, z_n=1.0, temperature=1.0, n_buffer=1, training_batch_size=256, training_epochs=10, learning_rate=0.001, weight_decay=0.0001, hidden_sizes=(128, 128), use_dropout=True, p_dropout=0.2, track_gradients=False, normalize_inputs=True, normalize_values=True, log_path=None, debug=False, use_queue_logger=False)[source]

Bases: BetaZero

ConstrainedZero: Neural MCTS for Chance-Constrained POMDPs.

Extends BetaZero with:

  1. 3-head network: Adds a failure probability head alongside policy and value.

  2. SPUCT selection: Safety-constrained PUCT that masks unsafe actions.

  3. Adaptive Delta (conformal inference): Calibrates the failure threshold during tree search using online conformal inference.

  4. Failure propagation: Tracks failure probability per action node using p = p_immediate + delta_compounding * (1 - p_immediate) * p_next.

  5. Constrained policy targets: Applies safety mask during target computation.

Parameters:
failure_fn

User-provided function state -> bool defining failure.

delta_0

Nominal failure probability threshold.

eta

Learning rate for adaptive Delta calibration.

delta_compounding

Discount factor for failure propagation.

Example

>>> import numpy as np
>>> np.random.seed(42)
>>> from POMDPPlanners.environments.tiger_pomdp import TigerPOMDP
>>> from POMDPPlanners.core.belief import get_initial_belief
>>> from POMDPPlanners.utils.action_samplers import DiscreteActionSampler
>>> from POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_zero import ConstrainedZero
>>>
>>> env = TigerPOMDP(discount_factor=0.95)
>>> sampler = DiscreteActionSampler(env.get_actions())
>>> planner = ConstrainedZero(
...     environment=env,
...     discount_factor=0.95,
...     depth=3,
...     name="CZ_Tiger",
...     action_sampler=sampler,
...     n_simulations=20,
...     state_dim=1,
...     failure_fn=lambda s: False,
... )
>>> belief = get_initial_belief(env, n_particles=10)
>>> actions, run_data = planner.action(belief)
>>> actions[0] in env.get_actions()
True
get_metric_keys()[source]

Return the loss-metric key names produced by train_step().

Return type:

List[str]

network: ConstrainedZeroNetwork

POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_zero_network module

Three-head neural network for ConstrainedZero.

This module extends the BetaZero dual-head network with an additional failure probability head. The failure head outputs a raw logit; sigmoid is applied during prediction to produce a probability in [0, 1].

References

Moss, R. J., Jamgochian, A., Fischer, J., Corso, A., & Kochenderfer, M. J. (2024). ConstrainedZero: Chance-Constrained POMDP Planning Using Learned Probabilistic Failure Surrogates and Adaptive Safety Constraints. Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence (IJCAI), 6752-6760. https://www.ijcai.org/proceedings/2024/746

Classes:

ConstrainedZeroNetwork: Shared-trunk network with policy, value, and failure heads.

class POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_zero_network.ConstrainedZeroNetwork(belief_dim, action_space_type, n_actions=None, action_dim=None, hidden_sizes=(128, 128), use_dropout=True, p_dropout=0.2)[source]

Bases: BetaZeroNetwork

Three-head neural network for ConstrainedZero.

Architecture:

  • Shared trunk: Linear(belief_dim, h) → ReLU → [Dropout] → ... → Linear(h, h) → ReLU → [Dropout]

  • Policy head: inherited from BetaZeroNetwork

  • Value head: inherited from BetaZeroNetwork

  • Failure head: Linear(h, h//2) -> ReLU -> Linear(h//2, 1)

The failure head outputs a raw logit. During predict(), sigmoid is applied to produce a failure probability in [0, 1].

Parameters:

  • belief_dim (int) – Dimensionality of the belief feature vector phi(b).

  • action_space_type (str) – "discrete" or "continuous".

  • n_actions (Optional[int]) – Number of discrete actions (required when action_space_type="discrete").

  • action_dim (Optional[int]) – Dimensionality of continuous actions (required when action_space_type="continuous").

  • hidden_sizes (Sequence[int]) – Tuple of hidden layer widths for the shared trunk.

  • use_dropout (bool) – If True, apply dropout after each ReLU in the shared trunk (default True).

  • p_dropout (float) – Dropout probability for trunk layers (default 0.2).

Example

>>> import numpy as np
>>> from POMDPPlanners.planners.mcts_planners.constrained_zero.constrained_zero_network import ConstrainedZeroNetwork
>>> net = ConstrainedZeroNetwork(belief_dim=4, action_space_type="discrete", n_actions=3, use_dropout=False)
>>> policy, value, failure_prob = net.predict(np.zeros(4, dtype=np.float32))
>>> policy.shape
(3,)
>>> isinstance(value, float)
True
>>> 0.0 <= failure_prob <= 1.0
True
forward(belief_features)[source]

Forward pass returning raw policy, value, and failure logit.

Parameters:

belief_features (Tensor) – Tensor of shape (batch, belief_dim) or (belief_dim,).

Return type:

Tuple[Tensor, Tensor, Tensor]

Returns:

Tuple of (policy_output, value, failure_logit) tensors.

predict(belief_features)[source]

Single-sample inference returning numpy policy, value, and failure probability.

Switches to eval mode before inference to disable dropout, then restores the original training mode.

Parameters:

belief_features (ndarray) – 1-D array of shape (belief_dim,).

Return type:

Tuple[ndarray, float, float]

Returns:

Tuple of (policy, value, failure_prob). - Discrete: policy is a probability vector summing to 1. - Continuous: policy is [mean, log_std]. - value is a Python float. - failure_prob is a Python float in [0, 1].

predict_batch(belief_features_batch)[source]

Batched inference returning numpy policy, value, and failure probability arrays.

Switches to eval mode before inference to disable dropout, then restores the original training mode.

Parameters:

belief_features_batch (ndarray) – 2-D array of shape (N, belief_dim).

Return type:

Tuple[ndarray, ndarray, ndarray]

Returns:

Tuple of (policies, values, failure_probs). - Discrete: policies is (N, n_actions) probability matrix. - Continuous: policies is (N, 2*action_dim) with [mean, log_std]. - values is (N,) array of floats. - failure_probs is (N,) array of floats in [0, 1].