Learn Dynamics with Diffusion Forcing
Boyuan Chen
Path to General EI: world models
Auto Regressive Models
Auto Regressive Models
- Flexible horizon
- Causal correctness
- better capture likelihood of "single-step transition"
- Better compositionality
- More efficient for search
Whole Sequence Diffusion
Whole Sequence Diffusion
- Flexible conditioning at test time wo/ retraining
- Guidance over long horizon (for planning)
- Better capture global, joint distribution
Auto Regressive Models
- Flexibility
- Capture causal dynamics
Whole Sequence Diffusion
- Long horizon guidance
- Capture global likelihood
Get good from both!
Diffusion Forcing (final name TBD) is a probabilistic sequence model that interleaves time axis of auto-regressive models and noise-level axis of diffusion models.
Prelim: Bayes Filter
Prior Model to predict state when there is no updated observation
Prelim: Bayes Filter
Posterior Model to predict state when new observation is made
What if we relax bayes filter from the binary duo of having an observation vs no observation?
No observation -> Noisy observation -> Has observation
Core idea: masking by noise
Masking is a fundamental technique in self-supervised learning
Core idea: masking by noise
What if we mask by noise, instead of zero padding?
Masked
Not masked
partially masked
Diffusion Forcing
Diffusion
Step (k)
Time step (t)
Diffusion Forcing
Noise
level (k)
Time step (t)
Diffusion Forcing
Diffusion Forcing
Diffusion Forcing
Diffusion Forcing
What does this give us:
We can diffuse a frame at far future, without fully diffusing the past!
(because we are trained on it)
Diffusion Forcing
And after two pages of math
Comparison
RNN is all you need?
50%
50%
RNN is all you need?
Optimal one step prediction in state space
with deterministic model
RNN is all you need?
Optimal one step prediction in pixel space
with deterministic model
RNN is all you need?
Imitate human trajectories
RNN is all you need?
Imitate human trajectories with deterministic model
Solution: Learn Distribution
Variational Inference
Variational Inference
A distribution we know
how to sample from
A distribution we know
how to sample from
target distribution
Types of variational model
Diffusion Model
Forward Diffusion process
\text{Given a data point sampled from a real data distribution}, \mathbf{x}_0 \sim q(\mathbf{x})
\text{define a forward diffusion process that gradually add gaussian noise}
Forward Diffusion process
Forward Diffusion process
Reverse Diffusion process
\text{Given a data point sampled from gaussian noise}, \mathbf{x}_T \sim N(0,I)
\text{Iteratively applies model to generate } x_{T-1} ... x_1, x_0
Reverse Diffusion process
Diffusion Training
Diffusion Training
Diffusion Training
Diffusion Training
Conditional Diffusion
x_T \bigoplus y
x_t \bigoplus y
x_{t-1} \bigoplus y
x_{0} \bigoplus y
Let y be a variable we want to condition on, e.g. y=a label='cat'
Conditional Diffusion
Let y be a variable that's learned? Yes!
It could be a differentiable latent code that encodes belief given all previous observations.
Warning!
Change of notation
Steps in diffusion are usually denoted by t. But from now on we will denote it as k, as we will introduce the actual time variable, t.
How to learn stochastic dynamics?
How to learn stochastic dynamics?
Learn dynamics and control with
Diffusion Forcing
Diffusion Forcing is a sequence model that
- have out-standing performance on continuous sequence modeling
- is a state-space model that's strictly markovian
- has explicit control over context length and thus compositionality
- has explicit control over generation horizon
- can take advantage of noisy observations
has explicit control over generation horizon
Full Trajectory rollout
Single step rollout
Our method can do both
State Space Models
\text{Let's omit action and } s_0, o_0 \text{ for now for simplicity. Learn} \\
\text{Prior model: } s_{t+1} \sim p(s_{t+1}|s_t) \\
\text{Posterior model: } s_{t+1} \sim p(s_{t+1}|s_t,o_{t+1}) \\
\text{Observation model: } o_t \sim p(o_t|s_t) \\\\
\text{that maximizes joint probability of observations} \ln p(o_{1:T})
ELBO 1
Given initial context \(s_0\) (ommited in notation for simplicity)
$$ \ln p(o_{1:T})=\ln\int_{s_{1:T}}p(o_{1:T},s_{1:T}) \ d{s_{1:T}}$$
$$=\ln\int_{s_{1:T}} \prod_{t=1}^T p(o_t,s_t|s_{1:t-1},o_{1:t-1}) \ d{s_{1:T}}$$
$$=\ln\int_{s_{1:T}} \prod_{t=1}^T p(o_t|s_{1:t},o_{1:t-1}) p(s_t|s_{1:t-1},o_{1:t-1}) \ d{s_{1:T}}$$
$$=\ln\int_{s_{1:T}} \prod_{t=1}^T p(o_t|s_{t}) p(s_t|s_{t-1}) \ d{s_{1:T}}$$
$$=\ln\mathop{\mathbb{E}}_{s_{1:T}}[\prod_{t=1}^T p(o_t|s_{t})] $$
ELBO 1
How to sample \( s_{1:T}\) as in the expectation \(\mathop{\mathbb{E}}_{s_{1:T}}[\prod_{t=1}^T p(o_t|s_{t})]\)? Candidate: free-running rollout
\(p(o_{1:T})=\ln\mathop{\mathbb{E}}_{s_{1:T}}[\prod_{t=1}^T p(o_t|s_{t})]\ge \sum_{t=1}^T \mathop{\mathbb{E}}_{s_{1:T}}[\ln p(o_t|s_{t})]\)
ELBO 2
What about importance sampling with posterior
$$\ln p(o_{1:T}) = \ln \mathop{\mathbb{E}}_{s_{1:T}}[\prod_{t=1}^T p(o_t|s_{t})]=\ln \mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1})}[\prod_{t=1}^T p(o_t|s_{t})]$$
$$=\ln\mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1},o_t)}[\prod_{t=1}^T p(o_t|s_{t}) p(s_t|s_{t-1}) / p(s_t|s_{t-1},o_t)]$$
$$=\ln\mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1},o_t)}[\prod_{t=1}^T p(o_t,s_t|s_{t-1}) / p(s_t|o_t,s_{t-1})]$$
$$=\ln\mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1},o_t)}[\prod_{t=1}^T p(o_t|s_{t-1}) ]$$
ELBO 2
$$\ln p(o_{1:T}) =\ln\mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1},o_t)}[\prod_{t=1}^T p(o_t|s_{t-1}) ]$$
$$\ge \mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1},o_t)}[\ln(\prod_{t=1}^T p(o_t|s_{t-1}) ])$$
$$\ge \mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1},o_t)}[\sum_{t=1}^T \ln p(o_t|s_{t-1}) ])$$
Where each \(p(o_t|s_{t-1})\) can be lower bounded
by ELBO of diffusion models, conditioned on s_{t-1}!
ELBO 2
\(\ln p(o_{1:T})\ge \mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1},o_t)}[\sum_{t=1}^T \ln p(o_t|s_{t-1}) ])\)
We can do importance sampling via teacher forcing
Compare ELBO 1 & 2
ELBO1: \(\ln p(o_{1:T})\ge \sum_{t=1}^T \mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1})}[\ln p(o_t|s_{t})]\)
ELBO2: \(\ln p(o_{1:T})\ge \sum_{t=1}^T \mathop{\mathbb{E}}_{s_t\sim p(s_t|s_{t-1},o_t)}[ \ln p(o_t|s_{t-1}) ])\)
ELBO1 samples from prior model \(p(s_t|s_{t-1})\)
ELBO2 samples from posterior model \( p(s_t|s_{t-1},o_t)\)
What about something in between? \( p(s_t|s_{t-1},N(o_t))\), where N(o_t) is a noised version of \(o_t\)
Diffusion Forcing
What about something in between?
\( p(s_t|s_{t-1},N(o_t))\), where \(N(o_t)\) is a noised version of \(o_t\)
\( p(s_t|s_{t-1},N(o_t)) \rightarrow p(s_t|s_{t-1},o_t)\) when \(N\) is zero noise
\( p(s_t|s_{t-1},N(o_t)) \rightarrow p(s_t|s_{t-1},o_t)\) when \(N(o_t)\) is pure noise
We constructed a smooth interpolation between prior and posterior!
Diffusion Forcing (ours)
Diffusion Forcing Sampling
Diffusion Forcing Sampling (horizon 1)
Diffusion
Step (k)
Time step (t)
Diffusion Forcing Sampling (horizon ∞)
Diffusion
Step (k)
Time step (t)
Diffusion Forcing Sampling (pyramid)
Diffusion
Step (k)
Time step (t)
Account for uncertainty under causal dynamics
Diffusion
Step (k)
Time step (t)
Video Prediction (72 frames)
Generated vs GT
Video Prediction (72 frames)
Generated vs GT
Infinite Rollout
train on 72, sample on 500 without sliding window
(500 is due to gif size, can do 3000+ wo blowing up)
64x64
128x128
Time Series Prediction
Control over context-horizon and compositionality
x=[o_t, a_t, r_t]
Planning with Diffusion Forcing
Diffusion Planning
Diffusion Planning
Diffusion Forcing
x=[o_{t+1}, a_t, r_t]
x_t=[o_t, a_t, r_t]
Why is this bad
Now you are at \(o_t\). You want an action. The best you can do is to diffuse \(x_t=[o_t, a_t, r_t]\).
But you already observed \(o_t\)!
Instead
x_t=[o_{t+1}, a_t, r_t]
Joint distribution of policy and dynamics
After receiving ground truth \(o_{t+1}\), update posterior state by \(x_t=[o_{GT, t+1}, a_t, r_t]\)
Now we learn the distribution of data
How to get good policy instead of copying suboptimal policy?
Classifier guidance
Classifier guidance
Classifier guidance
Value guidance
Define probability of task success as
\(\exp(V-V^*)\) where \(V^*\) denotes max possible value
[o_{t+1}, a_t, r_t]
[o_{t+2}, a_{t+1}, r_{t+1}]
[o_{T+1}, a_{T}, r_{T}]
...
Take \(V = r_t + r_{t+1} ... + r_{T}\), do value guidance
You can take \(V = r_t + r_{t+1} ... + r_{T}\) for guidance
But you can also take \(V = r_t\) for guidance
unlike previous methods
Diffusion Forcing is ANY horizon
This doesn't work immediately!
Repeated Sampling
[o_{t+1}, a_t, r_t]
[o_{t+2}, a_{t+1}, r_{t+1}]
[o_{T+1}, a_{T}, r_{T}]
...
x 10 samples
x 10 samples
x 10 samples
x 10 samples
Take \(V = avg(r_t + r_{t+1} ... + r_{T})\), do value guidance
This doesn't work immediately!
Why does this happen?
[o_{t+2}, a_{t+1}, r_{t+1}]
[o_{t+1}, a_t, r_t]
z_t
z_{t+1}
Noisy([o_{t+2}, a_{t+1}, r_{t+1}])
Noisy([o_{t+2}, a_{t+1}, r_{t+1}])
[o_{t+1}, a_t, r_t]
[o_{t+2}, a_{t+1}, r_{t+1}]
Why does this happen?
Perturb beginning of trajectory a bit changes later trajectory by a lot
Change \([o_{t+1}, a_t, r_t]\) should make \([o_{t+1}, a_t, r_t]\) more uncertain
While \(Noisy([o_{t+2}, a_{t+1}, r_{t+1}])\) asks it to be certain
Previous sampling
Diffusion
Step (k)
Time step (t)
N
Denote N=most noisy
N
N
N
N
N
N-1
N-1
N-1
N-1
N-1
N-1
N-2
N-2
N-2
N-2
N-2
N-2
Pyramid sampling
N
N
N
N
N
N
N-1
N-2
N
N
N
N
N
N-1
N
N
N
N
No gradient!
Pyramid sampling
Calculate guidance gradient without pyramid sampling but resample to higher noise level by adding noise
Result
Result
Result
Real Robot
Consider a table with 3 slots. An apple an an orange are at random slots upon initialization. Task is to swap them using the third slot.
This is not markovian!
Real Robot
Just do imitation learning on this, w/ or w/ memory
[o_{t+1}, a_t]
Diffusion Forcing can be used as diffusion policy but with flexible horizon memory
Real Robot
Real Robot
(All images shown are predicted by diffusion)
Result
Diffusion Forcing
- Flexible
- Flexible prediction horizon
- Flexible memory horizon
- Can we auto-regressive / full sequence diffusion at user's choice without retraining
- Easy to incorporate up-coming observation
- Efficient to do MCTS over
- Guidance
- Guidance with consideration of causal uncertainty
- Guidance over longer horizon
Group talk
By buoyancy99
