POMDPPlanners.environments.carla_pomdp.carla_generative_models package
Planner-side generative models paired with the forward-only CARLA world.
While carla_pomdp is the ground-truth
world (forward-only, no densities), a planner carries a generative model as
policy.environment. This subpackage holds that model: the abstract interface and its
concrete implementations, all sharing the CARLA state/observation schema defined by the
world.
- Classes:
CarlaModelPOMDP: Abstract generative-model interface over the CARLA schema. FactoredCarlaModelPOMDP: Concrete CARLA model with a factored observation model. KinematicCarlaModelPOMDP: Factored model with a kinematic bicycle transition. DreamerCarlaModelPOMDP: Concrete CARLA model backed by a Dreamer world model. DreamerWorldModel: Protocol a trained Dreamer RSSM must satisfy. CarDreamerWorldModel: DreamerV3-backed
DreamerWorldModelfrom a CarDreamer checkpoint.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.CarDreamerWorldModel(agent, action_dim=3, rng_seed=0)[source]
Bases:
objectTrained CarDreamer DreamerV3 world model exposed as a
DreamerWorldModel.Wraps a constructed DreamerV3 JAX agent (the object holding the fitted parameters in
agent.varibsand the world model inagent.agent.wm) and routes every protocol method onto its RSSM and prediction heads. Build one from a training checkpoint withfrom_checkpoint().- latent_dim
Width of a packed latent vector (
deterwidth + flattenedstochwidth), matching the flat state the planner carries.
Note
This class requires
jax,ninjax, and the CarDreamerdreamerv3package importable in the running environment. It is intentionally not unit-tested against a live model here; the framework-agnosticDreamerWorldModelprotocol is covered by a lightweight fake instead.- continue_prob(latents)[source]
Probability the episode continues for each of
(batch, latent_dim)latents.
- decode_log_prob(latents, observation)[source]
Log-density of one observation under each of
(batch, latent_dim)latents.
- encode(observation)[source]
Encode a real observation into a latent via the RSSM posterior (belief seed).
- classmethod from_checkpoint(checkpoint_path, obs_space, act_space, config_size='medium', config_updates=None, step=0, action_dim=3, rng_seed=0)[source]
Build the adapter from a CarDreamer/DreamerV3 training checkpoint.
Constructs the DreamerV3 config (defaults + the named size preset + any overrides), instantiates the agent over the given observation/action spaces, and restores the fitted parameters from
checkpoint_pathviaembodied.Checkpoint.- Parameters:
checkpoint_path (
str) – Path to a DreamerV3checkpoint.ckptwritten during training.obs_space (
Mapping[str,Any]) – The agent’s observation space ({name: embodied.Space}); must include the CARLA schema keysgnssandagents.act_space (
Mapping[str,Any]) – The agent’s action space ({name: embodied.Space}).config_size (
str) – DreamerV3 config size preset to load (e.g."small","medium","large"); must match the size the checkpoint was trained at.config_updates (
Optional[Mapping[str,Any]]) – Optional additional{"dreamerv3": {...}}config overrides.step (
int) – The environment step counter to seed the agent with.action_dim (
int) – Width of the control vector fed to the RSSM.rng_seed (
int) – Seed for the JAX PRNG driving the RSSM/head calls.
- Return type:
- Returns:
A
CarDreamerWorldModelwrapping the restored agent.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.CarlaModelPOMDP(discount_factor, action_presets=None, max_tracked_agents=5, observation_models=None, name=None)[source]
Bases:
DiscreteActionsEnvironmentAbstract generative-model interface paired with the forward-only CARLA world.
Concrete subclasses supply the dynamics —
sample_next_state(),transition_log_probability(),reward(),is_terminal(), and the initial-distribution hooks — while this base owns the fixed CARLA schema shared by every model (the ego + agent-slot state layout, the discrete action set, and the observation-dict hashing/equality) and the observation model. The observation is factored by channel: this base holds a{channel: CarlaObservationModel}map (observation_models) and composes it —sample_observation()perceives each channel of the clean observation built from the state,observation_log_probability()sums the per-channel densities, andencode_observation()(the single raw-observation seam) perceives each channel of the forward-only world’s raw reading into the same perceived space, so the belief filter and planner search operate on one consistent (encoded) observation. Two models differ in their observation only by the channel models they hold; a model that scores observations for a belief update needs every channel to provide a density (supports_density). A subclass whose observation is not a per-channel clean transform (e.g. a learned latent decoder) may instead overridesample_observation(),observation_log_probability()andencode_observation()directly.- Parameters:
- action_presets
Discrete
(throttle, steer, brake)control triples; the discrete action set is the indices into this list.
- max_tracked_agents
Number of fixed agent slots in the state/observation.
- observation_models
The per-channel
{channel: CarlaObservationModel}map composed to produce the observation, orNonefor a subclass that overrides the observation methods directly.
Note
This is an abstract base class and cannot be instantiated directly. See
FactoredCarlaModelPOMDPfor a concrete reference implementation.- encode_observation(observation)[source]
Perceive the world’s raw observation into the belief/planner observation space.
The forward-only world emits a raw, fully-detected observation; each channel model degrades its channel (per-slot detection gating, occlusion, additive noise) into the perceived observation the belief filter and planner search operate in. This is the single raw-observation seam — every other observation method works in the perceived (encoded) space. A subclass without channel models (
observation_models is None, e.g. a learned latent model) inherits the identity default.
- 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.
- 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.
- observation_log_probability(next_state, action, observations)[source]
Log-density of
observationsunder the per-channel perception givennext_state.
- abstractmethod sample_next_state(state, action, n_samples=1)[source]
Sample
n_samplesnext states for(state, action).
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.DreamerCarlaModelPOMDP(world_model, discount_factor, action_presets=None, max_tracked_agents=5, continue_threshold=0.5, initial_observation=None, name=None)[source]
Bases:
CarlaModelPOMDPConcrete CARLA generative model whose dynamics are a trained Dreamer world model.
The planner-side state is the Dreamer latent; transitions, observations, reward, and termination are served by the injected
DreamerWorldModel. The belief is seeded by encoding the world’s initial observation with the posterior.- Parameters:
- world_model
The trained Dreamer RSSM backing every dynamic quantity.
- continue_threshold
Termination fires when the continue head’s probability drops below this value.
Note
Reward comes from the world model’s learned reward head, not the analytic
driving_quality_reward(); a Dreamer model predicts reward directly from its latent.Example
>>> import numpy as np >>> >>> class _IdentityWorldModel: ... latent_dim = 4 ... def encode(self, observation): ... return np.zeros(self.latent_dim) ... def imagine(self, latents, controls): ... return np.asarray(latents, dtype=float) ... def decode(self, latents): ... batch = np.asarray(latents).shape[0] ... return {"gnss": np.zeros((batch, 3)), "agents": np.zeros((batch, 25))} ... def decode_log_prob(self, latents, observation): ... return np.zeros(np.asarray(latents).shape[0]) ... def reward(self, latents): ... return np.zeros(np.asarray(latents).shape[0]) ... def continue_prob(self, latents): ... return np.ones(np.asarray(latents).shape[0]) >>> >>> obs = {"gnss": np.zeros(3), "agents": np.zeros(25)} >>> env = DreamerCarlaModelPOMDP( ... _IdentityWorldModel(), discount_factor=0.95, initial_observation=obs) >>> >>> state = env.initial_state_dist().sample()[0] >>> action = env.get_actions()[0] >>> next_state, observation, reward = env.sample_next_step(state, action) >>> sorted(observation) ['agents', 'gnss'] >>> env.is_terminal(state) False
- 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_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-density of
observationsunder the per-channel perception givennext_state.
- 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.
- sample_next_state(state, action, n_samples=1)[source]
Sample
n_samplesnext states for(state, action).
- 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.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.DreamerWorldModel(*args, **kwargs)[source]
Bases:
ProtocolBatched operations a trained Dreamer RSSM must expose to back the CARLA model.
Every latent is a 1-D float vector of length
latent_dim(the packed deterministic + stochastic recurrent state). All methods are batched: they take a(batch, latent_dim)array of latents and return per-row results, so a single network call serves a whole particle set.- latent_dim
Width of a packed latent vector.
- continue_prob(latents)[source]
Probability the episode continues for each of
(batch, latent_dim)latents.
- decode_log_prob(latents, observation)[source]
Log-density of one observation under each of
(batch, latent_dim)latents.
- encode(observation)[source]
Encode a real observation into a latent via the posterior (belief seed).
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.FactoredCarlaModelPOMDP(discount_factor, action_presets=None, max_tracked_agents=5, perception_range=50.0, occlusion_radius=1.5, pose_std=0.5, gnss_std=1e-05, detect_prob=0.95, desired_speed=8.0, out_lane_thresh=2.0, collision_penalty=100.0, observation=None, name=None)[source]
Bases:
CarlaModelPOMDPConcrete CARLA generative model pairing placeholder dynamics with factored perception.
The observation is composed per channel (held as
self.observation_models): aFactoredAgentObservationModelonagents(detection gated by perception range and geometric occlusion, additive Gaussian pose noise) and aGnssObservationModelongnss. The reward is the shared gym-carla driving-quality reward. The transition is a documented identity placeholder to be replaced with real dynamics per study.- Parameters:
discount_factor (float)
action_presets (Sequence[Tuple[float, float, float]] | None)
max_tracked_agents (int)
perception_range (float | None)
occlusion_radius (float)
pose_std (float)
gnss_std (float)
detect_prob (float)
desired_speed (float)
out_lane_thresh (float)
collision_penalty (float)
name (str | None)
- observation_models
The per-channel
{channel: CarlaObservationModel}map carrying the observation parameters (perception_range,occlusion_radius,pose_std,gnss_std,detect_prob).
- desired_speed
Target longitudinal speed (m/s) used by the reward.
- out_lane_thresh
Lateral offset (m) beyond which the reward penalises out-of-lane.
- collision_penalty
Penalty scale applied on a terminal collision in the reward.
Example
>>> import numpy as np >>> np.random.seed(42) # For reproducible results >>> >>> from POMDPPlanners.environments.carla_pomdp.carla_pomdp import ( ... AGENT_SLOT_WIDTH, EGO_STATE_WIDTH) >>> env = FactoredCarlaModelPOMDP(discount_factor=0.95) >>> >>> width = EGO_STATE_WIDTH + env.max_tracked_agents * AGENT_SLOT_WIDTH >>> state = np.zeros(width) >>> action = env.get_actions()[0] >>> >>> next_state, observation, reward = env.sample_next_step(state, action) >>> sorted(observation) ['agents', 'gnss'] >>> env.is_terminal(state) False
- 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_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.
- 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.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.KinematicCarlaModelPOMDP(discount_factor, dt=0.05, action_presets=None, max_tracked_agents=5, wheelbase=2.8, max_steer_angle=0.6, accel=3.0, brake_decel=8.0, drag=0.05, collision_gap=5.0, collision_halfwidth=1.2, safe_distance=12.0, stop_gap=0.0, name=None, **kwargs)[source]
Bases:
FactoredCarlaModelPOMDPFactored CARLA model whose transition is a kinematic bicycle propagation.
- Parameters:
discount_factor (float)
dt (float)
action_presets (Sequence[Tuple[float, float, float]] | None)
max_tracked_agents (int)
wheelbase (float)
max_steer_angle (float)
accel (float)
brake_decel (float)
drag (float)
collision_gap (float)
collision_halfwidth (float)
safe_distance (float)
stop_gap (float)
name (str | None)
kwargs (Any)
- dt
Integration step (seconds); must match the world’s
fixed_delta_seconds.
- wheelbase
Bicycle-model wheelbase (m) mapping steer to yaw rate.
- max_steer_angle
Steering command of 1.0 maps to this front-wheel angle (rad).
- accel
Longitudinal acceleration per unit throttle (m/s^2).
- brake_decel
Longitudinal deceleration per unit brake (m/s^2).
- drag
Linear speed-proportional deceleration coefficient (1/s).
- collision_gap
Forward ego-frame distance (m) within which a present agent ahead is treated as a predicted collision by
is_terminal().
- collision_halfwidth
Lateral ego-frame half-corridor (m) within which a present agent ahead is treated as a predicted collision by
is_terminal().
- safe_distance
Lead gap (m) at/above which the reward targets the full
desired_speed; the obstacle-aware target ramps down below it.
- stop_gap
Lead gap (m) at/below which the obstacle-aware target speed is zero;
0.0keeps the flatdesired_speed.
Example
>>> import numpy as np >>> from POMDPPlanners.environments.carla_pomdp.carla_pomdp import ( ... AGENT_SLOT_WIDTH, EGO_STATE_WIDTH) >>> env = KinematicCarlaModelPOMDP(discount_factor=0.95, dt=0.05) >>> >>> width = EGO_STATE_WIDTH + env.max_tracked_agents * AGENT_SLOT_WIDTH >>> state = np.zeros(width) >>> throttle_action = 0 # (0.5, 0.0, 0.0) cruise straight >>> next_state = env.sample_next_state(state, throttle_action) >>> >>> bool(next_state[3] > 0.0) # throttle produced forward velocity True
- is_terminal(state)[source]
Whether a present agent occupies the ego’s footprint just ahead.
Because this model does predict ego and agent motion, it can foresee running into the vehicle ahead. Any present agent slot within
collision_gapmetres forward andcollision_halfwidthmetres laterally (ego frame) is treated as a collision, which the inherited reward turns into the terminalcollision_penalty(driving_quality_reward()).
- reward(state, action, next_state=None)[source]
Driving-quality reward whose target speed adapts to the nearest lead obstacle.
The bare parent reward tracks a fixed
desired_speedand only charges the collision once an agent is inside the terminal box — a cliff that cannot be braked for at speed, while a large fixed penalty instead freezes the ego in traffic. This override follows Roach’s obstacle-aware desired speed: the target equals the fulldesired_speedwhen the lead gap is at leastsafe_distance, ramps linearly to zero atstop_gap, and is zero closer in. The ego is thus rewarded for driving when the road is clear and for slowing as an obstacle nears — including the lidar/camera obstacle fused into the agent slots — without a separate penalty term that traps it at a standstill.stop_gap == 0keeps the flat parent behaviour.
- sample_next_state(state, action, n_samples=1)[source]
Sample
n_samplesnext states for(state, action).
- 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.
- POMDPPlanners.environments.carla_pomdp.carla_generative_models.build_cardreamer_model(checkpoint_path, obs_space, act_space, action_presets=None, **kwargs)[source]
Convenience wrapper around
CarDreamerWorldModel.from_checkpoint().Derives the RSSM action width from
action_presets(defaulting to the 3-wide CARLA control triple) and forwards the rest to the checkpoint loader.- Parameters:
checkpoint_path (
str) – Path to a DreamerV3 training checkpoint.obs_space (
Mapping[str,Any]) – The agent’s observation space.action_presets (
Optional[Sequence[Tuple[float,float,float]]]) – Discrete control triples; only their width (3) is used here.**kwargs (
Any) – Forwarded toCarDreamerWorldModel.from_checkpoint().
- Return type:
- Returns:
The constructed
CarDreamerWorldModel.
Submodules
POMDPPlanners.environments.carla_pomdp.carla_generative_models.cardreamer_world_model module
CarDreamer/DreamerV3 adapter for the DreamerWorldModel protocol.
DreamerCarlaModelPOMDP plans inside any object satisfying the
DreamerWorldModel
protocol. This module supplies one concrete backing: a trained DreamerV3 world model
from the CarDreamer project (https://github.com/ucd-dare/CarDreamer), whose RSSM is a
JAX/ninjax module.
The adapter bridges two representations:
The planner-side latent the POMDP carries is a flat 1-D float vector — this module packs it as
concat(deter, stoch.flatten())(the DreamerV3 recurrent state), and unpacks it back into the{deter, stoch}state dict the RSSM and heads consume.Each protocol method is batched over
(batch, latent_dim)and runs the relevant DreamerV3 component (encoder,rssm.obs_step/rssm.img_step, or a head) through a singleninjax.purecall, converting NumPy in and NumPy out.
JAX, ninjax, and dreamerv3 are imported lazily (inside the constructor and the
factory), so importing this module — and the example script that references it — never
requires the deep-learning stack. Only actually building a
CarDreamerWorldModel does.
- Precondition — schema alignment:
The trained checkpoint’s observation space must expose the CARLA schema keys this POMDP uses (
gnssandagents) with matching shapes, and its action space must accept the(throttle, steer, brake)control triple. A CarDreamer task configured with those observation handlers satisfies this by construction; a checkpoint trained on, e.g., birds-eye-view images does not and cannot be plugged in unchanged.- Classes:
CarDreamerWorldModel: DreamerV3-backed implementation of
DreamerWorldModel.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.cardreamer_world_model.CarDreamerWorldModel(agent, action_dim=3, rng_seed=0)[source]
Bases:
objectTrained CarDreamer DreamerV3 world model exposed as a
DreamerWorldModel.Wraps a constructed DreamerV3 JAX agent (the object holding the fitted parameters in
agent.varibsand the world model inagent.agent.wm) and routes every protocol method onto its RSSM and prediction heads. Build one from a training checkpoint withfrom_checkpoint().- latent_dim
Width of a packed latent vector (
deterwidth + flattenedstochwidth), matching the flat state the planner carries.
Note
This class requires
jax,ninjax, and the CarDreamerdreamerv3package importable in the running environment. It is intentionally not unit-tested against a live model here; the framework-agnosticDreamerWorldModelprotocol is covered by a lightweight fake instead.- continue_prob(latents)[source]
Probability the episode continues for each of
(batch, latent_dim)latents.
- decode_log_prob(latents, observation)[source]
Log-density of one observation under each of
(batch, latent_dim)latents.
- encode(observation)[source]
Encode a real observation into a latent via the RSSM posterior (belief seed).
- classmethod from_checkpoint(checkpoint_path, obs_space, act_space, config_size='medium', config_updates=None, step=0, action_dim=3, rng_seed=0)[source]
Build the adapter from a CarDreamer/DreamerV3 training checkpoint.
Constructs the DreamerV3 config (defaults + the named size preset + any overrides), instantiates the agent over the given observation/action spaces, and restores the fitted parameters from
checkpoint_pathviaembodied.Checkpoint.- Parameters:
checkpoint_path (
str) – Path to a DreamerV3checkpoint.ckptwritten during training.obs_space (
Mapping[str,Any]) – The agent’s observation space ({name: embodied.Space}); must include the CARLA schema keysgnssandagents.act_space (
Mapping[str,Any]) – The agent’s action space ({name: embodied.Space}).config_size (
str) – DreamerV3 config size preset to load (e.g."small","medium","large"); must match the size the checkpoint was trained at.config_updates (
Optional[Mapping[str,Any]]) – Optional additional{"dreamerv3": {...}}config overrides.step (
int) – The environment step counter to seed the agent with.action_dim (
int) – Width of the control vector fed to the RSSM.rng_seed (
int) – Seed for the JAX PRNG driving the RSSM/head calls.
- Return type:
- Returns:
A
CarDreamerWorldModelwrapping the restored agent.
- POMDPPlanners.environments.carla_pomdp.carla_generative_models.cardreamer_world_model.build_cardreamer_model(checkpoint_path, obs_space, act_space, action_presets=None, **kwargs)[source]
Convenience wrapper around
CarDreamerWorldModel.from_checkpoint().Derives the RSSM action width from
action_presets(defaulting to the 3-wide CARLA control triple) and forwards the rest to the checkpoint loader.- Parameters:
checkpoint_path (
str) – Path to a DreamerV3 training checkpoint.obs_space (
Mapping[str,Any]) – The agent’s observation space.action_presets (
Optional[Sequence[Tuple[float,float,float]]]) – Discrete control triples; only their width (3) is used here.**kwargs (
Any) – Forwarded toCarDreamerWorldModel.from_checkpoint().
- Return type:
- Returns:
The constructed
CarDreamerWorldModel.
POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_dreamer_model_pomdp module
Dreamer-backed concrete CARLA generative model.
DreamerCarlaModelPOMDP implements the
CarlaModelPOMDP
interface by delegating every dynamic quantity to a trained Dreamer world model (an
RSSM). The POMDP state carried through the planner is the Dreamer latent (the
packed deterministic + stochastic recurrent state); the interface methods map onto the
world model’s own components:
sample_next_state()-> RSSM imagination step (advance the recurrent state under the action’s control triple and sample the stochastic prior).sample_observation()-> decoder over the{gnss, agents}observation heads.observation_log_probability()-> decoder log-density (used to reweight particles in the belief update).reward()-> learned reward head.is_terminal()-> continue/termination head, thresholded.
The trained network is injected as a DreamerWorldModel — a small framework-
agnostic protocol — so this module carries no JAX/TF dependency and is testable with a
lightweight fake. Any concrete Dreamer implementation (e.g. a DreamerV3 RSSM) that exposes
those batched operations plugs in unchanged.
The discrete action set and the observation-dict hashing/equality are inherited from
CarlaModelPOMDP so the
world and the model agree on the schema by construction.
- Classes:
DreamerWorldModel: Protocol a trained Dreamer RSSM must satisfy. DreamerCarlaModelPOMDP: Concrete CARLA model backed by a Dreamer world model.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_dreamer_model_pomdp.DreamerCarlaModelPOMDP(world_model, discount_factor, action_presets=None, max_tracked_agents=5, continue_threshold=0.5, initial_observation=None, name=None)[source]
Bases:
CarlaModelPOMDPConcrete CARLA generative model whose dynamics are a trained Dreamer world model.
The planner-side state is the Dreamer latent; transitions, observations, reward, and termination are served by the injected
DreamerWorldModel. The belief is seeded by encoding the world’s initial observation with the posterior.- Parameters:
- world_model
The trained Dreamer RSSM backing every dynamic quantity.
- continue_threshold
Termination fires when the continue head’s probability drops below this value.
Note
Reward comes from the world model’s learned reward head, not the analytic
driving_quality_reward(); a Dreamer model predicts reward directly from its latent.Example
>>> import numpy as np >>> >>> class _IdentityWorldModel: ... latent_dim = 4 ... def encode(self, observation): ... return np.zeros(self.latent_dim) ... def imagine(self, latents, controls): ... return np.asarray(latents, dtype=float) ... def decode(self, latents): ... batch = np.asarray(latents).shape[0] ... return {"gnss": np.zeros((batch, 3)), "agents": np.zeros((batch, 25))} ... def decode_log_prob(self, latents, observation): ... return np.zeros(np.asarray(latents).shape[0]) ... def reward(self, latents): ... return np.zeros(np.asarray(latents).shape[0]) ... def continue_prob(self, latents): ... return np.ones(np.asarray(latents).shape[0]) >>> >>> obs = {"gnss": np.zeros(3), "agents": np.zeros(25)} >>> env = DreamerCarlaModelPOMDP( ... _IdentityWorldModel(), discount_factor=0.95, initial_observation=obs) >>> >>> state = env.initial_state_dist().sample()[0] >>> action = env.get_actions()[0] >>> next_state, observation, reward = env.sample_next_step(state, action) >>> sorted(observation) ['agents', 'gnss'] >>> env.is_terminal(state) False
- 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_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-density of
observationsunder the per-channel perception givennext_state.
- 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.
- sample_next_state(state, action, n_samples=1)[source]
Sample
n_samplesnext states for(state, action).
- 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.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_dreamer_model_pomdp.DreamerWorldModel(*args, **kwargs)[source]
Bases:
ProtocolBatched operations a trained Dreamer RSSM must expose to back the CARLA model.
Every latent is a 1-D float vector of length
latent_dim(the packed deterministic + stochastic recurrent state). All methods are batched: they take a(batch, latent_dim)array of latents and return per-row results, so a single network call serves a whole particle set.- latent_dim
Width of a packed latent vector.
- continue_prob(latents)[source]
Probability the episode continues for each of
(batch, latent_dim)latents.
- decode_log_prob(latents, observation)[source]
Log-density of one observation under each of
(batch, latent_dim)latents.
- encode(observation)[source]
Encode a real observation into a latent via the posterior (belief seed).
POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_factored_model_pomdp module
Reference concrete CARLA generative model pairing dynamics with factored perception.
FactoredCarlaModelPOMDP implements the
CarlaModelPOMDP
interface by composing per-channel observation models — a
FactoredAgentObservationModel
on the agents channel (per-slot detection with range + occlusion gating and additive Gaussian
pose noise) and a
GnssObservationModel
on the gnss channel — the observation methods themselves are inherited from the base and
driven by that map, together with the same gym-carla driving-quality reward the world uses. The
transition dynamics are a documented identity placeholder for a specific study to replace with
real (e.g. learned) motion.
State/observation layout and the shared reward are imported from
carla_pomdp so world and model agree by
construction.
- Classes:
FactoredCarlaModelPOMDP: Concrete CARLA model with a factored-perception observation model.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_factored_model_pomdp.FactoredCarlaModelPOMDP(discount_factor, action_presets=None, max_tracked_agents=5, perception_range=50.0, occlusion_radius=1.5, pose_std=0.5, gnss_std=1e-05, detect_prob=0.95, desired_speed=8.0, out_lane_thresh=2.0, collision_penalty=100.0, observation=None, name=None)[source]
Bases:
CarlaModelPOMDPConcrete CARLA generative model pairing placeholder dynamics with factored perception.
The observation is composed per channel (held as
self.observation_models): aFactoredAgentObservationModelonagents(detection gated by perception range and geometric occlusion, additive Gaussian pose noise) and aGnssObservationModelongnss. The reward is the shared gym-carla driving-quality reward. The transition is a documented identity placeholder to be replaced with real dynamics per study.- Parameters:
discount_factor (float)
action_presets (Sequence[Tuple[float, float, float]] | None)
max_tracked_agents (int)
perception_range (float | None)
occlusion_radius (float)
pose_std (float)
gnss_std (float)
detect_prob (float)
desired_speed (float)
out_lane_thresh (float)
collision_penalty (float)
name (str | None)
- observation_models
The per-channel
{channel: CarlaObservationModel}map carrying the observation parameters (perception_range,occlusion_radius,pose_std,gnss_std,detect_prob).
- desired_speed
Target longitudinal speed (m/s) used by the reward.
- out_lane_thresh
Lateral offset (m) beyond which the reward penalises out-of-lane.
- collision_penalty
Penalty scale applied on a terminal collision in the reward.
Example
>>> import numpy as np >>> np.random.seed(42) # For reproducible results >>> >>> from POMDPPlanners.environments.carla_pomdp.carla_pomdp import ( ... AGENT_SLOT_WIDTH, EGO_STATE_WIDTH) >>> env = FactoredCarlaModelPOMDP(discount_factor=0.95) >>> >>> width = EGO_STATE_WIDTH + env.max_tracked_agents * AGENT_SLOT_WIDTH >>> state = np.zeros(width) >>> action = env.get_actions()[0] >>> >>> next_state, observation, reward = env.sample_next_step(state, action) >>> sorted(observation) ['agents', 'gnss'] >>> env.is_terminal(state) False
- 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_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.
- 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.
POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_kinematic_model_pomdp module
Concrete CARLA model with a kinematic bicycle transition under the control preset.
KinematicCarlaModelPOMDP replaces the identity-placeholder transition of
FactoredCarlaModelPOMDP
with a real ego-motion model: it propagates the ego [x, y, yaw, vx, vy, lat,
heading_err] forward one tick under the selected (throttle, steer, brake) control
using a point-mass longitudinal model plus a bicycle yaw model, and closes the range on
tracked agents by the distance the ego travelled. The factored observation model, reward,
and terminal check are inherited unchanged.
This is what gives a planner a gradient toward accelerating: because throttle now visibly increases the along-lane speed the reward rewards, POMCPOW picks controls that actually move the car (the identity placeholder made every action look motionless).
- Classes:
KinematicCarlaModelPOMDP: Factored CARLA model with a kinematic ego transition.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_kinematic_model_pomdp.KinematicCarlaModelPOMDP(discount_factor, dt=0.05, action_presets=None, max_tracked_agents=5, wheelbase=2.8, max_steer_angle=0.6, accel=3.0, brake_decel=8.0, drag=0.05, collision_gap=5.0, collision_halfwidth=1.2, safe_distance=12.0, stop_gap=0.0, name=None, **kwargs)[source]
Bases:
FactoredCarlaModelPOMDPFactored CARLA model whose transition is a kinematic bicycle propagation.
- Parameters:
discount_factor (float)
dt (float)
action_presets (Sequence[Tuple[float, float, float]] | None)
max_tracked_agents (int)
wheelbase (float)
max_steer_angle (float)
accel (float)
brake_decel (float)
drag (float)
collision_gap (float)
collision_halfwidth (float)
safe_distance (float)
stop_gap (float)
name (str | None)
kwargs (Any)
- dt
Integration step (seconds); must match the world’s
fixed_delta_seconds.
- wheelbase
Bicycle-model wheelbase (m) mapping steer to yaw rate.
- max_steer_angle
Steering command of 1.0 maps to this front-wheel angle (rad).
- accel
Longitudinal acceleration per unit throttle (m/s^2).
- brake_decel
Longitudinal deceleration per unit brake (m/s^2).
- drag
Linear speed-proportional deceleration coefficient (1/s).
- collision_gap
Forward ego-frame distance (m) within which a present agent ahead is treated as a predicted collision by
is_terminal().
- collision_halfwidth
Lateral ego-frame half-corridor (m) within which a present agent ahead is treated as a predicted collision by
is_terminal().
- safe_distance
Lead gap (m) at/above which the reward targets the full
desired_speed; the obstacle-aware target ramps down below it.
- stop_gap
Lead gap (m) at/below which the obstacle-aware target speed is zero;
0.0keeps the flatdesired_speed.
Example
>>> import numpy as np >>> from POMDPPlanners.environments.carla_pomdp.carla_pomdp import ( ... AGENT_SLOT_WIDTH, EGO_STATE_WIDTH) >>> env = KinematicCarlaModelPOMDP(discount_factor=0.95, dt=0.05) >>> >>> width = EGO_STATE_WIDTH + env.max_tracked_agents * AGENT_SLOT_WIDTH >>> state = np.zeros(width) >>> throttle_action = 0 # (0.5, 0.0, 0.0) cruise straight >>> next_state = env.sample_next_state(state, throttle_action) >>> >>> bool(next_state[3] > 0.0) # throttle produced forward velocity True
- is_terminal(state)[source]
Whether a present agent occupies the ego’s footprint just ahead.
Because this model does predict ego and agent motion, it can foresee running into the vehicle ahead. Any present agent slot within
collision_gapmetres forward andcollision_halfwidthmetres laterally (ego frame) is treated as a collision, which the inherited reward turns into the terminalcollision_penalty(driving_quality_reward()).
- reward(state, action, next_state=None)[source]
Driving-quality reward whose target speed adapts to the nearest lead obstacle.
The bare parent reward tracks a fixed
desired_speedand only charges the collision once an agent is inside the terminal box — a cliff that cannot be braked for at speed, while a large fixed penalty instead freezes the ego in traffic. This override follows Roach’s obstacle-aware desired speed: the target equals the fulldesired_speedwhen the lead gap is at leastsafe_distance, ramps linearly to zero atstop_gap, and is zero closer in. The ego is thus rewarded for driving when the road is clear and for slowing as an obstacle nears — including the lidar/camera obstacle fused into the agent slots — without a separate penalty term that traps it at a standstill.stop_gap == 0keeps the flat parent behaviour.
- sample_next_state(state, action, n_samples=1)[source]
Sample
n_samplesnext states for(state, action).
- 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.
POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_kinematic_vectorized_model module
Torch, on-device vectorized generative model for the kinematic CARLA model.
This module provides CarlaKinematicVectorizedModel, a fully batched,
GPU-friendly implementation of
VectorizedGenerativeModel
for KinematicCarlaModelPOMDP.
It re-expresses the scalar model’s kinematic-bicycle transition, obstacle-aware
driving-quality reward, predicted-collision terminal check, and factored
perception (GNSS Gaussian noise plus per-slot agent detection with range and
occlusion gating and additive pose noise) 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 (control presets, kinematic
coefficients, reward weights, perception parameters) is read from a live
KinematicCarlaModelPOMDP 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 scalar model.
State layout is [ego(7)] + K*[present, rel_x, rel_y, rel_yaw, rel_speed] with
K = max_tracked_agents (default ds = 7 + 5*5 = 32); the observation drops
the ego block down to the GNSS position and keeps the agent slots
(do = 2 + K*5 = 27). Actions are integer indices into the discrete
(throttle, steer, brake) control presets.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_kinematic_vectorized_model.CarlaKinematicVectorizedModel(env, *, device=None, dtype=torch.float32, observation_resolution=0.5)[source]
Bases:
objectFully vectorized torch generative model for the kinematic CARLA model.
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. Actions are integer indices into the fixed
(throttle, steer, brake)control-preset table read from the scalar model.- Parameters:
env (KinematicCarlaModelPOMDP)
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 control presets.
- state_dim
Width of the state vectors (
7 + K*5).
- observation_dim
Width of the observation vectors (
2 + K*5).
Example
>>> import torch >>> from POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_kinematic_model_pomdp import ( ... KinematicCarlaModelPOMDP, ... ) >>> from POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_kinematic_vectorized_model import ( ... CarlaKinematicVectorizedModel, ... ) >>> torch.manual_seed(0) <torch._C.Generator object at ...> >>> env = KinematicCarlaModelPOMDP(discount_factor=0.95, dt=0.05) >>> model = CarlaKinematicVectorizedModel(env, device=torch.device("cpu")) >>> states = torch.zeros(3, model.state_dim) >>> actions = torch.zeros(3, dtype=torch.int64) # cruise straight >>> next_states = model.sample_next_states(states, actions) >>> rewards = model.rewards(states, actions, next_states) >>> tuple(next_states.shape), tuple(rewards.shape) ((3, 32), (3,))
- 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.carla_pomdp.carla_generative_models.carla_model_pomdp module
Abstract planner-side generative-model interface for the CARLA world.
CarlaPOMDP is a
forward-only world (no densities, no state injection). A planner instead carries a
generative model as policy.environment — one that can sample transitions from an
arbitrary state, score an observation against a state, and supply a reward. This module
defines the interface that model must satisfy: CarlaModelPOMDP owns only the
CARLA state/observation schema (agent-slot layout, discrete action set, observation-dict
hashing/equality) and leaves every dynamic quantity — transition, observation, reward,
terminal — abstract for a study- or task-specific subclass (e.g. a learned model) to fill
in.
A runnable reference implementation with a fixed factored observation model lives in
carla_factored_model_pomdp.
The schema (state/observation layout, action presets) is imported from
carla_pomdp so world and model agree by
construction.
- Classes:
CarlaModelPOMDP: Abstract generative-model interface over the CARLA schema.
- class POMDPPlanners.environments.carla_pomdp.carla_generative_models.carla_model_pomdp.CarlaModelPOMDP(discount_factor, action_presets=None, max_tracked_agents=5, observation_models=None, name=None)[source]
Bases:
DiscreteActionsEnvironmentAbstract generative-model interface paired with the forward-only CARLA world.
Concrete subclasses supply the dynamics —
sample_next_state(),transition_log_probability(),reward(),is_terminal(), and the initial-distribution hooks — while this base owns the fixed CARLA schema shared by every model (the ego + agent-slot state layout, the discrete action set, and the observation-dict hashing/equality) and the observation model. The observation is factored by channel: this base holds a{channel: CarlaObservationModel}map (observation_models) and composes it —sample_observation()perceives each channel of the clean observation built from the state,observation_log_probability()sums the per-channel densities, andencode_observation()(the single raw-observation seam) perceives each channel of the forward-only world’s raw reading into the same perceived space, so the belief filter and planner search operate on one consistent (encoded) observation. Two models differ in their observation only by the channel models they hold; a model that scores observations for a belief update needs every channel to provide a density (supports_density). A subclass whose observation is not a per-channel clean transform (e.g. a learned latent decoder) may instead overridesample_observation(),observation_log_probability()andencode_observation()directly.- Parameters:
- action_presets
Discrete
(throttle, steer, brake)control triples; the discrete action set is the indices into this list.
- max_tracked_agents
Number of fixed agent slots in the state/observation.
- observation_models
The per-channel
{channel: CarlaObservationModel}map composed to produce the observation, orNonefor a subclass that overrides the observation methods directly.
Note
This is an abstract base class and cannot be instantiated directly. See
FactoredCarlaModelPOMDPfor a concrete reference implementation.- encode_observation(observation)[source]
Perceive the world’s raw observation into the belief/planner observation space.
The forward-only world emits a raw, fully-detected observation; each channel model degrades its channel (per-slot detection gating, occlusion, additive noise) into the perceived observation the belief filter and planner search operate in. This is the single raw-observation seam — every other observation method works in the perceived (encoded) space. A subclass without channel models (
observation_models is None, e.g. a learned latent model) inherits the identity default.
- 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.
- 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.
- observation_log_probability(next_state, action, observations)[source]
Log-density of
observationsunder the per-channel perception givennext_state.
- abstractmethod sample_next_state(state, action, n_samples=1)[source]
Sample
n_samplesnext states for(state, action).