# ROS 2 planner wrapper

`ros2/wm_planner_ros` exposes two ROS 2 executables:

| Executable | Model | Observations |
|---|---|---|
| `wm_planner_bridge_node` | Standard `WorldModel`: CNN, DINOv2, or V-JEPA2 encoder with RSSM or MLP dynamics | Raw RGB, raw depth, odometry, IMU |
| `dino_wm_planner_bridge_node` | Separate DinoWM architecture with a token predictor | Raw RGB, odometry, IMU |

Both executables use the same Python bridge infrastructure in `bridge_adapter.py`: ROS subscriptions, goal handling, publishers, a bounded observation history, and a latest-observation-wins planning worker. `model_adapters.py` supplies the model-specific observation and planning methods. The standard adapter covers CNN+RSSM, DINOv2+RSSM, and JEPA-lite; the DinoWM adapter uses a three-frame history by default and keeps DinoWM's own CEM and TF path. JEPA-lite uses the standard `WorldModel` with a different training objective and normally MLP dynamics. See [World model](world_model.md) for the model architectures.

## Setup and launch

Follow [Getting started](getting_started.md) to build and source the wrapper. 

Install `image_transport`, `compressed_image_transport`, `rosbag2_transport`, and `rosbag2_storage_mcap` in the selected ROS distribution.
```bash
sudo apt-get install ros-$ROS_DISTRO-image-transport \
                     ros-$ROS_DISTRO-compressed-image-transport \
                     ros-$ROS_DISTRO-rosbag2_transport \
                     ros-$ROS_DISTRO-rosbag2-storage-mcap
```

Bag-in-the-loop launch:
```bash
ros2 launch wm_planner_ros bag_in_the_loop.launch.py \
  bag:=/path/to/recording.mcap robot:=<robot> \
  model:=cnn checkpoint_path:=/path/to/cnn_rssm/latest.pt device:=cuda
```

From the sourced workspace, replay an MCAP bag with a matching robot config and checkpoint:

```bash
ros2 launch wm_planner_ros bag_in_the_loop.launch.py \
  bag:=/path/to/recording.mcap robot:=<robot> \
  model:=cnn checkpoint_path:=/path/to/cnn_rssm/latest.pt device:=cuda
```

For DINOv2+RSSM, use `model:=dino` and its corresponding checkpoint. For JEPA-lite, first train the state decoder as shown in [Training](training.md), then launch with its updated checkpoint:

```bash
ros2 launch wm_planner_ros bag_in_the_loop.launch.py \
  bag:=/path/to/recording.mcap robot:=<robot> \
  model:=jepa-lite checkpoint_path:=/path/to/latest_with_rollout_heads.pt device:=cuda
```

DinoWM uses the same bag launch with its own checkpoint format and encoder name:

```bash
ros2 launch wm_planner_ros bag_in_the_loop.launch.py \
  bag:=/path/to/recording.mcap robot:=<robot> \
  model:=dinowm checkpoint_path:=/path/to/model_latest.pth \
  dino_model_name:=dinov2_vits14 num_hist:=3 device:=cuda
```

Pass the MCAP path and choose its robot config with `robot:=NAME`, or use `robot:=/path/to/robot.yaml`. `model:=dino` selects DINOv2+RSSM in the standard adapter; `model:=dinowm` selects the separate DinoWM model adapter. The JEPA-lite checkpoint metadata chooses its CNN or DINOv2 encoder and trained dynamics. `device:=cpu` is the portable default. Stop the launch with Ctrl+C.

| Launch argument | Default | Effect |
|---|---|---|
| `bag` | required | MCAP file to replay; relative paths resolve from the launch command's working directory. |
| `robot` | required | Name of an installed `configs/robots/*.yaml` file, or an explicit YAML path. The config supplies the node namespace and recorded sensor topics. |
| `model` | `cnn` | `cnn` selects CNN+RSSM; `dino` selects DINOv2+RSSM; `jepa-lite` selects JEPA-lite; `dinowm` selects DinoWM. |
| `checkpoint_path` | empty | Local `.pt` or `.pth` checkpoint. Required for JEPA-lite and DinoWM. For `cnn` or `dino`, an empty value requests `cnn_rssm_c1/latest.pt` or `dino_rssm_c1/latest.pt` from the default Hugging Face repository. |
| `norm_stats_path` | empty | Standard-model `norm_stats.pt`; if empty, the bridge looks beside the checkpoint. DinoWM uses its own normalization parameters. |
| `rate` | `0.1` | Positive multiplier for bag playback speed. |
| `dino_model_name` | `dinov2_vits14` | DinoWM encoder name; ignored by the standard adapter. |
| `num_hist` | `3` | DinoWM observation history length; ignored by the standard adapter. |
| `device` | `cpu` | PyTorch device, such as `cpu` or `cuda`. |

The launch validates the bag and any supplied checkpoint or normalization file paths. The three standard-WorldModel choices read `args` and `model_config` from training checkpoints to construct the trained backbone, dynamics, image size, and sensor/action dimensions before loading the saved `model` weights. For `cnn` and `dino`, the checkpoint architecture must match the launch choice. The bridge still accepts older plain state dict checkpoints using its configured defaults. If a standard-model checkpoint download fails, the node continues with random weights; provide a trained checkpoint for meaningful plans. Missing normalization statistics also degrade trajectory quality.

JEPA-lite training leaves the state and risk heads unsupervised. Use a rollout-trained checkpoint such as `latest_with_rollout_heads.pt` produced by `train_state_head_decoder.py` for goal-based planning. The launch sets `risk_weight:=0.0` for JEPA-lite, so its untrained risk head does not affect CEM scoring. The default `norm_stats.pt` lookup is beside the checkpoint; supply `norm_stats_path` if it is elsewhere.


### JEPA-lite latent planning

The original `wm_planner_bridge_node` now supports `planner_mode: latent`.
Use `ros2/wm_planner_ros/configs/jepa_latent.yaml` with a JEPA-lite training
checkpoint. This mode compares the rollout latent with an encoded goal
observation; it does not require a fitted state head. `trajectory_source:
actions` converts the selected `position_cmd` position dimensions into the
published path, beginning at the current odometry pose. The path represents
commanded setpoints, not decoded predicted robot positions.

The planner needs a goal observation before planning. Call the private
`~/capture_goal` Trigger service while the robot is at the destination, or
publish a `sensor_msgs/Image` on `~/goal_image` and optionally a
`geometry_msgs/PoseStamped` on `~/goal`. The image is paired with the latest
synchronized odometry and IMU; the pose overrides its destination position
and orientation. `~/clear_goal` removes the stored observation. In latent
mode, `~/planned_actions` publishes the selected action sequence as a flat
`Float32MultiArray` and `~/debug_metrics` publishes terminal latent distance,
mean path latent distance, current latent distance, CEM cost, goal reached
flag, and trajectory endpoint distance, in that order.

`planner_mode: auto` chooses latent mode for a JEPA-lite checkpoint without a
fitted state head, and state mode otherwise. `trajectory_source: auto` uses a
fitted state head when available and action position setpoints otherwise.
State-head availability is checked from the checkpoint's `head_training`
metadata, `latest_with_state_head.pt` filename, or a loaded
`state_head_decoder.pt`. An explicit `planner_mode: state` requires a fitted
head. Latent mode still requires a stored goal observation; Nexus currently
supplies only a goal pose in its MRP request, so populate the observation
through the goal service or image topic before requesting an MRP plan.
Keep that image paired with the same destination as the Nexus pose; the MRP
request replaces the stored goal position and orientation, not its image.

For JEPA-lite HF-dataset runs, set `depth_mode: zeros`,
`imu_layout: dataset`, and `norm_stats_mode: off` as in the supplied YAML.
Set `action_bounds_path` or place `action_bounds.pt` next to the checkpoint
for per-dimension CEM sampling; without it, the configured scalar action
bounds are used. Set `plan_horizon: 0` to use the training horizon. State
planning remains the default unless `planner_mode` is set explicitly.
The bridge rejects a planning horizon longer than the dynamics or state-head
training horizon. `allow_horizon_extrapolation: true` bypasses this check, but
is intended only for explicit experiments. A checkpoint whose action contract
uses window-start-relative future poses requires `goal_frame_mode: current_map`.

## What the launch starts

1. The selected executable in the namespace set by the robot config's `name`. `cnn`, `dino`, and `jepa-lite` start `wm_planner_bridge_node`; `dinowm` starts `dino_wm_planner_bridge_node`. Standard RSSM choices use their selected backbone, while JEPA-lite uses checkpoint metadata. Both share the buffered planner worker; the standard adapter needs one observation and DinoWM uses `num_hist` observations.
2. For configs with `sensors.color_encoding: compressed`, `image_transport republish` turns the recorded RGB stream into raw `sensor_msgs/Image` on the topic without the `/compressed` suffix. This requires the `compressed_image_transport` plugin. Raw color configs use the recorded topic directly.
3. `ros2 bag play` after a five-second startup delay, with MCAP storage, `/clock`, the selected rate, and `--loop`. The standard adapter replays RGB, depth, odometry, and IMU; DinoWM replays RGB, odometry, IMU, `/tf`, and `/tf_static`. Looping lets the bridge receive later passes if model loading takes longer than the first pass.

The launch reads `topics.color`, `topics.odom`, and `topics.imu` from the selected robot YAML file. For example, `robot:=race16` uses [race16.yaml](../configs/robots/race16.yaml).

| Stream | Config key | Consumer |
|---|---|---|
| RGB | `topics.color`, `sensors.color_encoding` | Both adapters receive raw `sensor_msgs/Image`. |
| Depth | `topics.depth` | Standard adapter only; raw `sensor_msgs/Image`. |
| Odometry | `topics.odom` | Both adapters; `nav_msgs/Odometry`. |
| IMU | `topics.imu` | Both adapters; `sensor_msgs/Imu`. |
| TF | Global `/tf`, `/tf_static` | DinoWM's TF listener only. |

Both adapters use an approximate-time synchronizer and the ROS sensor-data QoS profile. The shared history buffer holds one synchronized observation for the standard adapter and `num_hist` observations for DinoWM. A dedicated worker plans from the newest complete snapshot; incoming observations replace any pending snapshot while a plan is running. The standard adapter's `plan_on_observation` parameter can disable automatic planning for the MRP action path.

## Goals, outputs, and node parameters

The launch does not replay a goal topic. With the launch defaults, either bridge plans toward `goal_xyz:=[0.0, 0.0, 0.0]` until it receives a `geometry_msgs/PoseStamped` on its private `~/goal` topic. The standard node's name is `wm_planner_bridge`; DinoWM's is `dino_wm_planner_bridge`. The selected node publishes:

| Topic | Type | Purpose |
|---|---|---|
| `/<robot_name>/<node_name>/planned_trajectory` | `nav_msgs/Path` | Chosen trajectory. |
| `/<robot_name>/<node_name>/planned_trajectory/markers` | `visualization_msgs/MarkerArray` | Visualization markers for the chosen trajectory. |
| `/<robot_name>/<node_name>/debug_metrics` | `std_msgs/Float32MultiArray` | Standard model: goal distance, mean risk, goal weight, risk weight. DinoWM: goal distance and goal weight. |

The launch leaves `publish_frame` at each node's default (`body`). The standard node also supports `goal_topic`, `goal_xyz`, `use_goal_topic`, `publish_frame`, `queue_size` (8), image size (224×224), model dimensions (`odom_dim`, `imu_dim`, `act_dim`, `state_dim`, `z_vis_dim`, `z_prop_dim`, `z_dim`), and CEM parameters (`plan_horizon` 15, `plan_iters` 5, `plan_population` 128, `plan_elites` 16, action bounds, `cem_alpha`, `cem_temporal_noise`, `goal_weight`, `path_weight`, `reverse_progress_weight`, `trajectory_smoothness_weight`, `action_smoothness_weight`, and `risk_weight`). The path and smoothness weights default to zero for compatibility. Pose-action smoothness uses position second differences and quaternion geodesic-rate changes rather than direct quaternion subtraction. Training checkpoint metadata supplies the checkpoint-sensitive dimensions, image size, action semantics, and maximum trained horizon; plain state dict checkpoints use the node defaults. These are **node parameters**, not additional launch arguments; change them in the launch file or start the executable directly with a ROS parameter file.

The standard bridge can also expose the optional `~/plan_mrp` action via `enable_mrp_action:=true` when `phx_nav_msgs` is available. The bag launch disables that action. DinoWM loads a different `.pth` checkpoint containing separate encoder and predictor modules, uses `num_hist` frames (which must match training), and transforms its output through TF. Its normalization arrays and CEM settings are node parameters; the shared adapter leaves that model-specific math inside the DinoWM node.

For offline extraction and training from MCAP recordings, see [Data pipeline](data_pipeline.md).