POMDPPlanners.environments.rock_sample_pomdp package

RockSample POMDP Environment Module.

This module provides the RockSample POMDP environment implementation and related components for robot navigation and sampling tasks.

Classes:

RockSamplePOMDP: Main POMDP environment for rock sampling tasks RockSampleState: State representation with robot position and rock qualities RockSampleVisualizer: Visualization utilities for RockSample POMDP episodes

class POMDPPlanners.environments.rock_sample_pomdp.RewardModelType(*values)

Bases: Enum

Reward-model variants for RockSamplePOMDP.

Variants differ only in how the dangerous-area penalty is applied — base scoring (exit / sample / sense / step) is identical across all three. See POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp_utils.rock_sample_reward_models for the per-variant semantics.

CONSTANT_HAZARD_PENALTY = 'constant_hazard_penalty'

DISTANCE_DECAYED_HAZARD_PENALTY = 'distance_decayed_hazard_penalty'

ZERO_MEAN_HAZARD_SHOCK = 'zero_mean_hazard_shock'

class POMDPPlanners.environments.rock_sample_pomdp.RockSamplePOMDP(map_size=(5, 5), rock_positions=None, init_pos=(0, 0), sensor_efficiency=10.0, bad_rock_penalty=-10.0, good_rock_reward=10.0, step_penalty=0.0, sensor_use_penalty=0.0, exit_reward=10.0, dangerous_areas=None, dangerous_area_radius=1.0, dangerous_area_penalty=-5.0, dangerous_area_hit_probability=1.0, reward_model_type=RewardModelType.CONSTANT_HAZARD_PENALTY, penalty_decay=1.0, discount_factor=0.95, name='RockSample', output_dir=None, debug=False, use_queue_logger=False)

Bases: DiscreteActionsEnvironment

RockSample POMDP environment

This environment implements the classic rock sampling problem where a robot must navigate a grid, use sensors to evaluate rocks, and decide which ones to sample while balancing exploration costs and sampling rewards.

Stochasticity:

The dangerous-area penalty can be applied either deterministically (the default) or stochastically. When dangerous_area_hit_probability == 1.0 (default), the penalty is applied every time the robot’s next position lies inside a dangerous area, matching legacy behavior. When dangerous_area_hit_probability < 1.0, the penalty is applied only with that probability per reward() / reward_batch() call (one Bernoulli draw per state), producing a heavy-tailed return distribution suitable for benchmarking risk-sensitive planners (e.g. ICVaR-aware MCTS) against expected-value MCTS on the same env. Note that this makes reward(state, action) non-deterministic given a state-action pair, so any external caching that assumes deterministic rewards must be aware of this. transition_log_probability is unaffected.

Parameters:

• map_size (Tuple[int, int])

• rock_positions (List[Tuple[int, int]] | None)

• init_pos (Tuple[int, int])

• sensor_efficiency (float)

• bad_rock_penalty (float)

• good_rock_reward (float)

• step_penalty (float)

• sensor_use_penalty (float)

• exit_reward (float)

• dangerous_areas (List[Tuple[int, int]] | None)

• dangerous_area_radius (float)

• dangerous_area_penalty (float)

• dangerous_area_hit_probability (float)

• reward_model_type (RewardModelType)

• penalty_decay (float)

• discount_factor (float)

• name (str)

• output_dir (Path | None)

• debug (bool)

• use_queue_logger (bool)

map_size

Grid dimensions as (rows, cols)

rock_positions

List of rock positions as (row, col) tuples

init_pos

Initial robot position

sensor_efficiency

Sensor noise parameter (higher = less noise)

bad_rock_penalty

Penalty for sampling a bad rock

good_rock_reward

Reward for sampling a good rock

step_penalty

Cost for each action

sensor_use_penalty

Additional cost for using sensor

exit_reward

Reward for reaching the exit

Example

>>> import numpy as np
>>> np.random.seed(42)  # For reproducible results
>>>
>>> # Initialize environment
>>> env = RockSamplePOMDP(map_size=(5, 5), rock_positions=[(0, 0), (2, 2), (3, 3)])
>>>
>>> # Get initial state and actions
>>> initial_state = env.initial_state_dist().sample()[0]
>>> actions = env.get_actions()
>>>
>>> # Sample complete step using convenience method
>>> action = actions[0]
>>> next_state, observation, reward = env.sample_next_step(initial_state, action)
>>>
>>> # Check terminal condition
>>> env.is_terminal(initial_state)
False

cache_visualization(history, cache_path)

Cache visualization of episode history.

Parameters:

• history (List[StepData]) – Episode history containing states, actions, and rewards

• cache_path (Path) – Path where to save the visualization (must end with .gif)

Return type:

None

compute_metrics(histories)

Compute environment-specific metrics.

Return type:

List[MetricValue]

Parameters:

histories (List[History])

dangerous_areas: List[Tuple[int, int]]

get_actions()

Get all available actions.

Return type:

List[int]

get_metric_names()

Get names of RockSample POMDP specific metrics.

Returns:

avg_rocks_sampled, exit_success_rate, and average_dangerous_area_steps

Return type:

List[str]

hash_action(action)

Return a hashable key consistent with action equality.

Used by tree-search planners to index action children of a belief node in O(1). The returned key MUST satisfy:

action_a == action_b   (per env's notion of equality)
==> hash_action(action_a) == hash_action(action_b)

Subclasses with non-hashable actions (e.g. np.ndarray) must override to return a hashable surrogate (tobytes() is the standard choice for ndarray actions, which mirrors the np.array_equal semantics used by the linear-scan fallback).

Parameters:

action (Any) – Action to hash.

Return type:

Hashable

Returns:

A hashable key derived from action.

initial_observation_dist()

Get initial observation distribution.

Return type:

DiscreteDistribution

initial_state_dist()

Get initial state distribution.

Return type:

DiscreteDistribution

is_equal_observation(observation1, observation2)

Check if two observations are equal.

Return type:

bool

Parameters:

• observation1 (Any)

• observation2 (Any)

is_terminal(state)

Check if state is terminal.

Return type:

bool

Parameters:

state (ndarray)

observation_log_probability(next_state, action, observations)

Log-probability of each candidate observation under (next_state, action).

Returns np.ndarray of shape (N,) where N is the number of candidate observations. Subclasses must implement.

Return type:

ndarray

Parameters:

• next_state (ndarray)

• action (int)

• observations (Any)

observation_log_probability_per_state(next_states, action, observation)

Log-probability of one observation under each candidate next-state.

Used by particle filters: given N candidate next-states and ONE observation, return N log-likelihoods.

The default implementation falls back to a per-state Python loop delegating to observation_log_probability(). Native-backed envs (those whose observation kernel exposes batch_log_likelihood(next_states_array, observation_array)) should override to avoid the loop.

Parameters:

• next_states (Any) – A sequence (length N) or ndarray of shape (N, *dim) of candidate next-states.

• action (int) – The action that was executed.

• observation (Any) – A single observation.

Return type:

ndarray

Returns:

ndarray of shape (N,) with log-probabilities or log-PDFs.

reward(state, action, next_state=None)

Calculate immediate reward.

Uses the realised next_state when supplied (e.g. by Environment.sample_next_step()) so the dangerous-area penalty fires against the same outcome as the trajectory instead of a fresh draw.

Return type:

float

Parameters:

• state (ndarray)

• action (int)

• next_state (Any)

reward_batch(states, action, next_states=None)

Calculate rewards for a batch of states given a single action.

Threads caller-supplied next_states through to the dangerous-area position check so the batch path agrees with the scalar reward() whenever Environment.sample_next_step (or any other caller) pre-samples next states. When next_states is None, we fall back to closed-form reconstruction of the next robot position from (state, action); RockSample transitions are deterministic, so this fallback matches a fresh draw from sample_next_state(). The per-call Bernoulli refund for the dangerous-area penalty is preserved in both branches.

Return type:

ndarray

Parameters:

• states (Any)

• action (int)

• next_states (Any)

reward_model: BaseRockSampleRewardModel

sample_next_state(state, action, n_samples=1)

Sample one or more next states for (state, action).

Hot-path entry point used by MCTS planners and particle filters. Subclasses must implement.

Returns:

a single next state of the env’s native type. When n_samples > 1: an array-like of length n_samples (numeric envs return np.ndarray of shape (n_samples, *dim); structured envs return List[T]).

Return type:

Any

Parameters:

• state (ndarray)

• action (int)

• n_samples (int)

sample_next_state_batch(states, action)

Sample one next state per input state, all under the same action.

Used by particle filters: given N current particles and one action, draw N next states (one per particle) in a single vectorized call.

The default implementation falls back to a per-state Python loop delegating to sample_next_state(). Native-backed envs (those whose state-transition kernel exposes batch_sample(states_array)) should override to avoid the loop.

Parameters:

• states (Any) – A sequence (length N) or ndarray of shape (N, *dim) of input particles.

• action (int) – A single action to apply to every particle.

Returns:

np.ndarray of shape (N, *dim). For structured envs (Tiger strings, Pacman tuples): a list of length N.

Return type:

ndarray

sample_next_step(state, action)

Override to avoid reward() recomputing next state.

Return type:

Tuple[ndarray, str, float]

Parameters:

• state (ndarray)

• action (int)

sample_observation(next_state, action, n_samples=1)

Sample one or more observations for (next_state, action).

Hot-path entry point used by MCTS planners and particle filters. Subclasses must implement.

Returns:

a single observation. When n_samples > 1: an array-like of length n_samples.

Return type:

Any

Parameters:

• next_state (ndarray)

• action (int)

• n_samples (int)

simulate_random_rollout(state, action_sampler, max_depth, discount_factor, depth=0)

Random rollout via native C++ deterministic transition and reward kernel.

The C++ kernel applies the variant-aware dangerous-area reward term directly, so no Python fallback is required when danger zones are configured.

Parameters:

• state (Any) – Current RockSample state array.

• action_sampler (Any) – Object with a sample() method returning an integer action. Currently unused — actions are drawn uniformly by the native kernel.

• max_depth (int) – Maximum rollout depth.

• discount_factor (float) – Per-step discount factor.

• depth (int) – Depth already consumed by the search tree. Defaults to 0.

Return type:

float

Returns:

Discounted sum of immediate rewards along the sampled trajectory.

transition_log_probability(state, action, next_states)

Log-probability of each candidate next state under (state, action).

Returns np.ndarray of shape (N,) where N is the number of candidate next states. Subclasses must implement.

Return type:

ndarray

Parameters:

• state (ndarray)

• action (int)

• next_states (Any)

visualize_path(path, actions, cache_path)

Visualize robot path through the environment.

Parameters:

• path (List[ndarray]) – List of states representing the path

• actions (List[int]) – List of actions taken at each state

• cache_path (Path) – Path where to save the animation (must end with .gif)

Return type:

None

POMDPPlanners.environments.rock_sample_pomdp.RockSampleState

alias of ndarray

class POMDPPlanners.environments.rock_sample_pomdp.RockSampleVectorizedUpdater(map_rows, map_cols, num_rocks, rock_positions, sensor_efficiency)

Bases: VectorizedParticleBeliefUpdater

Vectorized particle belief updater for the RockSample POMDP.

Stores precomputed environment parameters and dispatches batched transitions and observation log-likelihood evaluations to the native C++ extension. State layout per particle is [robot_row, robot_col, rock_0_quality, ..., rock_{R-1}_quality].

Parameters:

• map_rows (int)

• map_cols (int)

• num_rocks (int)

• rock_positions (np.ndarray)

• sensor_efficiency (float)

map_rows

Number of grid rows.

map_cols

Number of grid columns.

num_rocks

Number of rocks in the environment.

rock_positions

Array of shape (R, 2) with rock (row, col) positions.

sensor_efficiency

Sensor noise parameter (higher = less noise).

batch_observation_log_likelihood(next_particles, action, observation)

Compute observation log-likelihoods for all particles.

Parameters:

• next_particles (ndarray) – Array of shape (N, 2 + num_rocks).

• action (ndarray) – Scalar action index.

• observation (ndarray) – Integer-encoded observation (0=none, 1=good, 2=bad).

Return type:

ndarray

Returns:

Log-likelihoods of shape (N,).

batch_transition(particles, action)

Transition all particles for the given action.

Parameters:

• particles (ndarray) – Array of shape (N, 2 + num_rocks).

• action (ndarray) – Scalar action index.

Return type:

ndarray

Returns:

Next-state particles of shape (N, 2 + num_rocks).

property config_id: str

Return a deterministic identifier for this updater configuration.

classmethod from_environment(env)

Construct an updater from a RockSamplePOMDP instance.

Return type:

RockSampleVectorizedUpdater

Parameters:

env (RockSamplePOMDP)

class POMDPPlanners.environments.rock_sample_pomdp.RockSampleVisualizer(env)

Bases: object

Handles visualization and animation for RockSample POMDP environments.

This class encapsulates all visualization logic for RockSample POMDP episodes, creating animated GIFs showing robot movement, rock sampling, sensor checks, dangerous areas, and exit behavior.

Parameters:

env (RockSamplePOMDP)

env

Reference to the RockSamplePOMDP environment instance

map_size

Grid dimensions as (rows, cols)

rock_positions

List of rock positions

action_names

Names of available actions

action_to_vector

Mapping from action indices to direction vectors

dangerous_areas

List of dangerous area center positions

dangerous_area_radius

Radius around dangerous area centers

create_visualization(history, cache_path)

Create animated visualization of a RockSample POMDP episode.

Creates an animated GIF showing the robot navigating, sampling rocks, using sensors, and exiting the grid.

Parameters:

• history (List[StepData]) – Episode history containing states, actions, and rewards

• cache_path (Path) – Path where to save the visualization (must end with .gif)

Raises:

• ValueError – If history is empty or cache_path doesn’t end with .gif

• TypeError – If cache_path is not a Path object or history is invalid

Return type:

None

visualize_path(path, actions, cache_path)

Visualize robot path through the environment.

Parameters:

• path (List[ndarray]) – List of states representing the path

• actions (List[int]) – List of actions taken at each state

• cache_path (Path) – Path where to save the animation

Return type:

None

POMDPPlanners.environments.rock_sample_pomdp.create_random_rock_sample(map_size=7, num_rocks=8, seed=None)

Create a random RockSample instance.

Parameters:

• map_size (int) – Size of square grid. Defaults to 7.

• num_rocks (int) – Number of rocks to place. Defaults to 8.

• seed (Optional[int]) – Random seed. Defaults to None.

Return type:

RockSamplePOMDP

Returns:

Randomly configured RockSample POMDP

POMDPPlanners.environments.rock_sample_pomdp.create_rock_sample_state(robot_pos, rocks)

Create a RockSample state as a numpy array.

Parameters:

• robot_pos (Tuple[int, int]) – Robot position as (row, col) tuple

• rocks (Tuple[bool, ...]) – Tuple of booleans indicating rock quality (True=good, False=bad)

Returns:

[robot_row, robot_col, rock_0, rock_1, …, rock_n] where rock values are 1.0 for good (True) and 0.0 for bad (False)

Return type:

ndarray

POMDPPlanners.environments.rock_sample_pomdp.create_rocksample_belief(env, belief_type=BeliefType.VECTORIZED_PARTICLE, n_particles=200, **kwargs)

Create a belief object for the RockSample POMDP.

Parameters:

• env (RockSamplePOMDP) – RockSample environment instance.

• belief_type (BeliefType) – Desired belief representation. Supports PARTICLE and VECTORIZED_PARTICLE.

• n_particles (int) – Number of particles. Defaults to 200.

• **kwargs (Any) – Reserved for future use.

Return type:

Belief

Returns:

A configured belief object.

Raises:

ValueError – If belief_type is not supported.

POMDPPlanners.environments.rock_sample_pomdp.get_robot_pos(state)

Extract robot position from state array.

Parameters:

state (ndarray) – State array

Return type:

Tuple[int, int]

Returns:

Robot position as (row, col) tuple

POMDPPlanners.environments.rock_sample_pomdp.get_rocks(state)

Extract rock qualities from state array.

Parameters:

state (ndarray) – State array

Return type:

Tuple[bool, ...]

Returns:

Tuple of booleans indicating rock quality

POMDPPlanners.environments.rock_sample_pomdp.states_equal(state1, state2)

Check if two states are equal.

Parameters:

• state1 (ndarray) – First state

• state2 (ndarray) – Second state

Return type:

bool

Returns:

True if states are equal

Subpackages

• POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp_beliefs package
  • RockSampleVectorizedUpdater
    • RockSampleVectorizedUpdater.map_rows
    • RockSampleVectorizedUpdater.map_cols
    • RockSampleVectorizedUpdater.num_rocks
    • RockSampleVectorizedUpdater.rock_positions
    • RockSampleVectorizedUpdater.sensor_efficiency
    • RockSampleVectorizedUpdater.batch_observation_log_likelihood()
    • RockSampleVectorizedUpdater.batch_transition()
    • RockSampleVectorizedUpdater.config_id
    • RockSampleVectorizedUpdater.from_environment()
  • create_rocksample_belief()
  • Submodules
  • POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp_beliefs.rocksample_belief_factory module
    • RockSampleVectorizedWeightedParticleBelief
      • RockSampleVectorizedWeightedParticleBelief.update()
    • create_rocksample_belief()
  • POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp_beliefs.rocksample_vectorized_updater module
    • RockSampleVectorizedUpdater
      • RockSampleVectorizedUpdater.map_rows
      • RockSampleVectorizedUpdater.map_cols
      • RockSampleVectorizedUpdater.num_rocks
      • RockSampleVectorizedUpdater.rock_positions
      • RockSampleVectorizedUpdater.sensor_efficiency
      • RockSampleVectorizedUpdater.batch_observation_log_likelihood()
      • RockSampleVectorizedUpdater.batch_transition()
      • RockSampleVectorizedUpdater.config_id
      • RockSampleVectorizedUpdater.from_environment()
• POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp_utils package
  • Submodules
  • POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp_utils.rock_sample_reward_models module
    • BaseRockSampleRewardModel
      • BaseRockSampleRewardModel.compute_reward()
      • BaseRockSampleRewardModel.compute_reward_batch()
    • RockSampleDistanceDecayedHazardPenaltyRewardModel
    • RockSampleRewardModel
      • RockSampleRewardModel.compute_reward()
      • RockSampleRewardModel.compute_reward_batch()
    • RockSampleZeroMeanHazardShockRewardModel

Submodules

POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp module

Module for RockSample POMDP environment.

This module provides the RockSample POMDP environment implementation based on the classic rock sampling problem.

The environment involves a robot navigating a grid world with rocks that are either good or bad. The robot must use a noisy sensor to determine rock quality and decide whether to sample them, balancing exploration and exploitation.

Classes:

RockSampleState: Represents the state of the environment RockSamplePOMDP: The main POMDP environment implementation

class POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp.RewardModelType(*values)

Bases: Enum

Reward-model variants for RockSamplePOMDP.

Variants differ only in how the dangerous-area penalty is applied — base scoring (exit / sample / sense / step) is identical across all three. See POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp_utils.rock_sample_reward_models for the per-variant semantics.

CONSTANT_HAZARD_PENALTY = 'constant_hazard_penalty'

DISTANCE_DECAYED_HAZARD_PENALTY = 'distance_decayed_hazard_penalty'

ZERO_MEAN_HAZARD_SHOCK = 'zero_mean_hazard_shock'

class POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp.RockSamplePOMDP(map_size=(5, 5), rock_positions=None, init_pos=(0, 0), sensor_efficiency=10.0, bad_rock_penalty=-10.0, good_rock_reward=10.0, step_penalty=0.0, sensor_use_penalty=0.0, exit_reward=10.0, dangerous_areas=None, dangerous_area_radius=1.0, dangerous_area_penalty=-5.0, dangerous_area_hit_probability=1.0, reward_model_type=RewardModelType.CONSTANT_HAZARD_PENALTY, penalty_decay=1.0, discount_factor=0.95, name='RockSample', output_dir=None, debug=False, use_queue_logger=False)

Bases: DiscreteActionsEnvironment

RockSample POMDP environment

This environment implements the classic rock sampling problem where a robot must navigate a grid, use sensors to evaluate rocks, and decide which ones to sample while balancing exploration costs and sampling rewards.

Stochasticity:

The dangerous-area penalty can be applied either deterministically (the default) or stochastically. When dangerous_area_hit_probability == 1.0 (default), the penalty is applied every time the robot’s next position lies inside a dangerous area, matching legacy behavior. When dangerous_area_hit_probability < 1.0, the penalty is applied only with that probability per reward() / reward_batch() call (one Bernoulli draw per state), producing a heavy-tailed return distribution suitable for benchmarking risk-sensitive planners (e.g. ICVaR-aware MCTS) against expected-value MCTS on the same env. Note that this makes reward(state, action) non-deterministic given a state-action pair, so any external caching that assumes deterministic rewards must be aware of this. transition_log_probability is unaffected.

Parameters:

• map_size (Tuple[int, int])

• rock_positions (List[Tuple[int, int]] | None)

• init_pos (Tuple[int, int])

• sensor_efficiency (float)

• bad_rock_penalty (float)

• good_rock_reward (float)

• step_penalty (float)

• sensor_use_penalty (float)

• exit_reward (float)

• dangerous_areas (List[Tuple[int, int]])

• dangerous_area_radius (float)

• dangerous_area_penalty (float)

• dangerous_area_hit_probability (float)

• reward_model_type (RewardModelType)

• penalty_decay (float)

• discount_factor (float)

• name (str)

• output_dir (Path | None)

• debug (bool)

• use_queue_logger (bool)

map_size

Grid dimensions as (rows, cols)

rock_positions

List of rock positions as (row, col) tuples

init_pos

Initial robot position

sensor_efficiency

Sensor noise parameter (higher = less noise)

bad_rock_penalty

Penalty for sampling a bad rock

good_rock_reward

Reward for sampling a good rock

step_penalty

Cost for each action

sensor_use_penalty

Additional cost for using sensor

exit_reward

Reward for reaching the exit

Example

>>> import numpy as np
>>> np.random.seed(42)  # For reproducible results
>>>
>>> # Initialize environment
>>> env = RockSamplePOMDP(map_size=(5, 5), rock_positions=[(0, 0), (2, 2), (3, 3)])
>>>
>>> # Get initial state and actions
>>> initial_state = env.initial_state_dist().sample()[0]
>>> actions = env.get_actions()
>>>
>>> # Sample complete step using convenience method
>>> action = actions[0]
>>> next_state, observation, reward = env.sample_next_step(initial_state, action)
>>>
>>> # Check terminal condition
>>> env.is_terminal(initial_state)
False

cache_visualization(history, cache_path)

Cache visualization of episode history.

Parameters:

• history (List[StepData]) – Episode history containing states, actions, and rewards

• cache_path (Path) – Path where to save the visualization (must end with .gif)

Return type:

None

compute_metrics(histories)

Compute environment-specific metrics.

Return type:

List[MetricValue]

Parameters:

histories (List[History])

dangerous_areas: List[Tuple[int, int]]

get_actions()

Get all available actions.

Return type:

List[int]

get_metric_names()

Get names of RockSample POMDP specific metrics.

Returns:

avg_rocks_sampled, exit_success_rate, and average_dangerous_area_steps

Return type:

List[str]

hash_action(action)

Return a hashable key consistent with action equality.

Used by tree-search planners to index action children of a belief node in O(1). The returned key MUST satisfy:

action_a == action_b   (per env's notion of equality)
==> hash_action(action_a) == hash_action(action_b)

Subclasses with non-hashable actions (e.g. np.ndarray) must override to return a hashable surrogate (tobytes() is the standard choice for ndarray actions, which mirrors the np.array_equal semantics used by the linear-scan fallback).

Parameters:

action (Any) – Action to hash.

Return type:

Hashable

Returns:

A hashable key derived from action.

initial_observation_dist()

Get initial observation distribution.

Return type:

DiscreteDistribution

initial_state_dist()

Get initial state distribution.

Return type:

DiscreteDistribution

is_equal_observation(observation1, observation2)

Check if two observations are equal.

Return type:

bool

Parameters:

• observation1 (Any)

• observation2 (Any)

is_terminal(state)

Check if state is terminal.

Return type:

bool

Parameters:

state (ndarray)

observation_log_probability(next_state, action, observations)

Log-probability of each candidate observation under (next_state, action).

Returns np.ndarray of shape (N,) where N is the number of candidate observations. Subclasses must implement.

Return type:

ndarray

Parameters:

• next_state (ndarray)

• action (int)

• observations (Any)

observation_log_probability_per_state(next_states, action, observation)

Log-probability of one observation under each candidate next-state.

Used by particle filters: given N candidate next-states and ONE observation, return N log-likelihoods.

The default implementation falls back to a per-state Python loop delegating to observation_log_probability(). Native-backed envs (those whose observation kernel exposes batch_log_likelihood(next_states_array, observation_array)) should override to avoid the loop.

Parameters:

• next_states (Any) – A sequence (length N) or ndarray of shape (N, *dim) of candidate next-states.

• action (int) – The action that was executed.

• observation (Any) – A single observation.

Return type:

ndarray

Returns:

ndarray of shape (N,) with log-probabilities or log-PDFs.

reward(state, action, next_state=None)

Calculate immediate reward.

Uses the realised next_state when supplied (e.g. by Environment.sample_next_step()) so the dangerous-area penalty fires against the same outcome as the trajectory instead of a fresh draw.

Return type:

float

Parameters:

• state (ndarray)

• action (int)

• next_state (Any)

reward_batch(states, action, next_states=None)

Calculate rewards for a batch of states given a single action.

Threads caller-supplied next_states through to the dangerous-area position check so the batch path agrees with the scalar reward() whenever Environment.sample_next_step (or any other caller) pre-samples next states. When next_states is None, we fall back to closed-form reconstruction of the next robot position from (state, action); RockSample transitions are deterministic, so this fallback matches a fresh draw from sample_next_state(). The per-call Bernoulli refund for the dangerous-area penalty is preserved in both branches.

Return type:

ndarray

Parameters:

• states (Any)

• action (int)

• next_states (Any)

reward_model: BaseRockSampleRewardModel

sample_next_state(state, action, n_samples=1)

Sample one or more next states for (state, action).

Hot-path entry point used by MCTS planners and particle filters. Subclasses must implement.

Returns:

a single next state of the env’s native type. When n_samples > 1: an array-like of length n_samples (numeric envs return np.ndarray of shape (n_samples, *dim); structured envs return List[T]).

Return type:

Any

Parameters:

• state (ndarray)

• action (int)

• n_samples (int)

sample_next_state_batch(states, action)

Sample one next state per input state, all under the same action.

Used by particle filters: given N current particles and one action, draw N next states (one per particle) in a single vectorized call.

The default implementation falls back to a per-state Python loop delegating to sample_next_state(). Native-backed envs (those whose state-transition kernel exposes batch_sample(states_array)) should override to avoid the loop.

Parameters:

• states (Any) – A sequence (length N) or ndarray of shape (N, *dim) of input particles.

• action (int) – A single action to apply to every particle.

Returns:

np.ndarray of shape (N, *dim). For structured envs (Tiger strings, Pacman tuples): a list of length N.

Return type:

ndarray

sample_next_step(state, action)

Override to avoid reward() recomputing next state.

Return type:

Tuple[ndarray, str, float]

Parameters:

• state (ndarray)

• action (int)

sample_observation(next_state, action, n_samples=1)

Sample one or more observations for (next_state, action).

Hot-path entry point used by MCTS planners and particle filters. Subclasses must implement.

Returns:

a single observation. When n_samples > 1: an array-like of length n_samples.

Return type:

Any

Parameters:

• next_state (ndarray)

• action (int)

• n_samples (int)

simulate_random_rollout(state, action_sampler, max_depth, discount_factor, depth=0)

Random rollout via native C++ deterministic transition and reward kernel.

The C++ kernel applies the variant-aware dangerous-area reward term directly, so no Python fallback is required when danger zones are configured.

Parameters:

• state (Any) – Current RockSample state array.

• action_sampler (Any) – Object with a sample() method returning an integer action. Currently unused — actions are drawn uniformly by the native kernel.

• max_depth (int) – Maximum rollout depth.

• discount_factor (float) – Per-step discount factor.

• depth (int) – Depth already consumed by the search tree. Defaults to 0.

Return type:

float

Returns:

Discounted sum of immediate rewards along the sampled trajectory.

transition_log_probability(state, action, next_states)

Log-probability of each candidate next state under (state, action).

Returns np.ndarray of shape (N,) where N is the number of candidate next states. Subclasses must implement.

Return type:

ndarray

Parameters:

• state (ndarray)

• action (int)

• next_states (Any)

visualize_path(path, actions, cache_path)

Visualize robot path through the environment.

Parameters:

• path (List[ndarray]) – List of states representing the path

• actions (List[int]) – List of actions taken at each state

• cache_path (Path) – Path where to save the animation (must end with .gif)

Return type:

None

class POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp.RockSamplePOMDPMetrics(*values)

Bases: Enum

Metric names for RockSample POMDP environment.

AVERAGE_DANGEROUS_AREA_STEPS = 'average_dangerous_area_steps'

AVG_ROCKS_SAMPLED = 'avg_rocks_sampled'

EXIT_SUCCESS_RATE = 'exit_success_rate'

POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp.create_random_rock_sample(map_size=7, num_rocks=8, seed=None)

Create a random RockSample instance.

Parameters:

• map_size (int) – Size of square grid. Defaults to 7.

• num_rocks (int) – Number of rocks to place. Defaults to 8.

• seed (Optional[int]) – Random seed. Defaults to None.

Return type:

RockSamplePOMDP

Returns:

Randomly configured RockSample POMDP

POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp.create_rock_sample_state(robot_pos, rocks)

Create a RockSample state as a numpy array.

Parameters:

• robot_pos (Tuple[int, int]) – Robot position as (row, col) tuple

• rocks (Tuple[bool, ...]) – Tuple of booleans indicating rock quality (True=good, False=bad)

Returns:

[robot_row, robot_col, rock_0, rock_1, …, rock_n] where rock values are 1.0 for good (True) and 0.0 for bad (False)

Return type:

ndarray

POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp.get_robot_pos(state)

Extract robot position from state array.

Parameters:

state (ndarray) – State array

Return type:

Tuple[int, int]

Returns:

Robot position as (row, col) tuple

POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp.get_rocks(state)

Extract rock qualities from state array.

Parameters:

state (ndarray) – State array

Return type:

Tuple[bool, ...]

Returns:

Tuple of booleans indicating rock quality

POMDPPlanners.environments.rock_sample_pomdp.rock_sample_pomdp.states_equal(state1, state2)

Check if two states are equal.

Parameters:

• state1 (ndarray) – First state

• state2 (ndarray) – Second state

Return type:

bool

Returns:

True if states are equal

POMDPPlanners.environments.rock_sample_pomdp.rock_sample_visualizer module

Visualization module for RockSample POMDP Environment.

This module provides visualization capabilities for RockSample POMDP episodes, creating animated GIFs showing robot movement, rock sampling, sensor usage, and exit behavior.

Classes:

RockSampleVisualizer: Handles all visualization logic for RockSample POMDP

class POMDPPlanners.environments.rock_sample_pomdp.rock_sample_visualizer.RockSampleVisualizer(env)

Bases: object

Handles visualization and animation for RockSample POMDP environments.

This class encapsulates all visualization logic for RockSample POMDP episodes, creating animated GIFs showing robot movement, rock sampling, sensor checks, dangerous areas, and exit behavior.

Parameters:

env (RockSamplePOMDP)

env

Reference to the RockSamplePOMDP environment instance

map_size

Grid dimensions as (rows, cols)

rock_positions

List of rock positions

action_names

Names of available actions

action_to_vector

Mapping from action indices to direction vectors

dangerous_areas

List of dangerous area center positions

dangerous_area_radius

Radius around dangerous area centers

create_visualization(history, cache_path)

Create animated visualization of a RockSample POMDP episode.

Creates an animated GIF showing the robot navigating, sampling rocks, using sensors, and exiting the grid.

Parameters:

• history (List[StepData]) – Episode history containing states, actions, and rewards

• cache_path (Path) – Path where to save the visualization (must end with .gif)

Raises:

• ValueError – If history is empty or cache_path doesn’t end with .gif

• TypeError – If cache_path is not a Path object or history is invalid

Return type:

None

visualize_path(path, actions, cache_path)

Visualize robot path through the environment.

Parameters:

• path (List[ndarray]) – List of states representing the path

• actions (List[int]) – List of actions taken at each state

• cache_path (Path) – Path where to save the animation

Return type:

None