| name | vbp-message-passing |
| description | Converts a standard Bethe BP implementation to the VBP (cross-entropy planning) message-passing scheme. Single action channel in the dynamics numerator, unmodified observation kernel, theta inferred via cavities. Reference paper Appendix E. |
VBP Message Passing (Cross-Entropy Planning)
VBP is standard Bethe BP with a single entropy correction -- the policy entropy -- reparameterized into the dynamics kernel via one action channel per time step. Substantially simpler than AIF-MP (1 channel vs 4).
Reference: paper Section 4.2, Appendix E (full derivation).
1. Objective (paper eq E.4)
$$\Delta F_{\mathrm{VBP}} = +\sum_{t=1}^{T} \mathbb{H}[q(u_t|x_{t-1})]$$
This penalizes action entropy, pushing the agent toward a committed policy (cross-entropy planning = expected reward maximization with policy commitment).
2. Single Channel
| Channel | Symbol | Normalizes over | Enters kernel |
|---|
| Policy | $r_{u\mid x,t}(u_t\mid x_{t-1})$ | $\sum_{u_t} r = 1;\forall x_{t-1}$ | Dyn numerator |
Setting $r_{u|x} \equiv$ uniform recovers standard BP. Since +H[q(u|x)] has a positive sign, the channel enters the numerator (multiplicative), unlike the AIF dynamics channel $r_{x|xu}$ which enters the denominator (paper Remark 17).
3. Modified Kernels (paper eq E.8)
| Factor | VBP kernel |
|---|
| $f_{\mathrm{obs}_t}$ | $p(y_t\mid x_t,\theta)$ -- unmodified |
| $f_{\mathrm{dyn}_t}$ | $p(x_t\mid x_{t-1},\theta,u_t) \cdot r_{u\mid x,t}(u_t\mid x_{t-1})$ |
4. Channel Update (paper Proposition 19, eq E.11)
$$r_{u|x,t}^*(u_t|x_{t-1}) = \frac{q_{\mathrm{pair},t}(x_{t-1}, u_t)}{q_{x_{t-1}}(x_{t-1})} = q_t(u_t|x_{t-1})$$
where $q_{\mathrm{pair}} = \sum_{x_t,\theta} q_{\mathrm{dyn},t}$ is the pair marginal from the dynamics factor belief.
5. Implementation Details (inference/vbp_channel.py)
- theta handling: Inferred via cavity messages. Uses
compute_theta_cavities_extended for per-timestep dyn/obs/pref cavity beliefs. This is different from AIF-MP which uses theta prior directly.
- Observation kernel: Raw B tensor, no channels. Tiled over time as
log_B_tiled.
- Damping: Geometric damping in log-space (same as AIF-MP):
damped = (1-alpha)*log_old + alpha*log_new. Applied to action channels only.
- Sparse transitions: Supports
T_idx (S, A, theta) int32 sparse index via sparse_reduced_weighted, sparse_dyn_to_theta_weighted, sparse_pair_marginal from messages.py.
- Preference factors: Supports 2D goal
C(x, theta) with pref->x and pref->theta messages via 3-way theta cavities.
- Return:
(action_dist, log_action_channels) where log_action_channels is (T, n_states, n_actions).
Iteration structure
Initialize: r_{u|x} = uniform, theta messages = zeros
For each iteration:
1. Theta cavities (exclude dyn[t] from theta belief for dyn, exclude obs[t] for obs)
2. Dyn kernels: T * r_{u|x} (channel in numerator)
3. Reduce: marginalize theta from dyn kernels using cavity beliefs
4. obs->x messages (raw B with obs cavity), pref->x if applicable
5. Forward pass
6. Backward pass + action marginals
7. dyn->theta messages (with kernel weight = action channel)
8. obs->theta messages (raw B)
9. pref->theta messages (if applicable)
10. Dyn region beliefs -> pair marginal -> action channel -> damp
Key functions used
compute_dyn_kernels_vbp(log_T_kernel, log_r_ux) -- T + r_ux (broadcast)compute_pair_marginal(log_dyn_regions) -- logsumexp over (x_new, theta)compute_action_channel(log_pair) -- normalize over actionsdamp_log_channel(old, new, damping, cond_axis) -- geometric damping
6. Properties (paper Section E.6)
- No min-max structure: Single channel in numerator only. Joint optimization over $(q, r_{u|x})$ is pure minimization, avoiding AIF's saddle-point convergence difficulties.
- Reduction to standard BP: Setting $r_{u|x} \equiv$ uniform removes the correction.
- Fixed-point interpretation: At convergence, dyn kernel becomes $p(x_t|x_{t-1},\theta,u_t) \cdot q(u_t|x_{t-1})$, reweighting the transition by the learned policy.
7. Verification Checklist
