POMDPPlanners.environments.safety_ant_velocity_pomdp package
Safety Ant Velocity POMDP Environment Package.
This package implements a safety-critical velocity control task where an agent must navigate while avoiding unsafe velocities.
- class POMDPPlanners.environments.safety_ant_velocity_pomdp.SafeAntVelocityPOMDP(discount_factor, safe_velocity_threshold=2.0, max_force=1.0, dt=0.1, mass=1.0, damping=0.1, position_noise=0.1, velocity_noise=0.2, safety_violation_penalty=-100.0, movement_reward_scale=1.0, name='SafeVelocityPOMDP', output_dir=None, debug=False, use_queue_logger=False)[source]
Bases:
DiscreteActionsEnvironmentSafety-critical velocity control task formulated as a POMDP.
This environment presents a safety-critical control problem where an agent must navigate while keeping velocity below a safety threshold. The challenge comes from balancing exploration rewards with safety constraints under noisy velocity observations.
Problem Structure: - State: [position_x, position_y, velocity_x, velocity_y] (continuous) - Actions: [0=no force, 1=small, 2=medium, 3=large force] (discrete) - Observations: Noisy position and velocity measurements (continuous) - Rewards: Movement reward - safety violation penalty (if unsafe) - Safety constraint: velocity magnitude ≤ safe_velocity_threshold - Termination: Velocity exceeds 1.5x safety threshold
Safety Features: - Tracks safety and critical violation rates - Heavy penalties for constraint violations - Configurable safety thresholds and penalties - Physics simulation with uncertainty in force direction
Example
>>> import numpy as np >>> np.random.seed(42) # For reproducible results >>> >>> # Initialize environment >>> env = SafeAntVelocityPOMDP(discount_factor=0.99) >>> >>> # 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
- Parameters:
- cache_visualization(history, output_dir, episode_index)[source]
Cache animated visualization of the safety ant velocity episode.
Creates an animated GIF showing the ant’s movement trajectory with velocity vectors, safety zones, force applications, and safety constraint violations.
- compute_metrics(histories)[source]
Compute environment-specific metrics from episode histories.
This method can be overridden by subclasses to provide custom metric calculations beyond standard return and episode length.
- get_actions()[source]
Get all possible actions in the discrete action space.
Note
Subclasses must implement this method to enumerate all possible actions. This is used by planning algorithms that need to iterate over actions.
- hash_action(action)[source]
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 thenp.array_equalsemantics used by the linear-scan fallback).
- initial_observation_dist()[source]
Get the initial observation distribution.
- Return type:
- Returns:
Distribution over initial observations
Note
Subclasses must implement this method to define initial observations.
- initial_state_dist()[source]
Get the initial state distribution.
- Return type:
- Returns:
Distribution over initial states
Note
Subclasses must implement this method to define the starting distribution.
- is_equal_observation(observation1, observation2)[source]
Check if two observations are equal.
- Parameters:
- Return type:
- Returns:
True if observations are considered equal, False otherwise
Note
Subclasses must implement this method to define observation equality. This is particularly important for discrete observation spaces.
- is_terminal(state)[source]
Check if a state is terminal.
- Parameters:
state (
ndarray) – State to check for terminal condition- Return type:
- Returns:
True if the state is terminal, False otherwise
Note
Subclasses must implement this method to define terminal conditions.
- observation_log_probability(next_state, action, observations)[source]
Log-probability of each candidate observation under
(next_state, action).Returns
np.ndarrayof shape(N,)where N is the number of candidate observations. Subclasses must implement.
- observation_log_probability_per_state(next_states, action, observation)[source]
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 exposesbatch_log_likelihood(next_states_array, observation_array)) should override to avoid the loop.
- reward(state, action, next_state=None)[source]
Calculate the immediate reward for a state-action(-next_state) tuple.
next_stateis the realised post-transition state when known (e.g. threaded bysample_next_step()), allowing rewards that depend on stochastic transition outcomes to use the same draw as the trajectory instead of resampling. Subclasses whose reward is a pure function of(state, action)may ignore it; subclasses whose reward depends on the realised next state (collision penalties, win bonuses) should consume it when provided and fall back to drawing/computing one whenNone.- Parameters:
- Return type:
- Returns:
Immediate reward value.
Note
Subclasses must implement this method to define reward structure.
- reward_batch(states, action, next_states=None)[source]
Calculate rewards for a batch of states given a single action.
Provides a loop-based default that subclasses can override with vectorized numpy implementations for better performance.
- Parameters:
- Return type:
- Returns:
1-D array of reward values with shape
(N,).
- sample_next_state(state, action, n_samples=1)[source]
Sample one or more next states for
(state, action).Hot-path entry point used by MCTS planners and particle filters. Subclasses must implement.
- sample_next_state_batch(states, action)[source]
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 exposesbatch_sample(states_array)) should override to avoid the loop.
- sample_next_step(state, action)[source]
Sample a complete state transition step.
This convenience method combines state transition, observation generation, and reward calculation in a single operation.
- sample_observation(next_state, action, n_samples=1)[source]
Sample one or more observations for
(next_state, action).Hot-path entry point used by MCTS planners and particle filters. Subclasses must implement.
- simulate_random_rollout(state, action_sampler, max_depth, discount_factor, depth=0)[source]
Random rollout via native C++ physics and reward kernel.
Pre-draws action indices and runs the full rollout in a single C++ call, avoiding per-step Python dispatch.
- Parameters:
state (
Any) – Current 4-D state[px, py, vx, vy].action_sampler (
Any) – Accepted for interface compatibility with the basesimulate_random_rolloutsignature; the native rollout draws action indices vianp.random.randintdirectly and never invokes the sampler.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:
- Returns:
Discounted sum of immediate rewards along the sampled trajectory.
- class POMDPPlanners.environments.safety_ant_velocity_pomdp.SafeAntVelocityVisualizer(env)[source]
Bases:
objectVisualizer for Safety Ant Velocity POMDP episodes.
This class creates animated visualizations showing the ant’s movement trajectory, velocity vectors, force applications, safety zones, and safety constraint violations.
- env
The SafeAntVelocityPOMDP environment instance
- safe_velocity_threshold
Maximum safe velocity magnitude
- max_force
Maximum force that can be applied
- create_animation(history, cache_path)[source]
Create animated visualization of the safety ant velocity episode.
Creates an animated GIF showing the ant’s movement trajectory with velocity vectors, safety zones, force applications, and safety constraint violations.
- Parameters:
- Raises:
ValueError – If history is empty or cache_path doesn’t end with .gif
TypeError – If cache_path is not a Path object
- Return type:
Subpackages
- POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_pomdp_beliefs package
SafetyAntVelocityVectorizedUpdaterSafetyAntVelocityVectorizedUpdater.obs_distSafetyAntVelocityVectorizedUpdater.dtSafetyAntVelocityVectorizedUpdater.massSafetyAntVelocityVectorizedUpdater.dampingSafetyAntVelocityVectorizedUpdater.max_forceSafetyAntVelocityVectorizedUpdater.force_scalesSafetyAntVelocityVectorizedUpdater.batch_observation_log_likelihood()SafetyAntVelocityVectorizedUpdater.batch_transition()SafetyAntVelocityVectorizedUpdater.config_idSafetyAntVelocityVectorizedUpdater.from_environment()
create_safety_ant_velocity_belief()- Submodules
- POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_pomdp_beliefs.safety_ant_velocity_belief_factory module
- POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_pomdp_beliefs.safety_ant_velocity_vectorized_updater module
SafetyAntVelocityVectorizedUpdaterSafetyAntVelocityVectorizedUpdater.obs_distSafetyAntVelocityVectorizedUpdater.dtSafetyAntVelocityVectorizedUpdater.massSafetyAntVelocityVectorizedUpdater.dampingSafetyAntVelocityVectorizedUpdater.max_forceSafetyAntVelocityVectorizedUpdater.force_scalesSafetyAntVelocityVectorizedUpdater.batch_observation_log_likelihood()SafetyAntVelocityVectorizedUpdater.batch_transition()SafetyAntVelocityVectorizedUpdater.config_idSafetyAntVelocityVectorizedUpdater.from_environment()
Submodules
POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_pomdp module
Safety Ant Velocity POMDP Environment Implementation.
This module implements a safety-critical velocity control task where an agent must navigate while avoiding unsafe velocities. The challenge is balancing exploration and movement rewards with safety constraints under partial observability.
The Safety Ant Velocity POMDP features: - Continuous 4D state space: [position_x, position_y, velocity_x, velocity_y] - Discrete action space: [0 (no force), 1 (small), 2 (medium), 3 (large force)] - Physics-based dynamics with force application and damping - Noisy observations of both position and velocity - Safety constraints on maximum velocity magnitude - Safety-focused metrics tracking violation rates
Key aspects: - Rewards encourage movement but heavily penalize safety violations - Episode terminates if velocity becomes critically high - Force direction is randomized to create uncertainty - Safety metrics track violation rates over episodes
- Classes:
SafeAntVelocityPOMDP: Main safety-critical velocity control environment
- class POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_pomdp.SafeAntVelocityPOMDP(discount_factor, safe_velocity_threshold=2.0, max_force=1.0, dt=0.1, mass=1.0, damping=0.1, position_noise=0.1, velocity_noise=0.2, safety_violation_penalty=-100.0, movement_reward_scale=1.0, name='SafeVelocityPOMDP', output_dir=None, debug=False, use_queue_logger=False)[source]
Bases:
DiscreteActionsEnvironmentSafety-critical velocity control task formulated as a POMDP.
This environment presents a safety-critical control problem where an agent must navigate while keeping velocity below a safety threshold. The challenge comes from balancing exploration rewards with safety constraints under noisy velocity observations.
Problem Structure: - State: [position_x, position_y, velocity_x, velocity_y] (continuous) - Actions: [0=no force, 1=small, 2=medium, 3=large force] (discrete) - Observations: Noisy position and velocity measurements (continuous) - Rewards: Movement reward - safety violation penalty (if unsafe) - Safety constraint: velocity magnitude ≤ safe_velocity_threshold - Termination: Velocity exceeds 1.5x safety threshold
Safety Features: - Tracks safety and critical violation rates - Heavy penalties for constraint violations - Configurable safety thresholds and penalties - Physics simulation with uncertainty in force direction
Example
>>> import numpy as np >>> np.random.seed(42) # For reproducible results >>> >>> # Initialize environment >>> env = SafeAntVelocityPOMDP(discount_factor=0.99) >>> >>> # 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
- Parameters:
- cache_visualization(history, output_dir, episode_index)[source]
Cache animated visualization of the safety ant velocity episode.
Creates an animated GIF showing the ant’s movement trajectory with velocity vectors, safety zones, force applications, and safety constraint violations.
- compute_metrics(histories)[source]
Compute environment-specific metrics from episode histories.
This method can be overridden by subclasses to provide custom metric calculations beyond standard return and episode length.
- get_actions()[source]
Get all possible actions in the discrete action space.
Note
Subclasses must implement this method to enumerate all possible actions. This is used by planning algorithms that need to iterate over actions.
- hash_action(action)[source]
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 thenp.array_equalsemantics used by the linear-scan fallback).
- initial_observation_dist()[source]
Get the initial observation distribution.
- Return type:
- Returns:
Distribution over initial observations
Note
Subclasses must implement this method to define initial observations.
- initial_state_dist()[source]
Get the initial state distribution.
- Return type:
- Returns:
Distribution over initial states
Note
Subclasses must implement this method to define the starting distribution.
- is_equal_observation(observation1, observation2)[source]
Check if two observations are equal.
- Parameters:
- Return type:
- Returns:
True if observations are considered equal, False otherwise
Note
Subclasses must implement this method to define observation equality. This is particularly important for discrete observation spaces.
- is_terminal(state)[source]
Check if a state is terminal.
- Parameters:
state (
ndarray) – State to check for terminal condition- Return type:
- Returns:
True if the state is terminal, False otherwise
Note
Subclasses must implement this method to define terminal conditions.
- observation_log_probability(next_state, action, observations)[source]
Log-probability of each candidate observation under
(next_state, action).Returns
np.ndarrayof shape(N,)where N is the number of candidate observations. Subclasses must implement.
- observation_log_probability_per_state(next_states, action, observation)[source]
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 exposesbatch_log_likelihood(next_states_array, observation_array)) should override to avoid the loop.
- reward(state, action, next_state=None)[source]
Calculate the immediate reward for a state-action(-next_state) tuple.
next_stateis the realised post-transition state when known (e.g. threaded bysample_next_step()), allowing rewards that depend on stochastic transition outcomes to use the same draw as the trajectory instead of resampling. Subclasses whose reward is a pure function of(state, action)may ignore it; subclasses whose reward depends on the realised next state (collision penalties, win bonuses) should consume it when provided and fall back to drawing/computing one whenNone.- Parameters:
- Return type:
- Returns:
Immediate reward value.
Note
Subclasses must implement this method to define reward structure.
- reward_batch(states, action, next_states=None)[source]
Calculate rewards for a batch of states given a single action.
Provides a loop-based default that subclasses can override with vectorized numpy implementations for better performance.
- Parameters:
- Return type:
- Returns:
1-D array of reward values with shape
(N,).
- sample_next_state(state, action, n_samples=1)[source]
Sample one or more next states for
(state, action).Hot-path entry point used by MCTS planners and particle filters. Subclasses must implement.
- sample_next_state_batch(states, action)[source]
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 exposesbatch_sample(states_array)) should override to avoid the loop.
- sample_next_step(state, action)[source]
Sample a complete state transition step.
This convenience method combines state transition, observation generation, and reward calculation in a single operation.
- sample_observation(next_state, action, n_samples=1)[source]
Sample one or more observations for
(next_state, action).Hot-path entry point used by MCTS planners and particle filters. Subclasses must implement.
- simulate_random_rollout(state, action_sampler, max_depth, discount_factor, depth=0)[source]
Random rollout via native C++ physics and reward kernel.
Pre-draws action indices and runs the full rollout in a single C++ call, avoiding per-step Python dispatch.
- Parameters:
state (
Any) – Current 4-D state[px, py, vx, vy].action_sampler (
Any) – Accepted for interface compatibility with the basesimulate_random_rolloutsignature; the native rollout draws action indices vianp.random.randintdirectly and never invokes the sampler.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:
- Returns:
Discounted sum of immediate rewards along the sampled trajectory.
- class POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_pomdp.SafeAntVelocityPOMDPMetrics(*values)[source]
Bases:
EnumMetric names for Safety Ant Velocity POMDP environment.
- CRITICAL_VIOLATION_RATE = 'critical_violation_rate'
- SAFETY_VIOLATION_RATE = 'safety_violation_rate'
- TOTAL_CRITICAL_VIOLATIONS = 'total_critical_violations'
- TOTAL_SAFETY_VIOLATIONS = 'total_safety_violations'
POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_vectorized_model module
Torch, on-device vectorized generative model for Safety Ant Velocity.
This module provides SafetyAntVelocityVectorizedModel, a fully
batched, GPU-friendly implementation of
VectorizedGenerativeModel
for SafeAntVelocityPOMDP.
It re-expresses the environment’s damped-force transition (a deterministic integration step driven by a uniformly-sampled force direction), its identity-mean diagonal-Gaussian observation model, and its speed-based reward and terminal rules as torch tensor kernels so a vectorized planner (VOPP) can run tens of thousands of parallel simulations on the GPU without a host/device sync. Every constant (time step, mass, damping, max force, force scales, noise standard deviations, safety threshold, penalty, movement scale) is read from a live environment instance, so the environment stays the single source of truth for configuration; only the numeric kernels are duplicated in torch. The accompanying parity test pins these kernels to the environment’s native (C++) implementations.
The Safety Ant Velocity action space is already discrete – four force levels
indexed 0..3 – so actions pass straight through as integer indices into
the environment’s DEFAULT_FORCE_SCALES table; no preset representative
action table is synthesized. The only supported configuration is one whose
discrete action set is exactly range(len(force_scales)) (so that action
index i selects force_scales[i]); any other action set raises
NotImplementedError.
- class POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_vectorized_model.SafetyAntVelocityVectorizedModel(env, *, device=None, dtype=torch.float32, observation_resolution=0.1)[source]
Bases:
objectFully vectorized torch generative model for the Safety Ant Velocity POMDP.
The model batches the transition, observation, reward, terminal, and observation-likelihood kernels over a leading particle dimension and keeps every tensor on a single device. States and observations are the 4-D vector
[position_x, position_y, velocity_x, velocity_y]. Actions are integer indices0..num_actions-1into the environment’s force-scale table.- Parameters:
env (SafeAntVelocityPOMDP)
device (torch.device | None)
dtype (torch.dtype)
observation_resolution (float)
- device
Device every tensor argument and return value lives on.
- dtype
Floating dtype used for state / observation / reward tensors.
- num_actions
Number of discrete force levels (rows in the force table).
Example
>>> import torch >>> from POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_pomdp import ( ... SafeAntVelocityPOMDP, ... ) >>> from POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_vectorized_model import ( ... SafetyAntVelocityVectorizedModel, ... ) >>> torch.manual_seed(0) <torch._C.Generator object at ...> >>> env = SafeAntVelocityPOMDP(discount_factor=0.99) >>> model = SafetyAntVelocityVectorizedModel(env, device=torch.device("cpu")) >>> states = torch.tensor([[0.0, 0.0, 0.5, 0.0], [1.0, 0.0, 0.0, 0.3]]) >>> actions = torch.tensor([2, 2]) >>> next_states = model.sample_next_states(states, actions) >>> observations = model.sample_observations(next_states, actions) >>> rewards = model.rewards(states, actions, next_states) >>> tuple(next_states.shape), tuple(observations.shape), tuple(rewards.shape) ((2, 4), (2, 4), (2,))
- observation_log_probs(next_states, actions, observations)[source]
- Return type:
Tensor- Parameters:
next_states (torch.Tensor)
actions (torch.Tensor)
observations (torch.Tensor)
- rewards(states, actions, next_states)[source]
- Return type:
Tensor- Parameters:
states (torch.Tensor)
actions (torch.Tensor)
next_states (torch.Tensor)
- sample_next_states(states, actions)[source]
- Return type:
Tensor- Parameters:
states (torch.Tensor)
actions (torch.Tensor)
POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_visualizer module
Visualization utilities for Safety Ant Velocity POMDP Environment.
This module provides visualization capabilities for the Safety Ant Velocity POMDP, creating animated GIF visualizations of episode trajectories with safety zones, velocity vectors, and safety constraint violations.
- Classes:
SafeAntVelocityVisualizer: Creates animated visualizations of episodes
- class POMDPPlanners.environments.safety_ant_velocity_pomdp.safety_ant_velocity_visualizer.SafeAntVelocityVisualizer(env)[source]
Bases:
objectVisualizer for Safety Ant Velocity POMDP episodes.
This class creates animated visualizations showing the ant’s movement trajectory, velocity vectors, force applications, safety zones, and safety constraint violations.
- env
The SafeAntVelocityPOMDP environment instance
- safe_velocity_threshold
Maximum safe velocity magnitude
- max_force
Maximum force that can be applied
- create_animation(history, cache_path)[source]
Create animated visualization of the safety ant velocity episode.
Creates an animated GIF showing the ant’s movement trajectory with velocity vectors, safety zones, force applications, and safety constraint violations.
- Parameters:
- Raises:
ValueError – If history is empty or cache_path doesn’t end with .gif
TypeError – If cache_path is not a Path object
- Return type: