POMDPPlanners.environments.isaac_lab_pomdp package
IsaacLab POMDP wrapper environment module.
This module provides a forward-only adapter exposing a registered IsaacLab task as a ground-truth world for the POMDPPlanners episode loop.
- Classes:
IsaacLabPOMDP: Forward-only adapter exposing an IsaacLab task as a world. IsaacLabPOMDPVisualizer: RGB-frame-to-
.mp4video writer for episodes. IsaacLabModelPOMDP: Planner-side generative model POMCPOW searches inside. GaussianObservationModel: Additive-Normal observation model (obs = state + noise). TransitionModel: Interface for a state-transition model. GaussianRandomWalkTransition: Action-ignoring Gaussian random-walk transition. LinearGaussianTransition: Fit-from-data linear-Gaussian action-conditioned transition. RewardModel: Interface for a reward model. LinearRewardModel: Fit-from-data linear reward model POMCPOW optimizes.
- class POMDPPlanners.environments.isaac_lab_pomdp.GaussianObservationModel(observation_dim, noise_std=0.1)[source]
Bases:
objectAdditive-Normal observation model:
observation = state + N(0, Sigma).A fixed-covariance multivariate normal whose mean is the state, mirroring the continuous light-dark
NORMAL_NOISEmodel. The covariance is diagonal and parameterized per channel, so proprioceptive channels can carry tight noise and exteroceptive ones looser noise without any per-task code.- dim
Dimensionality of the observation/state vector.
Example
Sample and score an observation for a 4-D state:
model = GaussianObservationModel(observation_dim=4, noise_std=0.1) obs = model.sample([0.0, 1.0, 2.0, 3.0]) log_p = model.log_probability([0.0, 1.0, 2.0, 3.0], obs)
- log_probability(state, observations)[source]
Gaussian log-density of
observationscentered onstate.
- class POMDPPlanners.environments.isaac_lab_pomdp.GaussianRandomWalkTransition(dim, process_noise_std=0.05)[source]
Bases:
TransitionModelAction-ignoring Gaussian random walk:
next_state = state + N(0, Sigma).The trivial placeholder transition. Because it ignores the action, a planner’s lookahead cannot distinguish actions under it — use it only for wiring/tests or when a learned transition is not yet available.
Example
Draw a next state near the current one:
transition = GaussianRandomWalkTransition(dim=4, process_noise_std=0.05) nxt = transition.sample_next_state([0.0, 1.0, 2.0, 3.0], action=None)
- log_probability(state, action, next_states)[source]
Log-density of each of
next_statesunder the transition for(state, action).
- class POMDPPlanners.environments.isaac_lab_pomdp.IsaacLabModelPOMDP(observation_dim, action_presets, discount_factor, observation_noise_std=0.1, process_noise_std=0.05, transition=None, reward_model=None, name=None)[source]
Bases:
DiscreteActionsEnvironmentDiscrete-action generative model POMCPOW searches inside for IsaacLab.
State and observation share one space; the observation is the state seen through
GaussianObservationModel, and the transition is anyTransitionModel(defaulting toGaussianRandomWalkTransition). Actions are a finite set of continuous control vectors applied verbatim to the world.- Parameters:
- observation_dim
Dimensionality of the shared state/observation vector.
- action_presets
The finite set of action vectors the planner chooses among.
Example
Build a model over a 4-D observation with three 2-D action presets:
import numpy as np presets = [np.zeros(2), np.ones(2), -np.ones(2)] model = IsaacLabModelPOMDP(observation_dim=4, action_presets=presets, discount_factor=0.99) actions = model.get_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).
- hash_observation(observation)[source]
Return a hashable key consistent with
is_equal_observation().Used by tree-search planners to index belief children by observation in O(1). The returned key MUST satisfy the contract:
is_equal_observation(a, b) implies hash_observation(a) == hash_observation(b)
- Parameters:
observation (
Any) – Observation to hash.- Returns:
the observation itself when it is already hashable).
- Return type:
- Raises:
NotImplementedError – If the observation is not hashable and the subclass has not provided an override. Subclasses with non-hashable observations (e.g.
np.ndarray) MUST override.
- 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 (
Any) – 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.
- 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.
- 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_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.
- class POMDPPlanners.environments.isaac_lab_pomdp.IsaacLabPOMDP(task_id, discount_factor, num_envs=1, device='cuda', env_cfg_kwargs=None, state_asset='robot', observation_sensor='lidar', state_extractor=None, observation_extractor=None, action_space_type=SpaceType.CONTINUOUS, observation_space_type=SpaceType.CONTINUOUS, headless=True, render_mode=None, seed=None, name=None, reward_range=None, output_dir=None, debug=False, use_queue_logger=False)[source]
Bases:
EnvironmentForward-only adapter exposing an IsaacLab task as a world POMDP.
The wrapper drives a registered IsaacLab task as the ground-truth world of an episode. It calls
env.step(action)exactly once per real interaction and serves the resulting next state, observation and reward from a small cache, because the POMDPPlanners episode loop requests those three quantities through separate method calls while IsaacLab produces them atomically. The state is read from the physics engine (scene[state_asset].data) and the observation from a sensor (scene[observation_sensor].data), giving a genuineobservation = h(state)split.Note
This is a world environment, not a generative model. It cannot sample a transition from an arbitrary state, so belief particle propagation and density queries are unsupported and raise
NotImplementedError/RuntimeError. Pair it with a generative model environment on the planner (policy.environment).- Parameters:
task_id (str)
discount_factor (float)
num_envs (int)
device (str)
state_asset (str)
observation_sensor (str)
action_space_type (SpaceType)
observation_space_type (SpaceType)
headless (bool)
render_mode (str | None)
seed (int | None)
name (str | None)
output_dir (Path | None)
debug (bool)
use_queue_logger (bool)
- task_id
Registered IsaacLab task id passed to
gymnasium.make.
- num_envs
Number of parallel sim envs; must be 1 for a world adapter.
- device
Torch device the simulator runs on (e.g.
"cuda").
- env_cfg_kwargs
Extra keyword arguments forwarded to
parse_env_cfg.
- state_asset
Scene key for the ground-truth articulation.
- observation_sensor
Scene key for the observation sensor.
- headless
Whether to launch the simulator without a GUI.
- seed
Optional seed applied to the first
resetfor reproducibility.
Example
Constructed like a Gymnasium world but reading state and sensor buffers directly from the IsaacLab scene (illustrative — requires a working Isaac Sim install, so it is not run as a doctest):
from POMDPPlanners.environments.isaac_lab_pomdp import IsaacLabPOMDP world = IsaacLabPOMDP( task_id="Isaac-Velocity-Flat-Anymal-C-v0", discount_factor=0.99, observation_sensor="lidar", ) state = world.initial_state_dist().sample()[0] action = world.space_info # supply a valid task action here next_state, observation, reward = world.sample_next_step(state, action) terminal = world.is_terminal(next_state)
- 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).
- hash_observation(observation)[source]
Return a hashable key consistent with
is_equal_observation().Used by tree-search planners to index belief children by observation in O(1). The returned key MUST satisfy the contract:
is_equal_observation(a, b) implies hash_observation(a) == hash_observation(b)
- Parameters:
observation (
Any) – Observation to hash.- Returns:
the observation itself when it is already hashable).
- Return type:
- Raises:
NotImplementedError – If the observation is not hashable and the subclass has not provided an override. Subclasses with non-hashable observations (e.g.
np.ndarray) MUST override.
- 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 (
Any) – 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.
- render()[source]
Return the current simulator viewport as an RGB frame.
Renders what the simulator sees (
env.render()) as an(H, W, 3)uint8array, suitable for assembling into a video withIsaacLabPOMDPVisualizer.- Return type:
- Returns:
An
(H, W, 3)uint8RGB frame of the simulator viewport.- Raises:
RuntimeError – If the world was not constructed with
render_mode="rgb_array"or the simulator returns no frame.
- 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.
- 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_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.
- class POMDPPlanners.environments.isaac_lab_pomdp.IsaacLabPOMDPVisualizer(environment=None)[source]
Bases:
objectRGB-frame-to-
.mp4video writer for an IsaacLabPOMDP episode.The visualizer needs only the list of RGB frames captured while an episode was rolled forward; it stores the environment for parity with the other environment visualizers but does not require any geometry from it.
Example
Rendering is driven from the RGB frames captured during an episode:
visualizer = IsaacLabPOMDPVisualizer(world) visualizer.frames_to_video(frames, Path("isaac_episode.mp4"))
- Parameters:
environment (Any)
- frames_to_video(frames, cache_path, fps=10)[source]
Write a sequence of RGB frames to an
.mp4video.- Parameters:
- Raises:
TypeError – If
framesis not a list orcache_pathis not a Path.ValueError – If
framesis empty orcache_pathdoes not end with.mp4.
- Return type:
- class POMDPPlanners.environments.isaac_lab_pomdp.IsaacLabSimulatorTransition(task_id, dim, state_writer, num_envs=1, device='cuda', env_cfg_kwargs=None, headless=True, process_noise_std=0.01, state_reader=None, state_asset='robot', action_space_type=SpaceType.CONTINUOUS)[source]
Bases:
TransitionModelState-transition model that steps the IsaacLab simulator itself.
next_state = f_sim(state, action) + N(0, Sigma): the flat state vector is written back into the physics engine, the sim is advanced one step underaction, and the resulting state is read out as the transition mean. A small diagonal process noise is added so the transition is a proper, samplable density — the bare physics step is a deterministic point mass with no density for the belief filter to weight against.Unlike
LinearGaussianTransition, this is the true nonlinear dynamics rather than a linear fit, so POMCPOW plans against exact contact physics. The cost is one GPU sim step per query.- Constraints:
One SimulationApp per process. This transition builds its own batched IsaacLab env through the same launch singleton the world uses, so it cannot coexist with an IsaacLab world in the same process — pair it with a non-IsaacLab world (real robot or a cheaper simulator).
State writer is task-specific. Writing a flat state back into the sim is the inverse of the world’s state extractor and depends on the asset/frame conventions, so it is injected via
state_writerrather than guessed.
- Parameters:
- dim
Dimensionality of the shared state/observation vector.
- num_envs
Number of parallel sim envs the batched model env is built with.
Example
Wrap a task, supplying the inverse of the world’s state extractor:
def write_state(env, states): # states: (batch, dim) env.unwrapped.write_root_state_to_sim(to_torch(states)) transition = IsaacLabSimulatorTransition( task_id="Isaac-Reach-Franka-v0", dim=18, state_writer=write_state, device="cpu", process_noise_std=0.01, ) nxt = transition.sample_next_state(state, action)
- log_probability(state, action, next_states)[source]
Log-density of each of
next_statesunder the transition for(state, action).
- class POMDPPlanners.environments.isaac_lab_pomdp.LinearGaussianTransition(weight_state, weight_action, bias, covariance)[source]
Bases:
TransitionModelLearned linear-Gaussian transition
next = A @ state + B @ action + b + N(0, Sigma).A first-order, action-conditioned dynamics model whose parameters are fit from
(state, action, next_state)rollouts (seefit()). Analytic, so both sampling and log-density reuse the pre-factorized normal. It is deliberately a system-identification baseline — a linear approximation of the true nonlinear contact dynamics — but, unlike the random walk, it makes actions move the predicted state so POMCPOW’s planning is meaningful.- dim
State dimensionality.
- action_dim
Action dimensionality.
Example
Fit from rollouts, then sample:
transition = LinearGaussianTransition.fit(states, actions, next_states) nxt = transition.sample_next_state(states[0], actions[0])
- classmethod fit(states, actions, next_states, regularization=0.0001, min_variance=1e-06)[source]
Fit
A,B,band a diagonal residual covariance via ridge least squares.- Parameters:
states (
ndarray) – Array of shape(N, dim)of source states.actions (
ndarray) – Array of shape(N, action_dim)of applied actions.next_states (
ndarray) – Array of shape(N, dim)of resulting states.regularization (
float) – Ridge penalty added to the normal-equations diagonal.min_variance (
float) – Floor on each residual variance to keep the covariance positive definite.
- Return type:
- Returns:
A fitted
LinearGaussianTransition.- Raises:
ValueError – If fewer than two transitions are supplied or shapes disagree.
- log_probability(state, action, next_states)[source]
Log-density of each of
next_statesunder the transition for(state, action).
- class POMDPPlanners.environments.isaac_lab_pomdp.LinearRewardModel(weight_state, weight_action, weight_next_state, bias)[source]
Bases:
RewardModelLearned linear reward
r = w_s . state + w_a . action + w_n . next_state + b.Fit from
(state, action, next_state, reward)rollouts via ridge least squares (seefit()). It is a first-order approximation of the task’s true (often nonlinear) reward, but it gives POMCPOW a real objective to optimize — without it the planner has no signal and produces undirected behavior.Example
Fit from rollouts, then score a transition:
reward_model = LinearRewardModel.fit(states, actions, next_states, rewards) value = reward_model.reward(states[0], actions[0], next_states[0])
- classmethod fit(states, actions, next_states, rewards, regularization=0.0001)[source]
Fit the reward coefficients via ridge least squares.
- Parameters:
states (
ndarray) – Array of shape(N, dim)of source states.actions (
ndarray) – Array of shape(N, action_dim)of applied actions.next_states (
ndarray) – Array of shape(N, dim)of resulting states.rewards (
ndarray) – Array of shape(N,)of observed rewards.regularization (
float) – Ridge penalty added to the normal-equations diagonal.
- Return type:
- Returns:
A fitted
LinearRewardModel.- Raises:
ValueError – If fewer than two transitions are supplied or shapes disagree.
- class POMDPPlanners.environments.isaac_lab_pomdp.RewardModel[source]
Bases:
ABCInterface for a reward model over the shared state/observation space.
A concrete reward model scores a
(state, action, next_state)transition. The planner-side model needs one because the forward-only world cannot be queried for the reward of a hypothetical in-tree transition.Note
This is an abstract base class and cannot be instantiated directly.
- class POMDPPlanners.environments.isaac_lab_pomdp.TransitionModel[source]
Bases:
ABCInterface for a state-transition model over the shared state/observation space.
A concrete transition supplies a generative next-state sampler and the matching log-density, both conditioned on
(state, action). Implementations here are analytic Gaussians; a learned world model can implement the same two methods.Note
This is an abstract base class and cannot be instantiated directly.
- abstractmethod log_probability(state, action, next_states)[source]
Log-density of each of
next_statesunder the transition for(state, action).
Submodules
POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp module
Planner-side generative model for the IsaacLab POMDP.
IsaacLabPOMDP
is a forward-only world: it steps a live IsaacLab task and emits the real
observation, but it cannot resample from an arbitrary state or score observation
densities. A belief filter and a planner such as POMCPOW therefore need a
separate generative model to search inside — this module provides it.
Design (see the project discussion that motivated it): rather than hand-splitting
each of IsaacLab’s ~200 tasks into privileged-state vs sensor-observation terms,
the model keeps state and observation in the same space and treats the
observation as the state seen through additive Gaussian sensor noise —
observation = state + N(0, Sigma). This makes the observation model generic
(one model for every task), trivially samplable from any state, and guarantees
the state explains the observation by construction. Truly unmeasurable constant
parameters (friction, added mass, actuator gains) are handled as domain
randomization on the world side, not as observation channels.
The Normal-noise mechanism reuses
CovarianceParameterizedMultivariateNormal
— the same fixed-covariance-varying-mean multivariate normal the continuous
light-dark POMDP uses for its NORMAL_NOISE observation model and its state
transition — so the Cholesky factorization is computed once and reused.
Transition model — a swappable seam. IsaacLab’s true contact dynamics are not
analytic, so the transition is injected via the TransitionModel
interface. Two concretes ship here:
GaussianRandomWalkTransition— the trivial, action-ignoring default (next = state + N(0, Sigma)); planning is cosmetic under it.LinearGaussianTransition— a first-order learned dynamics modelnext = A @ state + B @ action + b + N(0, Sigma)whose parameters are fit from(state, action, next_state)rollouts via ridge least squares. It is action-conditioned, so POMCPOW’s lookahead sees that actions move the state.
A learned, history-conditioned world model (RSSM/Dreamer) can be dropped in as a
third TransitionModel without touching the observation model or the
planner.
Reward model — the objective POMCPOW optimizes. The forward-only world cannot be
queried for the reward of a hypothetical in-tree transition, so the reward is also
injected, via RewardModel. LinearRewardModel is fit from the same
warm-up rollouts as the transition; without it the model’s reward is a flat zero
and the planner has no objective, so the task is never solved.
- Classes:
TransitionModel: Interface for a state-transition model (sample + log-density). GaussianRandomWalkTransition: Action-ignoring Gaussian random-walk transition. LinearGaussianTransition: Fit-from-data linear-Gaussian action-conditioned transition. IsaacLabSimulatorTransition: Steps the IsaacLab simulator as the true transition. RewardModel: Interface for a reward model over the state/observation space. LinearRewardModel: Fit-from-data linear reward model POMCPOW optimizes. GaussianObservationModel: Additive-Normal observation model (obs = state + noise). IsaacLabModelPOMDP: Discrete-action generative model POMCPOW searches inside.
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp.GaussianObservationModel(observation_dim, noise_std=0.1)[source]
Bases:
objectAdditive-Normal observation model:
observation = state + N(0, Sigma).A fixed-covariance multivariate normal whose mean is the state, mirroring the continuous light-dark
NORMAL_NOISEmodel. The covariance is diagonal and parameterized per channel, so proprioceptive channels can carry tight noise and exteroceptive ones looser noise without any per-task code.- dim
Dimensionality of the observation/state vector.
Example
Sample and score an observation for a 4-D state:
model = GaussianObservationModel(observation_dim=4, noise_std=0.1) obs = model.sample([0.0, 1.0, 2.0, 3.0]) log_p = model.log_probability([0.0, 1.0, 2.0, 3.0], obs)
- log_probability(state, observations)[source]
Gaussian log-density of
observationscentered onstate.
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp.GaussianRandomWalkTransition(dim, process_noise_std=0.05)[source]
Bases:
TransitionModelAction-ignoring Gaussian random walk:
next_state = state + N(0, Sigma).The trivial placeholder transition. Because it ignores the action, a planner’s lookahead cannot distinguish actions under it — use it only for wiring/tests or when a learned transition is not yet available.
Example
Draw a next state near the current one:
transition = GaussianRandomWalkTransition(dim=4, process_noise_std=0.05) nxt = transition.sample_next_state([0.0, 1.0, 2.0, 3.0], action=None)
- log_probability(state, action, next_states)[source]
Log-density of each of
next_statesunder the transition for(state, action).
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp.IsaacLabModelPOMDP(observation_dim, action_presets, discount_factor, observation_noise_std=0.1, process_noise_std=0.05, transition=None, reward_model=None, name=None)[source]
Bases:
DiscreteActionsEnvironmentDiscrete-action generative model POMCPOW searches inside for IsaacLab.
State and observation share one space; the observation is the state seen through
GaussianObservationModel, and the transition is anyTransitionModel(defaulting toGaussianRandomWalkTransition). Actions are a finite set of continuous control vectors applied verbatim to the world.- Parameters:
- observation_dim
Dimensionality of the shared state/observation vector.
- action_presets
The finite set of action vectors the planner chooses among.
Example
Build a model over a 4-D observation with three 2-D action presets:
import numpy as np presets = [np.zeros(2), np.ones(2), -np.ones(2)] model = IsaacLabModelPOMDP(observation_dim=4, action_presets=presets, discount_factor=0.99) actions = model.get_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).
- hash_observation(observation)[source]
Return a hashable key consistent with
is_equal_observation().Used by tree-search planners to index belief children by observation in O(1). The returned key MUST satisfy the contract:
is_equal_observation(a, b) implies hash_observation(a) == hash_observation(b)
- Parameters:
observation (
Any) – Observation to hash.- Returns:
the observation itself when it is already hashable).
- Return type:
- Raises:
NotImplementedError – If the observation is not hashable and the subclass has not provided an override. Subclasses with non-hashable observations (e.g.
np.ndarray) MUST override.
- 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 (
Any) – 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.
- 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.
- 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_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.
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp.IsaacLabSimulatorTransition(task_id, dim, state_writer, num_envs=1, device='cuda', env_cfg_kwargs=None, headless=True, process_noise_std=0.01, state_reader=None, state_asset='robot', action_space_type=SpaceType.CONTINUOUS)[source]
Bases:
TransitionModelState-transition model that steps the IsaacLab simulator itself.
next_state = f_sim(state, action) + N(0, Sigma): the flat state vector is written back into the physics engine, the sim is advanced one step underaction, and the resulting state is read out as the transition mean. A small diagonal process noise is added so the transition is a proper, samplable density — the bare physics step is a deterministic point mass with no density for the belief filter to weight against.Unlike
LinearGaussianTransition, this is the true nonlinear dynamics rather than a linear fit, so POMCPOW plans against exact contact physics. The cost is one GPU sim step per query.- Constraints:
One SimulationApp per process. This transition builds its own batched IsaacLab env through the same launch singleton the world uses, so it cannot coexist with an IsaacLab world in the same process — pair it with a non-IsaacLab world (real robot or a cheaper simulator).
State writer is task-specific. Writing a flat state back into the sim is the inverse of the world’s state extractor and depends on the asset/frame conventions, so it is injected via
state_writerrather than guessed.
- Parameters:
- dim
Dimensionality of the shared state/observation vector.
- num_envs
Number of parallel sim envs the batched model env is built with.
Example
Wrap a task, supplying the inverse of the world’s state extractor:
def write_state(env, states): # states: (batch, dim) env.unwrapped.write_root_state_to_sim(to_torch(states)) transition = IsaacLabSimulatorTransition( task_id="Isaac-Reach-Franka-v0", dim=18, state_writer=write_state, device="cpu", process_noise_std=0.01, ) nxt = transition.sample_next_state(state, action)
- log_probability(state, action, next_states)[source]
Log-density of each of
next_statesunder the transition for(state, action).
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp.LinearGaussianTransition(weight_state, weight_action, bias, covariance)[source]
Bases:
TransitionModelLearned linear-Gaussian transition
next = A @ state + B @ action + b + N(0, Sigma).A first-order, action-conditioned dynamics model whose parameters are fit from
(state, action, next_state)rollouts (seefit()). Analytic, so both sampling and log-density reuse the pre-factorized normal. It is deliberately a system-identification baseline — a linear approximation of the true nonlinear contact dynamics — but, unlike the random walk, it makes actions move the predicted state so POMCPOW’s planning is meaningful.- dim
State dimensionality.
- action_dim
Action dimensionality.
Example
Fit from rollouts, then sample:
transition = LinearGaussianTransition.fit(states, actions, next_states) nxt = transition.sample_next_state(states[0], actions[0])
- classmethod fit(states, actions, next_states, regularization=0.0001, min_variance=1e-06)[source]
Fit
A,B,band a diagonal residual covariance via ridge least squares.- Parameters:
states (
ndarray) – Array of shape(N, dim)of source states.actions (
ndarray) – Array of shape(N, action_dim)of applied actions.next_states (
ndarray) – Array of shape(N, dim)of resulting states.regularization (
float) – Ridge penalty added to the normal-equations diagonal.min_variance (
float) – Floor on each residual variance to keep the covariance positive definite.
- Return type:
- Returns:
A fitted
LinearGaussianTransition.- Raises:
ValueError – If fewer than two transitions are supplied or shapes disagree.
- log_probability(state, action, next_states)[source]
Log-density of each of
next_statesunder the transition for(state, action).
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp.LinearRewardModel(weight_state, weight_action, weight_next_state, bias)[source]
Bases:
RewardModelLearned linear reward
r = w_s . state + w_a . action + w_n . next_state + b.Fit from
(state, action, next_state, reward)rollouts via ridge least squares (seefit()). It is a first-order approximation of the task’s true (often nonlinear) reward, but it gives POMCPOW a real objective to optimize — without it the planner has no signal and produces undirected behavior.Example
Fit from rollouts, then score a transition:
reward_model = LinearRewardModel.fit(states, actions, next_states, rewards) value = reward_model.reward(states[0], actions[0], next_states[0])
- classmethod fit(states, actions, next_states, rewards, regularization=0.0001)[source]
Fit the reward coefficients via ridge least squares.
- Parameters:
states (
ndarray) – Array of shape(N, dim)of source states.actions (
ndarray) – Array of shape(N, action_dim)of applied actions.next_states (
ndarray) – Array of shape(N, dim)of resulting states.rewards (
ndarray) – Array of shape(N,)of observed rewards.regularization (
float) – Ridge penalty added to the normal-equations diagonal.
- Return type:
- Returns:
A fitted
LinearRewardModel.- Raises:
ValueError – If fewer than two transitions are supplied or shapes disagree.
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp.RewardModel[source]
Bases:
ABCInterface for a reward model over the shared state/observation space.
A concrete reward model scores a
(state, action, next_state)transition. The planner-side model needs one because the forward-only world cannot be queried for the reward of a hypothetical in-tree transition.Note
This is an abstract base class and cannot be instantiated directly.
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp.TransitionModel[source]
Bases:
ABCInterface for a state-transition model over the shared state/observation space.
A concrete transition supplies a generative next-state sampler and the matching log-density, both conditioned on
(state, action). Implementations here are analytic Gaussians; a learned world model can implement the same two methods.Note
This is an abstract base class and cannot be instantiated directly.
- abstractmethod log_probability(state, action, next_states)[source]
Log-density of each of
next_statesunder the transition for(state, action).
POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_pomdp module
IsaacLab POMDP world wrapper environment.
This module adapts a registered IsaacLab
task (Isaac-*-v0) to the POMDPPlanners
Environment interface so it can serve as
the ground-truth world in an
EpisodeRunner.
Like a Gymnasium env, an IsaacLab env is forward-only: it exposes
reset() / step(action) and is a black-box physics simulator with no
transition/observation density. It therefore cannot act as a planner’s
generative model, and IsaacLabPOMDP.transition_log_probability() /
IsaacLabPOMDP.observation_log_probability() intentionally raise
NotImplementedError. In the two-environment episode design the planner
keeps its own generative model (policy.environment) and this wrapper only
advances the single true state forward, one step per real interaction.
Unlike Gymnasium, IsaacLab gives a genuine observation = h(state) split: the
ground-truth robot pose/joint state is read directly from the physics engine
(scene[asset].data), while the observation is a simulated sensor buffer
(e.g. a RayCaster LiDAR). The two extractors are configurable per task.
One wrapper covers the many registered tasks because they share one entry point
(parse_env_cfg + gymnasium.make), one step API
((obs_dict, reward, terminated, truncated, info)) and one scene accessor
(env.unwrapped.scene[...]). Task-specific asset/sensor names become override
hooks (state_extractor / observation_extractor).
- Caveats:
One SimulationApp per process. IsaacLab launches a single global
SimulationApp; two IsaacLab envs cannot coexist in one process, so the world andpolicy.environmentcannot both be IsaacLab in-process, and the multiprocessing task managers cannot fork the GPU simulator. Drive single-processEpisodeRunnerruns.num_envs must be 1. A world drives one true trajectory; batched (particle) use is a model-side concern this world adapter does not cover.
Density methods are unsupported (forward-only world).
- Classes:
IsaacLabPOMDP: Forward-only adapter exposing an IsaacLab task as a world.
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_pomdp.IsaacLabPOMDP(task_id, discount_factor, num_envs=1, device='cuda', env_cfg_kwargs=None, state_asset='robot', observation_sensor='lidar', state_extractor=None, observation_extractor=None, action_space_type=SpaceType.CONTINUOUS, observation_space_type=SpaceType.CONTINUOUS, headless=True, render_mode=None, seed=None, name=None, reward_range=None, output_dir=None, debug=False, use_queue_logger=False)[source]
Bases:
EnvironmentForward-only adapter exposing an IsaacLab task as a world POMDP.
The wrapper drives a registered IsaacLab task as the ground-truth world of an episode. It calls
env.step(action)exactly once per real interaction and serves the resulting next state, observation and reward from a small cache, because the POMDPPlanners episode loop requests those three quantities through separate method calls while IsaacLab produces them atomically. The state is read from the physics engine (scene[state_asset].data) and the observation from a sensor (scene[observation_sensor].data), giving a genuineobservation = h(state)split.Note
This is a world environment, not a generative model. It cannot sample a transition from an arbitrary state, so belief particle propagation and density queries are unsupported and raise
NotImplementedError/RuntimeError. Pair it with a generative model environment on the planner (policy.environment).- Parameters:
task_id (str)
discount_factor (float)
num_envs (int)
device (str)
state_asset (str)
observation_sensor (str)
action_space_type (SpaceType)
observation_space_type (SpaceType)
headless (bool)
render_mode (str | None)
seed (int | None)
name (str | None)
output_dir (Path | None)
debug (bool)
use_queue_logger (bool)
- task_id
Registered IsaacLab task id passed to
gymnasium.make.
- num_envs
Number of parallel sim envs; must be 1 for a world adapter.
- device
Torch device the simulator runs on (e.g.
"cuda").
- env_cfg_kwargs
Extra keyword arguments forwarded to
parse_env_cfg.
- state_asset
Scene key for the ground-truth articulation.
- observation_sensor
Scene key for the observation sensor.
- headless
Whether to launch the simulator without a GUI.
- seed
Optional seed applied to the first
resetfor reproducibility.
Example
Constructed like a Gymnasium world but reading state and sensor buffers directly from the IsaacLab scene (illustrative — requires a working Isaac Sim install, so it is not run as a doctest):
from POMDPPlanners.environments.isaac_lab_pomdp import IsaacLabPOMDP world = IsaacLabPOMDP( task_id="Isaac-Velocity-Flat-Anymal-C-v0", discount_factor=0.99, observation_sensor="lidar", ) state = world.initial_state_dist().sample()[0] action = world.space_info # supply a valid task action here next_state, observation, reward = world.sample_next_step(state, action) terminal = world.is_terminal(next_state)
- 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).
- hash_observation(observation)[source]
Return a hashable key consistent with
is_equal_observation().Used by tree-search planners to index belief children by observation in O(1). The returned key MUST satisfy the contract:
is_equal_observation(a, b) implies hash_observation(a) == hash_observation(b)
- Parameters:
observation (
Any) – Observation to hash.- Returns:
the observation itself when it is already hashable).
- Return type:
- Raises:
NotImplementedError – If the observation is not hashable and the subclass has not provided an override. Subclasses with non-hashable observations (e.g.
np.ndarray) MUST override.
- 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 (
Any) – 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.
- render()[source]
Return the current simulator viewport as an RGB frame.
Renders what the simulator sees (
env.render()) as an(H, W, 3)uint8array, suitable for assembling into a video withIsaacLabPOMDPVisualizer.- Return type:
- Returns:
An
(H, W, 3)uint8RGB frame of the simulator viewport.- Raises:
RuntimeError – If the world was not constructed with
render_mode="rgb_array"or the simulator returns no frame.
- 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.
- 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_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.
POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_vectorized_model module
Torch, on-device vectorized generative model for the Isaac Lab planner model.
This module provides IsaacLabVectorizedModel, a fully batched,
GPU-friendly implementation of
VectorizedGenerativeModel
built from the fitted planner-side model of
POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp.
The Isaac Lab world is a physics simulator, so the model VOPP searches is not
the simulator itself but the linear-Gaussian surrogate fit from a warm-up of
random transitions: a LinearGaussianTransition, a
GaussianObservationModel, and a LinearRewardModel. Because
those three are already linear / Gaussian, they re-express exactly as batched
torch kernels:
transition
s' = A s + B a + b + N(0, Sigma_tr),observation
o = s' + N(0, Sigma_obs)(state and observation share one space, sods == do),reward
r = w_s . s + w_a . a + w_n . s' + b_r, anda terminal test that is always
False(the Isaac velocity tasks never terminate the model).
Actions are integer indices into a fixed preset table of continuous action
vectors (the VOPP representative-action assumption); the ground-truth Isaac
action for index i is action_presets[i]. An accompanying parity test
pins every kernel to the fitted numpy models.
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_vectorized_model.IsaacLabVectorizedModel(transition, observation_model, reward_model, action_presets, *, device=None, dtype=torch.float32, observation_resolution=0.5)[source]
Bases:
objectFully vectorized torch generative model for the Isaac Lab surrogate model.
The model batches the linear-Gaussian transition, additive-Gaussian observation, linear reward, always-false terminal, and observation log-likelihood kernels over a leading particle dimension, keeping every tensor on one device. Actions are integer indices into a fixed preset table of continuous Isaac actions.
- Parameters:
transition (LinearGaussianTransition)
observation_model (GaussianObservationModel)
reward_model (LinearRewardModel)
action_presets (ndarray)
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 preset actions (rows of the action table).
Example
>>> import numpy as np >>> import torch >>> from POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_model_pomdp import ( ... GaussianObservationModel, ... LinearGaussianTransition, ... LinearRewardModel, ... ) >>> transition = LinearGaussianTransition( ... np.eye(3), np.zeros((3, 2)), np.zeros(3), 0.01 * np.eye(3) ... ) >>> observation = GaussianObservationModel(3, noise_std=0.1) >>> reward = LinearRewardModel(np.zeros(3), np.zeros(2), np.ones(3), 0.0) >>> presets = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]]) >>> model = IsaacLabVectorizedModel( ... transition, observation, reward, presets, device=torch.device("cpu") ... ) >>> states = torch.zeros(4, 3) >>> actions = torch.tensor([0, 1, 2, 1]) >>> next_states = model.sample_next_states(states, actions) >>> tuple(next_states.shape), model.num_actions ((4, 3), 3)
- property action_vectors: torch.Tensor
The
[num_actions, action_dim]table of continuous action vectors.
- 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.isaac_lab_pomdp.isaac_lab_visualizer module
IsaacLab POMDP episode video visualizer.
Renders an episode as an .mp4 video of the simulator viewport: each frame is
the RGB image the IsaacLab simulator produced for that step (captured via
render()),
so the video shows what is seen in the simulator rather than an abstract plot
of the state or observation.
Video (rather than GIF) output is used because IsaacLab frames are full-colour,
high-resolution renders where an animated GIF would be large and heavily
quantized; .mp4 keeps the file small and the colours faithful. Encoding pipes
the raw RGB frames straight into the system ffmpeg binary (located via
matplotlib’s configured animation.ffmpeg_path), which adds no extra Python
dependency and avoids the per-frame matplotlib figure round-trip that dominates a
FuncAnimation-based encode.
- Classes:
IsaacLabPOMDPVisualizer: RGB-frame-to-
.mp4video writer for IsaacLabPOMDP.
- class POMDPPlanners.environments.isaac_lab_pomdp.isaac_lab_visualizer.IsaacLabPOMDPVisualizer(environment=None)[source]
Bases:
objectRGB-frame-to-
.mp4video writer for an IsaacLabPOMDP episode.The visualizer needs only the list of RGB frames captured while an episode was rolled forward; it stores the environment for parity with the other environment visualizers but does not require any geometry from it.
Example
Rendering is driven from the RGB frames captured during an episode:
visualizer = IsaacLabPOMDPVisualizer(world) visualizer.frames_to_video(frames, Path("isaac_episode.mp4"))
- Parameters:
environment (Any)
- frames_to_video(frames, cache_path, fps=10)[source]
Write a sequence of RGB frames to an
.mp4video.- Parameters:
- Raises:
TypeError – If
framesis not a list orcache_pathis not a Path.ValueError – If
framesis empty orcache_pathdoes not end with.mp4.
- Return type: