Goal-Conditioned Transition Modeling with Energy-Based Models

This note describes a goal-conditioned transition model with two prediction modes:

1. Action-conditioned prediction: an action is known or sampled, so the model predicts a specific next latent state.
2. Action-free prediction: no action is given, so the model represents a distribution over possible next latent states with an energy function.

The main idea is that the action-free model should not directly output a probability vector over arbitrary future states. Instead, it should assign an energy to any candidate next latent state. Lower energy means the candidate is more plausible.

1. Setup and Notation

Given a current observation $x$ and goal $g$, define a goal-conditioned latent state:

where:

• $x$ is the current observation,
• $g$ is the goal,
• $F_\psi$ is a goal-conditioned encoder,
• $s \in \mathbb{R}^d$ is the goal-conditioned latent state.

For the next observation $x’$, define the target next latent state:

$F_{\bar{\psi}}$ can be:

• the same encoder as $F_\psi$,
• a detached copy of the encoder,
• an EMA target encoder.

In practice, a detached or EMA target encoder is often safer because it can reduce representation collapse.

2. Model Components

2.1 Action-Conditioned Predictor

If an action is available or sampled, define:

where $\epsilon$ is noise if $H_\eta$ is stochastic, for example when $H_\eta$ is a diffusion model or another generative action sampler.

The action-conditioned predictor outputs one specific next latent state:

This corresponds to modeling:

2.2 Action-Free Energy Model

If no action is given, the model should represent the distribution over possible next latent states:

An energy-based model represents this distribution indirectly:

where $\tilde{s}’$ is a candidate next latent state.

The conditional distribution is:

with partition function:

The EBM is queried with both the current latent $s$ and a candidate next latent $\tilde{s}’$. It outputs a scalar energy:

Low energy means $\tilde{s}’$ is a plausible next latent state from $s$. High energy means it is unlikely.

A simple implementation is:

where $[s,\tilde{s}’]$ denotes concatenation.

3. Relationship Between the Two Modes

When the action is unknown, the true action-free transition distribution can be viewed as the action-marginal distribution:

This means that many actions may be possible from $s$, and each action may lead to a different next state. If the action is unobserved, the next-state distribution is a mixture over possible actions.

Explicitly computing the marginal would require modeling both:

and:

The EBM avoids this explicit decomposition. It learns the final compatibility function directly from observed transition data:

At optimum, the energy can be interpreted approximately as:

where $C(s)$ is an arbitrary function of $s$.

For a fixed $s$, $C(s)$ does not affect optimization over $s’$, so:

corresponds to finding a likely next latent state.

4. Exact EBM Likelihood

For one positive transition pair:

the negative log-likelihood is:

Using the EBM definition:

Therefore:

This objective has two forces:

• decrease the energy of real next states,
• increase energy elsewhere through the partition-function term.

4.1 Gradient of the Exact Likelihood

Differentiate the exact loss:

Because:

we have:

Thus:

This is the classical EBM learning rule:

• lower energy at real data samples,
• raise energy at model-generated samples.

5. Contrastive Training with In-Batch Negatives

Given a batch of $B$ transitions:

define pairwise energies:

For each $s_i$, the positive target is $s_i’^+$. The negatives are the other batch targets:

Define logits:

where $\tau>0$ is a temperature.

The contrastive probability is:

The loss for sample $i$ is:

Equivalently:

The batch loss is:

This is a practical starting point for training the conditional EBM.

5.1 Gradient of the Contrastive Loss

Recall:

and:

The standard cross-entropy derivative is:

Since:

we obtain:

For the positive pair $j=i$:

Under gradient descent, this lowers the positive energy $e_{ii}$.

For a negative pair $j\neq i$:

Under gradient descent, this raises the negative energy $e_{ij}$.

The parameter gradient is:

This is the finite-negative version of the exact EBM gradient.

6. Inference with the Action-Free EBM

Once $E_\theta(s,\tilde{s}’)$ is trained, there are two main inference modes.

6.1 MAP Prediction

If you want one likely next latent state, solve:

This can be optimized using gradient descent:

This produces a likely next latent state, usually one mode of the learned distribution.

6.2 Sampling Diverse Futures

If you want diverse samples from:

you can use Langevin dynamics:

where:

Without the noise term, this becomes pure energy minimization.

6.3 Useful Derivative for Optimization and Sampling

For fixed $s$:

Because $Z_\theta(s)$ does not depend on $\tilde{s}’$:

and:

7. Action-Conditioned Losses

When an action is available or sampled:

and:

If a ground-truth next latent $s’^+$ is available, use a prediction loss:

Equivalently:

The derivative with respect to $\hat{s}’$ is:

The parameter gradient for $\phi$ is:

where $\nabla_\phi G_\phi$ is the Jacobian of the predictor output with respect to its parameters.

8. Compatibility Between the Predictor and the EBM

Because $G_\phi(s,a)$ produces a candidate next latent state, you can encourage it to lie in a low-energy region of the action-free EBM:

Let:

Then:

The derivative with respect to $\theta$ is:

The derivative with respect to $\hat{s}’$ is:

The derivative with respect to predictor parameters $\phi$ is:

If the action is generated by:

then the derivative with respect to $\eta$ is:

This means the action sampler can be trained to generate actions whose predicted next states have low energy.

9. Full Training Objective

A reasonable full objective is:

where:

• $\mathcal{L}_{\text{EBM}}$ trains the action-free energy landscape,
• $\mathcal{L}_{\text{pred}}$ trains the action-conditioned predictor,
• $\mathcal{L}_{\text{compat}}$ aligns action-conditioned predictions with the energy landscape,
• $\mathcal{L}_{\text{reg}}$ can include latent normalization, energy regularization, entropy regularization for the action sampler, smoothness penalties, or other stabilizers.

A practical contrastive EBM loss is:

10. Negative Sampling Options

For the contrastive EBM loss, useful negative candidates $\tilde{s}’^-$ include:

11. Minimal Implementation Sketch

# x, g, x_next: batch data
# F_enc: encoder
# F_target: target encoder
# E_model: energy model
# G_act: action-conditioned predictor
# H: action sampler

import torch
import torch.nn.functional as F

s = F_enc(x, g)                         # [B, d]
s_next = F_target(x_next, g).detach()   # [B, d]

B, d = s.shape

# ------------------------------------------------------------
# EBM contrastive loss
# energies[i, j] = E_model(s_i, s_next_j)
# ------------------------------------------------------------

s_i = s[:, None, :].expand(B, B, d)
sp_j = s_next[None, :, :].expand(B, B, d)

energies = E_model(s_i, sp_j)           # [B, B]
logits = -energies / tau

labels = torch.arange(B, device=s.device)
loss_ebm = F.cross_entropy(logits, labels)

# ------------------------------------------------------------
# Action-conditioned mode
# ------------------------------------------------------------

eps = torch.randn(B, noise_dim, device=s.device)
a = H(s, eps)
s_pred = G_act(s, a)

loss_pred = 0.5 * ((s_pred - s_next) ** 2).sum(dim=-1).mean()

# ------------------------------------------------------------
# Compatibility with EBM
# ------------------------------------------------------------

loss_compat = E_model(s, s_pred).mean()

# ------------------------------------------------------------
# Full loss
# ------------------------------------------------------------

loss = loss_ebm + lambda_pred * loss_pred + lambda_compat * loss_compat
loss.backward()

12. Recommended Starting Point

Start with the simplest stable version:

1. Encode the current latent:

   \[s = F(x,g).\]

2. Encode the target next latent:

   \[s'^+ = F_{\bar{\psi}}(x',g).\]

3. Train the action-free EBM:

   \[E_\theta(s,\tilde{s}').\]

4. Use in-batch negatives:

   \[E_\theta(s_i,s_j'^+), \quad j\neq i.\]

5. Train with a contrastive softmax loss.

6. For no-action MAP prediction, optimize:

   \[\hat{s}' = \arg\min_{\tilde{s}'} E_\theta(s,\tilde{s}').\]

7. For diverse no-action futures, sample with Langevin dynamics:

   \[\tilde{s}'_{k+1} = \tilde{s}'_k - \frac{\alpha}{2} \nabla_{\tilde{s}'}E_\theta(s,\tilde{s}'_k) + \sqrt{\alpha}\epsilon_k.\]

13. Key Takeaways

Do not write the action-free model as:

if this suggests a direct finite probability output over arbitrary next latent states.

Instead, write it as:

The distribution is implicit:

The two prediction modes are:

The action-conditioned model learns:

while the EBM learns the action-marginal transition:

The EBM does not need to compute this integral explicitly. It learns the marginal compatibility directly from transition data by assigning low energy to observed next latents and high energy to unlikely alternatives.