Latent Diffusion Planning for Imitation Learning
Authors: Amber Xie
Oleh Rybkin
Dorsa Sadigh
Chelsea Finn
Venue: ICML 2025 spotlightposter (2025)
Keywords: imitation learning
diffusion
planning
robotics
Primary area: applications->robotics
TL;DR: We learn a diffusion-based planner and inverse dynamics model in latent space for imitation learning.

Abstract: Recent progress in imitation learning has been enabled by policy architectures that scale to complex visuomotor tasks, multimodal distributions, and large datasets. However, these methods often rely on learning from large amount of expert demonstrations.
To address these shortcomings, we propose Latent Diffusion Planning (LDP), a modular approach consisting of a planner which can leverage action-free demonstrations, and an inverse dynamics model which can leverage suboptimal data, that both operate over a learned latent space. First, we learn a compact latent space through a variational autoencoder, enabling effective forecasting of future states in image-based domains. Then, we train a planner and an inverse dynamics model with diffusion objectives. By separating planning from action prediction, LDP can benefit from the denser supervision signals of suboptimal and action-free data.
On simulated visual robotic manipulation tasks, LDP outperforms state-of-the-art imitation learning approaches, as they cannot leverage such additional data.

Latent Diffusion Planning for Imitation Learning
Amber Xie Oleh Rybkin? 2 Dorsa Sadigh Chelsea Finn'
Abstract
collect. While it is often easier to collect in-domain data that is suboptimal or action-free, these methods not designed to use such data, as they rely on directly modeling optimal actions.
Recent progress in imitation learning has been enabled by policy architectures that scale com- plex visuomotor tasks, multimodal distributions, and large datasets. However, these methods of- ten rely on learning from large amount of expert demonstrations. To address these shortcomings, we propose Latent Diffusion Planning (LDP), a modular approach consisting of planner which can leverage action-free demonstrations, and an inverse dynamics model which can leverage sub- optimal data, that both operate over a learned la- tent space. First, we learn compact latent space through a variational autoencoder, enabling effec- tive forecasting of future states in image-based domains. Then, we train planner and an inverse dynamics model with diffusion objectives. By separating planning from action prediction, LDP can benefit from the denser supervision signals of suboptimal and action-free frre data. On simulated vi- sual robotic manipulation tasks, LDP outperforms state-of-the-art imitation learning approaches, as they cannot leverage such additional data."
Prior works in offline RL, reward-conditioned policies, or imitation learning from suboptimal demonstrations attempt to leverage suboptimal trajectories, though they are still un- able utilize action-free data. Notably, these works often make restrictive assumptions like access to either reward labels (Kumar et al., 2020; Chen et al., 2021; Kumar et al., 2019a). the optimality of demonstrations (Beliaev et al., 2022), or similar metrics (Zhang et al., 2022), which can be impractical or noisy to label. Other works implicitly attempt to use unlabelled, suboptimal data via pretraining on such data and later fine-tuning the policy on the optimal data (Ra- dosavovic et al., 2023; Wu et al., 2023b; Cui et al., 2024). While these approaches can potentially learn representations during pretraining, it does not necessarily improve planning capabilities of these methods.
Our key idea is to take a modular approach, where we sepa- rate learning a video planner from learning an inverse dy- namics model. Each one of these two components can leverage different types of data. For instance, planner can benefit from action-free data, while an IDM can lever- age unlabelled suboptimal data. While using a modular approach has been proposed in recent prior work (Du et al., 2023a; Black et al., 2023), prior approaches focus on high- level decision making by forecasting subgoals, limiting its capabilities for closed loop re-planning in roboties tasks. These works also operate across images, which are high- dimensional and expensive to generate. To create an efficient modular approach that can benefit from all forms of data (suboptimal, action-free, and optimal), we propose learn- ing the planner and inverse-dynamics model over a learned, compact latent space, allowing for closed-loop robot poli- cies. Our imitation learning objective consists of forecasting dense trajectory of latent states, scaling up gracefully to vision-based domains without the computational complexi- ties of video generation.
1. Introduction
Combining large-scale expert datasets and powerful imita- tion learning policies has been a promising direction for robot lcarning. Recent methods using transformer back- bones or diffusion heads (Octo Model Team al., 2024: Kim et al., 2024; Zhao et al., 2024; Chi et al., 2023) have cap- italized on new robotics datasets pooled together from many institutions (Khazatsky et al., 2024; Open X-Embodiment Collaboration et al., 2023), showing potential for learning generalizable robot policies. However, this recipe is funda- mentally limited by expert data, as robotics demonstration data can be challenging, time-consuming. and expensive to
'Stanford 'UC Berkeley. Correspondence to: Amber Xie <am- berxie @ stanford.edu...
We propose Latent Diffusion Planning (LDP), which learns planner that can be trained on action-free data; and an in- verse dynamics model (IDM) that can be trained on data that may be suboptimal. First, it trains variational autoencoder
Proceedings of the 4rnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Project Website and Code: ioluub.oo//dpp
https:/ wwwebxbbbiiiiceeem e

---

Latent Diffusion Planning for Imitation Learning
LDP can use diverse data sources for enhanced training
Latent Diffusion Planning
1. Latent Encoder
II. Planner and IDM
Planner
IDM
a
IDM
IDM
Enc
+
ag
aH
2
+
Action- Pree Demos
Expert Demos
Suboptimal Data
Dec
Enc
2
A
Planner
Enc
Dec
Figure Latent Diffusion Planning. Left: LDP separates control problem into forecasting future states with a diffusion-baset planner, and extracting actions with liffusion-based inverse dynamics model (IDM). This design enables training on heterogeneous sources of data, including suboptimal data and action-free data. Righr: Unlike action imitation methods such as diffusion policy, LDP based on forecasting a dense temporal sequenee of latent states well as actions. Using powerful diffusion models for both of these objectives enables L.PP to have competitive performance to state-of-the-art imitation learning. Further, unlike prior work on forecasting subgoals, LDP predicts dense temporal sequence of latent states, which enables scalable closed-loop planning.
with an image reconstruction loss, producing compressed latent embeddings that used by the planner and inverse dynamics model. Then, it learns an imitation learning policy through two components: (1) a planner, which consumes demonstration state sequences, which may be action-free, and (2) an inverse dynamics model, trained on in-domain, possibly suboptimal, environment interactions. As diffusion objectives have proven be effective for imitation learning in robotics tasks (Chi et al., 2023), use diffusion for both forecasting plans (planner) and extracting actions (IDM), which enables competitive performance. Our method is closed-loop and reactive, as planning over latent space is much faster than generating visually and physically consis- tent video frames.
2. Related Work
Imitation Learning Robotics. One common approach to learning robot control policies is imitation learning, where policies are learned from expert-collected demonstration datasets. This is most commonly done via behavior cloning, which reduces policy learning to supervised learning ob- jective of mapping states to actions. Recently, Diffusion Policy (Chi et al., 2023) and Action Chunking with Trans- formers (Zhao et al., 2023) have shown successful results in complex manipulation tasks using action chunking and more expressive architectures. Diffusion models have also been successful in capturing multimodal human behavior (Pearce et al., 2023). Similarly, Behavior Transformer (Shafiullah et al., 2022) and VQ-BeT (Lee et al., 2024) improve the ability of policies to capture multimodal behaviors. In this work, we focus on forecasting sequence of future states instead of actions, and use diffusion to capture multimodal trajectories.
In summary, our main contributions are threefold:
We propose novel imitation learning algo- rithm, Latent Diffusion Planning, a simple, diffusion planning- basedd method comprised of learned vi- sual encoder, latent planner, and an inverse dynamics model.
Learning from Unlabelled Suboptimal and Action-Free Data. Learning from suboptimal data has long been goal of many robot learning methods, including reinforcement learning. A typical approach is offline reinforcement learn- ing, which considers solving a Markov decision process from an offline dataset of states, actions, and reward (Levine et al., 2020; Kumar et al., 2020; Kostríkov al., 2021; Hansen-Estruch et al., 2023; Yu et al., 2022). Particularly relevant the approaches that supervised learning conditioned on rewards (Schmidhuber, 2019; Kumar et al., 2019a; Chen et al., 2021). In this work, we want to leverage
We show that Latent Diffusion Planning can be trained suboptimal or action- cetion-free data, and improves from learning on such data in the regime where demonstra- tion data is limited.
We experimentally show that our method outperforms prior video planning- based work by leveraging tempo- rally dense predictions a latent space, which enables fast inference for closed-loop planning.
2

---

Latent Diffusion Planning for Imitation Learning
suboptimal, reward-free data, such as play data or failed tra- jectories. In addition, would like to avoid the additional complexity of annotating the data with rewards or training value function which the offline RL methods rely on.
3. Background
Diffusion Models Diffusion models, such as Denoising Diffusion Probabilistic Models (DDPMs), are likelihood- based generative models that learn an iterative denoising process from Gaussian prior to a data distribution (Sohl- Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020). During training time, DDPMs are trained to reverse single noising step. Then, at sampling time, to reverse the diffusion process, the model iteratively denoises a sample drawn from the known Gaussian prior.
Several works have also addressed learning from action- free data, such as using inverse models (Torabi et al., 2018; Baker et al.. 2022). latent action models (Edwards et al., 2019; Schmeckpeper et al.. 2020; Bruce al., 2024), or representation learning (Radosavovic et al., 2023; Wu et al., 2023b; Cui et al., 2024). In this work we focus on a simple recipe for robotic imitation learning that is naturally able to leverage action-free data through state forecasting.
Diffusion models may also be conditioned on additional context. For example, text-to-image generative models are conditioned on text, Diffusion Policy is conditioned on vi- sual observations, and Decision Diffuser can be conditioned on reward, skills, and constraints.
Diffusion and Image Prediction in Robot Learning. Dif- fusion models, due to their expressivity and training and sampling stability, have been applied to robot learning tasks. Diffusion has been used in offline reinforcement learning (Hansen-Fstruuch et al., 2023) and imitation learn- ing (Chi et al., 2023). Diffuser (Janner et al., 2022) learns a denoising diffusion model on trajectories, including both states and actions, in a model-based reinforcement learn- ing setting. Decision Diffuser (Ajay et al., 2023) extends Diffuser by showing compositionality over skills, rewards, and constraints, and instead diffuses over states and uses an inverse dynamics model to extract actions from the plan. Due to the complexity of modeling image trajectories, Dif- fuser and Decision Diffuser restrict their applications to low-dimensional states.
Recent generative models have used Latent Diffusion Mod- els, which trains a diffusion model in a learned, compressed latent space (Rombach et al., 2022; Peebles & Xie, 2023; Blattmann et al., 2023) to improve computational and mem- ory efficiency. The latent space is typically learned via an autoencoder, with encoder & and decoder 'D trained to re- construct x R â = D(E(x)). Instead of diffusing over x, the diffusion model is trained on diffusing over z = E(x).
Imitation Learning In the imitation learning framework, we assume access dataset of expert demonstrations, (ET,aT) aT) }, generated by TEg an
D á { Ho, a.),..., (ST, IT, expert policy. T4, as correspond to the state, image, and action at timestep i respectively. The imitation learning ob- jective is to extract policy ((aa 8, x) that most closely imi- tates Ty. In robotics, this is typically approached through behavior cloning, which learns the mapping between states and actions directly via supervised learning. We consider single-task imitation, where the dataset corresponds to single task.
To scale up to diffusing over sigees-dimensional plans, UniPi (Du et al., 2023a; Ko et al., 2023) adapts video models for planning. Unlike works that rely foundation models and video models for planning (Du et al., 2023b; Yang et al., 2024; Zhou et al., 2024), our method avoids computational and modeling complexities of generative video modeling by planning over latent embeddings instead.
Previous works have used world models to plan over images in compact latent space (Hansen et al., 2024; Hafner et al., 2019; 2020). In contrast with these works, we focus single task imitation instead of reinforcement learning.
Diffusion Policy (Chi et al., 2023) is an instantiation of diffusion models for imitation learning that has shown suc- cess in simulated and real-world tasks. Diffusion Policy uses a DDPM objective model the distribution of action sequences, conditioned on observations. The CNN instan- tiation uses a Conditional U-Net Architecture, based on the 1D Temporal CNN in (Janner et al., 2022), which en- courages temporal consistency due to the inductive biases of convolutions. LDP's s planner architecture is based on the CNN- Diffusion Policy, though we forecast latent states instead of actions.
Many prior works argue that state forecasting objectives are uniquely suitable for robotics to improve planning quality with trajectory optimization or reinforcement learning (Finn & Levine, 2017: Yang et al., 2023), by using the model directly to plan future states (Du et al., 2023b;a), as well as representation learning (Wu et al., 2023a; Radosavovic et al., 2023). We follow this line of work by proposing a planning- based method competitive to state-of-the-art robotic imita- tion learning that can leverage heterogeneous data sources.
Datasets of expert demonstrations often do not provide suffi- cient state distribution coverage to effectively solve a given task with imitation learning. However, there often exists additional data in the form of action-free or suboptimal data, which may consist of failed policy rollouts, play data, or miscellaneous environment interactions. Unfortunately.
3

---

Latent Diffusion Planning for Imitation Learning
behavior cloning assumes access data annotated with optimal actions, such additional data cannot be easily incorporated into training.
Algorithm 1 Inference with Latent Diffusion Planning
Input: Encoder &, Planner cg. IDM Eç, Planner Dif- fusion Timesteps Tр IDM Diffusion Timesteps TIMM, Planning Horizon Hp. Action Horizon Ha
4. Latent Diffusion Planning
Latent Diffusion Planning consists of three parts, as shown Figure 1: (1) Training an image encoder via an image reconstruction loss, (2) learning inverse dynamics model to extract actions a, from pairs of latent states zt. Z++11 and (3) learning a planner to forecast future latents
2: Observe imitial state and image roz k = 0
3: while not done do
Z + (E(xx), sk)
Diffuse over latent embedding plan
5:
z k++ K+H (0, for / T Tp... 1.......
6:
1do
Inverse Dynamics Model
7:
é E Evi (2 (+kt..... H, Zk.
L IDM
8:
Update &. k411..., 2 k+ Hp using DDPM update with
Enc
end for
IDM Diffusion Model
// Diffuse over actions between latent embeddings for ;- 0 Ho
Sui
Enc
10:
do
11:
ākts ~ (0 // Predict action for each timestep in
action horizon
Latent Planner
12:
for t TIDM..... do
S+11
S.
13:
14:
Update āR+t using DDPM update with
Enc @
Enc
15:
end for
16:
end for
Zyes
Z
Execute actions
planner
17:
for i = 0... ...... Ha do
2
18: 5k+i+1 f- env.step(*k ax+i)
Planner Diffusion Model
Enc
2
19: end for
20: k Ha
21:
end while
Figure 2. After training the encoder, Latent Diffusion Planning trains two diffusion models. Top: We train a inverse dynamics model (IDM) with diffusion objective to directly extract the actions that will be used for control from pairs of latent states. Bottom: train a powerful latent diffusion model to forecast chunk of future latent states. The planner and the IDM are used together to produce an action chunk, similar (Chi et al., 2023).
& Welling, 2014; Rezende al., 2014) to obtain latent encoder & and decoder D. Specifically, we optimize the B VAE (Higgins et al., 2017) objective, where x is our original image, z our learned latent representation of the image, are the parameters for our decoder, & are the parameters for our encoder, and is the weight for the KL regularization term:
4.1. Learning the Latent Space
We circumvent planning over high- nigh-dimensional image ob- servations by planning over learned latent space. Similar prior work planning with world models (Watter al., 2015; Ha & Schmidhuber, 2018; Hafner et al., 2020), learn this latent space using an image reconstruction objec-
(x|xx) (log (x zz]
(I)
BDLL x))lp(z))
In practical scenarios, we may have limited expert demon- stration dataset, but much larger and diverse suboptimal or action-free free datasets. In this phase of learning, we can make use of the visual information in such datasets for training a more robust latent encoder.
tive. Our planner thus becomes similar to video models that forecast image frames in a learned latent space (Yan et al., 2021; Hong et al., 2022: Blattmann et al., 2023).
In this work, we train a variational autoencoder (Kingma

---

Latent Diffusion Planning for Imitation Learning
4.2. Planner and Inverse Dynamics Model
Is Latent Diffusion Planning comparable state-of- the-art imitation learning algorithms that leverage sub- optimal data?
Our policy consists of two separate modules: (1) a plan- ner over latent embeddings and (2) an inverse dynamics model similarly operating over the latent embeddings. The planner and IDM are both parameterized as DDPM models, motivated by the expressivity that diffusion models offer.
Can Latent Diffusion Planning be an effective imita- tion learning method in real-world robotics system, where there may be suboptimal or action-free data?
The planner is conditioned on the current latent embedding, which consists of the concatenated latent image embedding and robot propriaception, and diffuses over horizon of future embeddingg. We use Diffusion Policy's Conditional U-Net architecture, with CNN backbone. Concretely, we optimize the following objective:
5.1. Experimental Setup
Simulated Tasks We focus our experiments on 4 image- based imitation learning tasks: (1) Robomimic Lift, (2) Robomimic Can, (3) Robomimic Square, and (4) ALOHA Sim Transfer Cube. Robomimic (Mandlekar al., 2021) robotic manipulation and imitation benchmark, including the tasks Lift, Can, and Square. The Transfer Cube task is simulated bimanual ALOHA task, in which one ViperX 6-DoF arm grabs a block and transfers it to the other arm (Zhao et al., 2023).
Cptamner (1, z) = (f. 4k. - <|1-1
(2)
where Zy is the latent embedding timestep of the tra- jectory; ž ..+1....;2... Zz+H is the noised latent embedding sequence, with corresponding noise e; H is the maximum horizon of the forecasted latent plan; is the diffussion noise timestep; and & are the parameters of the planner diffusion model.
Lift
Can
Square
ALOHA Cube
Franka Lift
Our inverse dynamics model trained to reconstruct the ac- tion between a pair of states, conditioned on their associated latent embeddings. We use the MLPResNet architecture from IDQL EHansen-Estruch et al., 2023) to diffuse actions, as it is more lightweight. We optimize the loss:
To demonstrate the effectiveness of Latent Diffusion Plan- ning, we assume low demonstration data regime, such that additional suboptimal or action-free data can improve performance. For Can and Square, we use 100 out of the 200 demonstrations in the Robomimic datasets; for Lift, we use 3 demonstrations out of the 200 total; and for Trans- fer Cube, we use 25 demonstrations. To further emphasize the importance of suboptimal data, these demonstrations cover limited state space of the environment. Our subopti- mal data consists of 500 failed trajectories from an under- trained behavior cloning agent. Our action free data consists of 100 demonstrations for Lift, Can, and Square from the Robomimic dataset, and 25 demonstrations for Cube. We evaluate the success rate out of 50 trials, using the best checkpoint from the last 5 saved checkpoints, with 2 seeds.
(EDM(E z) = [lee (āx; (aaiZz, Zk, Ze+1, () - el|2]
where z is the latent embedding at timestep k of the trajec- tory; 0g is the noised action, with corresponding noise e; is the diffusion noise timestep; and & are the parameters of the inverse dynamics diffusion model.
Because our latent embedding is frozen from the learned VAE, the planner and IDM do not share parameters and can be trained separately (Figure 2). Then, at inference time, the two modules are combined to extract action sequences. First, the planner forecasts a future horizon of states via DDPM sampling (Alg. 1 Lines 6-9). Then, we use the inverse dynamics model to extract actions from latent embedding pairs produced by the planner, also via DDPM Diffusion (Alg. 1 Lines 12-15). Like Diffusion Policy, we employ siceding-horiznt control (Mayne & Michalska, 1988), and execute for shorter horizon than the full forecasted horizon (Alg. Lines 17-19).
Real World Task We create a real world implementation of the Robomimic Lift task, where the task is to pick up a red block from randomly initialized position. We use Frank Panda 7 degree of freedom robot arm, with a wrist-mounted Zed camera. We use the DROID setup (Khazatsky et al., 2024) and teleoperate via the Oculus Quest 2 headset. We use cartesian pose control.
Ocalaa Inteoperation
Red wlnt camers
5. Experiments
candorily initialized
We seek to answer the following questions:
Does Latent Diffusion Planning leverage action-free data for improved planning?
5

---

Latent Diffusion Planning for Imitation Learning
Table Leveraging Action-Free Data. LDP is able to leverage suboptimal and action-frte free data. We compare against DP baselines that leverage action-free data by using an IDM to relabel actions (DP-VPT) and video planning models that may use action-frve free data for the video planner (UniPi). We find LDP to perform better than both approaches, especially when combined with suboptimal data.
Method
Lift
Can
Square
ALOHA Cube
Average
DP
0.60
0.00
0.63 0.01
0.48 + 0.00
0.32 + 0.00
0.51
DP-VPT
0.69 0.01
0.75 0.01
0.48 4: 0.04
0.45 4: 0.03
0.59
UniPi-OL + Action-Free
0.09 + 0.05
0.23 + 0.03
0.07 +: 0.03
0.02 :: 0.00
0.11
UniPi-CL.+ + Action-Free
0.14 0.02
0.32 1: 0.04
0.09 1: 0.01
0.17 1: 0.03
0.18
LDP
0.69
0.03
0.70 1: 0.02
0.46 1: 0.00
0.64
1: 0.04
0.65
LDP + Action-Free
0.67
+ 0.01
0.78
# 0.04
0,47
+: 0.03
0.70
= 0.02
0.66
LDP + Action-Free + Subopt
1.00 +: 0.00
0.98 +: 0.00
0.83 +: 0.01
0.97 :: 0.01
0.95
We collect 82 demonstrations, collect 84 suboptimal trajectories, and 12 action-free demonstrations.
demos. This is representative of methods that lever- age suboptimal data through learning trajectory-level features.
Our suboptimal data consists of failed trajectories from pol- icy evaluations. This is an effective way to reuse the data generated during iterations of training, that algorithms mod- eling actions, such Diffusion Policy, cannot use. Action- Free data may consist of kinesthetic demonstrations, human videos, or handheld demonstrations (Chi et al., 2024). In our case, we collect teleoperated demonstrations with actions removed.
Diffussion Policy with Video PreTraining (DP-VPT) trains an inverse dynamics model to relabel action- free data, inspired by VPT (Baker al., 2022) and BCO (Torabi et al., 2018).
Open-Loop UniPi (UniPi-OL) is based off of UniPi (Du et al., 2023a), a video planner for robot manipulation. UniPi-OL generates a single video tra- jectory, extracts actions, and executes the actions in an open-loop fashion. We use a goal-conditioned behavior cloning agent to reach generated subgoals (Wen et al., 2024).
To evaluate our policies, we calculate the success rate across 45 evaluation trials. To thoroughly evaluate performance across the initial state space, we evaluate across grid of 3x3 points, with 5 attempts per point. We evaluate 3 seeds per method.
Closed-Loop UniPi (UniPi-CL) is modification that allows UniPi to perform closed-loop replanning over image chunks. Like LDP, UniPi-CL generates dense plans instead of waypoints, though in image space. We learn an inverse dynamics model to extract actions.
Baselines We consider two main categories of baselines: (1) Imitation learning with suboptimal or action-free data (DP, RC-DP, DP+Repr, DP PT FT, DP-VPT), and (2) Video planning (UniPi-OL,, UniPi-CL.).
Diffusion Policy (DP) is a state-of-the-art imitation learning algorithm.
5.2. Imitation Learning with Action-Free Data
In Table 1, we examine how action-free data can be used to improve imitation learning policies. Imitation learning poli- cies that model actions, such as DP, are unable to natively use action- free data, while planning-basee approaches can benefit from this additional data.
Diffusion Policy (RC-DP) uti- lizes suboptimal actions by conditioning the pol- icy on a binary value indicating whether the action chunk comes from optimal demonstrations or not. This method is inspired by reward-conditioned ap- proaches (Kumar et al., 2019b; Chen et al., 2021).
One approach is to relabel action-fres data using an inverse dynamics model. We find this to be effective for most tasks, showing that the reannotated actions are useful for policy improvement. However, we see that LDP is better able to leverage action-free data by directly using it for the planner, rather than generating possibly inaccurate actions to subsequently learn from.
Diffusion Policy with Representation Learning (DP+Repr) uses a VAE pretrained on demonstration, suboptimal, and action-free data as the observation en- coder. This is representative of methods that leverage suboptimal data through representation learning.
Diffusion Policy Pretrain + Finetune (DP PT + FT) pretrains on suboptimal trajectories and finetunes on
Like LDP, video planning methods can directly use action- free data for improving the planner, which for UniPi the
6

---

Latent Diffusion Planning for Imitation Learning
Table Leveraging Suboptimal Data. Latent Diffusion Planning can utilize suboptimal data for improved performance across a suite tasks. Video planning methods like UniPi consistently struggle, whereas reward-condition conditoond or pretraining DP approaches are able to improve DP policy performance. LDP best able to use suboptimal data, especially when combined with action-free data.
Method DP
Lift
Can
Square
ALOHA Cube
Average
0.60 0.00
0.63
0.01
0.48 + 0.00
0.32 = 0.00
0.51
RC-DP
0.40 0.04
0.73
0.03
0.66 +: 0.02
0.60 :: 0.04
0.60
DP+Repr DP PT
0.66 0.04
0.61
+ 0.01
0,44 +: 0.02
0.25 +: 0.03
0.49
0.52 0.02
0.67 = 0.01
0.57 + 0.03
0.78 = 0.00
0.64
UniPi-OL
0.12 0.06
0.28 :: 0.02
0.07 1: 0.01
0.00 1: 0.00
0.12
UniPi-CL
0.12 0.02
0.30 0.02
0.10 4: 0.04
0.15 +: 0.07
0.17
LDP
0.69 +: 0.03
0.70 4: 0.02
0,46 4: 0.00
0.64 :: 0.04
0.65
LDP + Subopt
0.84 = 0.06
0.68 = 0.02
0.55 + 0.03
0.71
+ 0.03
0.70
LDP + Action-Free Subopt
1.00 0.00
0.98 + 0.00
0.83 + 0.01
0.97 + 0.01
0.95
video generation model. We see limited improvement when using action-free data, and both DP and I.DP still strongly outperform UniPi-CL and UniPi-OL. UniPi methods par- ticularly struggle with more complex tasks like Square and ALOHA Cube, implying that video planning is still ineffec- tive for tasks that require more precise manipulation skills.
motions of reaching toward or grasping the object, which are partially covered by the suboptimal dataset, provides useful visuomotor prior for the policy, ALOHA Cube sees significant improvement, possibly because the larger action space of bimanual control benefits from additional reward-labelled data.
Qualitatively, for UniPi-OL, we find that while the goal- conditioned agent is able to follow goals effectively, the policy still struggles with the difficult parts of the task, such as grasping the object. Forecasting goals does not provide the dense supervision for exactly how to grasp an object, and furthermore, UniPi-OL does not support replanning when grasp is missed. UniPi-CL is able to address this by dense image forecasting, and consistently outperforms UniPi-OL. However, this closed-loop method is not only slow, but faces issues with video generation, such as regenerating static frames during parts of the task with less movement, leading the agent to be stuck in certain positions. This is especially noticeable for the ALOHA Cube task, where the agent is often stuck right before picking up the cube. Compared to UniPi-OL and UniPi-CL, LDP is able to circumvent many of these issues due to its latent planning and dense forecasting.
Next, we explore using suboptimal data for feature-learning. DP + Repr uses suboptimal data for pretraining a vision- encoder, and DP PT FT for pretraining the visuomotor policy. DP+ + Repr only improves policy performance for the Lift task, implying that end-to-end training of the vi- sion encoder learns stronger features for complex manipula- tion tasks. DP PT + FT is particularly successful for tasks that require more precise manipulation, such as Square and ALOHA Cube, implying that learned prior motions is a use- ful policy initialization. Both the success of RC-DP and DP PT FT suggest that leveraging the suboptimal trajectories is useful way improve imitation learning results. LDP, which leverages suboptimal data for the latent encoder and IDM, has higher overall performance than these methods, averaging across the suite of simulated tasks.
Next, we compare against UniPi, which plans over image subgoals (OL.) or image chunks (CL.) Due to the low demon- stration data regime, learning effective and accurate video policies is difficult, and LDP strongly outperforms UniPi- OL and UniPi-CL. In addition, we notice that UniPi-CL outperforms UniPi-OI. for all tasks, implying that dense forecasting, even within the image domain, is more effec- tive than goal-conditioned methods
5.3. Imitation Learning with Suboptimal Data
In Table 2. we present imitation learning results with sub- optimal data. First, LDP outperforms DP, which can only utilize data with optimal actions. We notice, especially, that LDP with suboptimal data typically improves further upon LDP, showing the potential of leveraging diverse data sources outside of the demonstration dataset.
Finally, we find that LDP with action-free and suboptimal data leads to the strongest performance. We find a signif- icant improvement in all of the simulated tasks, including compared to the variants of LDP that only use either action- free or suboptimal data. This suggests that the recipe for combining these two data sources can lead to a stronger
Next, RC-DP, a conditional variant of DP that utilizes sub- optimal data, outperforms DP. By learning from suboptimal data, RC-DP can learn priors of robot motions while dis- tinguishing optimal action sequences. We hypothesize that for the Can, Square, and ALOHA Cube tasks, the primitive
7

---

Latent Diffusion Planning for Imitation Learning
in (Chi et al., 2024), can be scalable direction.
Table Real World Results For the Franka Lift task, we find that LDP consistently outperforms DP. Suboptimal and action-free data, that DP is unable use, further improves performance.
5.5. Ablation: LDP Hierarchical
Method
Success Rate
In an attempt to understand the effect of LDP's dense fore- casting, we compare LDP with hierarchical version of LDP (LDP Hierarchical). In this implementation, LDP Hi- erarchical plans over subgoals steps apart, and extracts 4 actions between pairs of forecasted latent states. Thus, LDP Hierarchical is closed-loop, yet the planner operates at a slightly abstracted level compared to non- non-hierarchical LDP.
DP
69.6 =
LDP
73.3 =
LDP + Subopt
74.8
LDP Action-Free
79.3
planner and a more robust inverse dynamics model, leading to the best performing model for these tasks.
In Table 4, we find that LDP outperforms LDP Hierarchical across the 3 Robomimic tasks. This suggests that dense forecasting is important part and contribution of LDP. This also reflects UniPi results in Table I and Table 2 that show closed-loop planning outperforms open-loop planning.
5.4. Imitation Learning in the Real World
Real world data is more expensive and consuming to collect;
hence, examining the effect of casier-to-collect sub- optimal and action-free data provides insights for scalable learning. In Table 3, provide results Franka Lift Cube task. In this task, examine the performance of DP, which can only leverage action-labeled data, with our method, which can use suboptimal and action- free data.
Table 4. Hierarchical Ablation. LDP's dense forecasting oulper- forms a hicrarchical variant of LDP.
Method
Lift
Can
Square
LDP
0.69 0.03
0.70 +: 0.02
0.46 : 0.00
LDP Hier.
0.53 0.05
0.62 0.04
0.31 0.01
LDP + Subopt
0.84
0.68 4:
0.55 0.03
LDP Hier. + Subopt
0.65 0.05
0.60 + 0.02
0,43 # 0.05
6. Discussion
We presented Latent Diffusion Planning, a simple planning- based method for imitation learning. We show that our design using powerful diffusson models for latent state fore- casting enables competitive performance with state-of-the- imitation learning. We further show this latent state forecasting objective enables us easily leverage heteroge- neous data sources. In the low-demonstration data imitation regime, LDP outperforms prior imitation learning work that does not leverage such additional data as effectively.
Limitations. One limitation the current approach is that the latent space for planning is simply learned with vari- ational autoencoder and might not learn the most useful features for control. Future work will explore different rep- resentation learning objectives. Further, method requires diffusing over states, which incurs additional computational overhead as compared to diffusing actions. However, we expect continued improvements in hardware and inference speed will mitigate this drawback. Finally, we did not ex- plore applying recent improvements in diffusion models (Peebles & Xie, 2023; Lipman et al., 2022), which may be important in large data regimes.
Figure 3. Visualizations of Generated Plans. LDP produces dense, closed-loop plans. Here, we visualize decoded latents se- lected from an LDP trajectory for the Lift, Can, Square, ALOHA Cube, and Franka Lift tasks.
We find that LDP is able to consistently outperform DP, pecially with the addition of action-free data. For real-world systems, this is promising, as collecting high-quality demos can be difficult and time intensive, whereas utilizing subop- timal trajectories, often a byproduct of evaluating policies, or collecting action-free trajectories more efficiently, such
Future work. We have validated in simulation and real the hypothesis that latent state forecasting can leverage het- erogencous data sources. Future work can also evaluate whether this can be used to further improve more complex
8

---

Latent Diffusion Planning for Imitation Learning
real-world tasks. One direction is to use a diverse dataset of human collected data, such with handheld data collec- tion tools (Young et al., 2021), Another approach would be to use autonomously collected robotic data (Bousmalis* et al., 2023). As these alternative data sources are easier to collect than demonstrations, they represent a different scaling paradigm that can outperform pure behavior cloning approaches. By presenting a method that can leverage such data, we believe this work makes step toward more perfor- mant and general robot policies.
Bousmalis*, K., Vezzani*, G., Rao*, D., Devin C., Lee
A. X., Bauza*, M., Davchev T., Zhou*, Y., Gupta", A., Raju, A., Laurens, A., Fantacci, C., Dalibard, V., Zambelli, M., Martins, M., Pevceviciute, R., Blokzijl, M., Denil, M., Batchelor, N., Lampe, T., Parisotto, E.,, Zotna, K., Reed, S., Colmenarejo, S. G., Scholz, J., Abdolmaleki,
A., Groth, O., Regli, J.-B., Sushkov, O., Rothörl, T., Chen, J. E., Aytar, Y., Barker, D., Ortiz, J., Riedmiller, M.. Springenberg, J. T., Hadsellt, R., Norit, F., and Heess, N. Robocat: A self-improving foundation agent for robotic manipulation. ArXiv, 2023.
Impact Statement
Bruce, J., Dennis, M. D., Edwards, A., Parker-Holder J., Shi. Y., Hughes, E., Lai, M., Mavalankar, A., Steiger- wald, R., Apps, C., et al. Genie: Generative interactive environments. In Forty-first International Conference on Machine Learning, 2024.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be
specifically highlighted here.
Chen, L., Lu, K., Rajeswaran, A.. Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. 2021. URL https: http///axxvvvv///// org/abs///// 2106. 01345.
Acknowledgments
We would like to thank members of ILIAD and IRIS at Stan- ford for fruitful discussions and feedback. We are grateful to Archit Sharma, Annie Chen, Austin Patel, Hengyuan Hu, Jensen Gao, Joey Hejna, Mallika Parulekar, Mengda Xu, Philippe Hansen-Estruch, Shuang Li, Suvir Mirchandani, and Zeyi Liu for helpful conversations.
Chi, C., Feng, S., Du, Y., Xu, Z., Cousineau, E., Burchficl, B., and Song, S. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
References
Chi, C., Xu, Z., Pan, C., Cousineau, E., Burchfiel, B., Feng, S., Tedrake, R., and Song, S. Universal manipulation interface: In-the-wild robot teaching without in-the-wild robots. In Proceedings of Robotics: Science and Systems (RSS), 2024.
Ajay, A., Du, Y., Gupta, A., Tenenbaum, J., Jaakkola, T., and Agrawal, P. Is conditional generative modeling all you need for decision-making', 2023. URL https: //araaaa //aaxivvv.. www.aiiirrg/aa..aaaaao...aoooeae abs/ /2211 15657.
Baker, B., Akkaya, I., Zhokov, P., Huizinga, J., Tang, J., Ecoffet, A., Houghton, B., Sampedro, R., and Clune, Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances Neural Information Processing Systems, 35:24639-24655 2022.
Cui, Z. J., Pan, H., Iyer, A., Haldar, S., and Pinto, L. Dy- namo: In-domain dynamics pretraining for visuo-motor control. arXiv preprint 10X1122409.1219,, 2024.
Darcet, T., Oquab, M., Mairal, J., and Bojanowski, P. Vision transformers need registers, 2023.
Beliaev, M., Shih, A., Ermon, S., Sadigh, D., and Pedarsani, R. Imitation learning by estimating expertise of demon- strators, 2022. URL https: /atpiv//axvvvvv// org/ abs/ 2202.
Du, Y., Yang, M., Dai, B., Dai, H., Nachum, O., Tenenbaum, J. B., Schuurmans, D.,, and Abbeel, Learning universal policies via texted video generation. 2023a. URL https:/ / araii. org abs / 220.2 00111.
Black, K., Nakamoto, M., Atreya, P., Walke, H., Finn, C., Kumar, A., and Levine, S. Zero-shot robotic manipulation with pretrained image-editing diffusion models, 2023. URL https:tp////.. /ttp//axiivv//////vv///////..ava org/
Du, Y., Yang. M., Florence, P., Xia, F., Wahid, A., Ichter, B., Sermanet, P., Yu, T., Abbeel, P., Tenenbaum, J. B., Kaelbling, L., Zeng, A., and Tompson, J. Video language planning. 2023b. URL https: /ttp:::////vvvv.a//aaeeaeeaasaaa /aap/axxvvvvoooa////a org.io.orrr.....o org/abs//////o/ooo///......b 2310 10625.
10639.
Blattmann, A., Rombach, R., Ling, H., Dockhorn, T., Kim, S. W., Fidler, S., and Kreis, K. Align your latents: High- resolution video synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22563-225755 2023.
Edwards, A., Sahni, H., Schroecker, Y., and Isbell, C. Imi- tating latent policies from observation. In International conference on machine learning, pp. 1755-1763 PMLR, 2019.

---

Latent Diffusion Planning for Imitation Learning
Finn, C. and Levine, S. Deep visual foresight for planning robot motion. In 2017 IEEE International Conference Robotics and Automation (ICRA), pp. 2786-2793. IEEE, 2017.
Kingma, D. P. and Welling, M. Auto-encoding varia- tional bayes, 2014. URL https://aaxiv..rrrr
abs/1312.611
Ko, P.-C., Mao, J., Du, Y., Sun, S.-H., and Tenenbaum, J. Learning to Act from Actionless Videos through Dense Correspondences. arXiv:2310.08577 2023.
Ha, D. and Schmidhuber, J. World models. arXiv preprint arXiv: 1803. 10122, 10122, 2018.
Hafner, D., Lillicrap. T., Fischer, I,, Villegas, R., Ha, D.,
Kostrikov, I., Nair, A., and Levine, S. Offline reinforcement
Lee, H., and Davidson, J. Learning latent dynamics for planning from pixels, 2019. URL https: aariiv. org/aba abs 111811.11 04551
learning with implicit q-learning, 2021. URL https: //arxiv. org/ abs /2100.2111..0000 06169.
Kumar, A., Peng, X. B., and Levine, S. Reward-conditioned policies. urXiv preprint arXiv: 13465, 2019a.
Hafner, D., Lillierap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination, 2020.
URL Uttps:::: //ae:i///xiivv
org/ abs 91122 2. 0100016600 01.03.
Kumar, A., Peng, X. B., and Levine, Reward-conditioned policies, 2019b. URL, 1912. 13465.
Hansen, N., Su, H., and Wang, X. Td-mpe2: Scalable, robust world models for continuous control, 2024. URL patpss//arxxr org/ abs /2l0/ 6828.
Kumar, A., Zhou, A., Tucker, G., and Levine, S. Conserva- tive q-learning for offline reinforcement learning, 2020. URL https: /ariii///aaai////aa//eaaeaiaiaaa http///g////a/oo//////oo
Hansen-Estruch, P., Kostrikov, I., Janner, M., Kuba, J. G., and Levine, S. Idql: Implicit q-learming as an actor- critic method with diffusson policies, 2023. URL htts:::: //arxiv. org/ abs 2los/2360 4304 10573.
04779.
Lee, S., Wang, Y., Etukuru. H., Kim, H. J., Shafiullah, N. M. M., and Pinto, I.. Behavior generation with latent actions, 2024. URL https: htrps://arxiv.org org/bbs/ 2403.0311 03181
Higgins, I. Matthey, L. Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. beta- VAE: Learning basic visual concepts with a con- strained variational framework. In International Confer- ence on Learning Representations, 2017. URL ht htsps: //openreview net forum? 2d-
Levine, S., Kumar, A., Tucker, G., and Fu, Offline rein- forcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint wwXiv:2005.0666 2020.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. arXiv preprint arXiv:2210.02744 2022.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion prob- abilistic models, 2020. URL.
abs/2006. 11239.
Mandlekar, A., Xu, D., Wong, J., Nasiriany, S., Wang, C., Kulkarni, R., Fei-Fei, I.., Savarese, S., Zhu, Y., and Martin, R. What matters in learning from of- fline human demonstrations for robot manipulation. In Conference on Robot Learning (CoRL), 2021.
Hong, W., Ding, M., Zheng, W., Liu, X., and Tang, J. Cogvideo: Large-scale pretraining for text-to- text-to-video gener- ation via transformers. arXiv preprint arXiv:2205. 15868, 2022.
Janner, M., Du Y., Tenenbaum, J. B., and Levine, S. Plan- ning with diffusion for flexible behavior synthesis, 2022.
Mayne, D. Q. and Michalska, H. Receding horizon control nonlinear systems. In Proceedings of the 27th IEEE Conference on Decision and Control, pp. 464-465. IEEE, 1988.
URL httpt::///aaaae //arxiv///vv.v.a org/ abs/ 220.. 09991.
Khazatsky, A., Pertsch, K., Nair, S., Balakrishna, A.,
Dasari, S., Karamcheti, S., Nasiriany, S., Srirama, M. K., Chen, L.. Y., Ellis, K., et al. Droid: A large-scale in-the-wild robot manipulation dataset. arXiv preprint arXiv:2403.12945 2024.
Octo Model Team, Ghosh, D., Walke, H., Pertsch, K., Black, K., Mees, O., Dasari, S., Hejna, J., Xu, C., Luo, J., Kreiman, T., Tan, Y., Chen, L. Y., Sanketi, P., Vuong, Q., Xiao, T., Sadigh, D., Finn, C., and Levine, S. Octo: An open-source generalist robot policy. Proceedings of Science and Systems, Delft, Netherlands, 2024.
Kim, M., Pertsch, K., Karamcheti, S., Xiao, T., Balakr- ishna, A., Nair, S., Rafailov, R., Foster, E., Lam, G., Sanketi, P., Vuong, Q., Kollar, T., Burchfiel, B., Tedrake, R., Sadigh, D., Levine, S., Liang, P., and Finn, C. Open- model. arXiv
Robotics:
Open X-Embodiment Collaboration, O' O'ill,,, A., Rehman, A., Gupta, A., Maddukuri, A., Gupta, A., Padalkar, A., Lee, A., Pooley, A., Gupta, A., Mandlekar, A., Jain, A.,
vla: An open-source disiion-annuaeee. preprint Xoviv:2406 09246, 2024.
Tung, A., Bewley, A., Herzog, A., Irpan, A., Khazatsky,
10

---

Latent Diffusson Planning for Imitation Learning
A., Rai, A., Gupta, A., Wang, A., Kolobov, A., Singh, A., Garg, A., Kembhavi, A., Xie, A., Brohan, A., Raffin, A., Sharma, A., Yavary, A., Jain, A.. Balakrishna, A., Wahid, A., Burgess-Limerice B., Kim, B., Schölkopf, B., Wulfe, B., Ichter, B., Lu, C., Xu, C., Le, C., Finn, C., Wang. C.,, Xu, C., Chi, C., Huang, C., Chan, C., Agia, C., Pan, C., Fu, C., Devin, C., Xu, D., Morton, D., Driess, D., Chen, D., Pathak, D., Shah, D., Büchler, D., Jayaraman, D., Kalashnikov, D., Sadigh, D., Johns, E., Foster, E.,, Liu, F. Ceola, F., Xia, F., Zhao, F., Frujeri, F. V., Stulp, F,, Zhou, G., Sukhatme, G. S., Salhotra, G., Yan, G., Feng, G. Schiavi, G., Berseth, G., Kahn, G., Yang, G., Wang, G. Su, H., Fang, H.-S., Shi, H., Bao, H., Amor, H. B., Christensen,
Z., Cui, Z. J., Zhang, Z., Fu, Z., and Lin, Z. Open X- Embodiment: Robotic learning datasets and RT-X models. https: http:///xxvvvrr.g...ee wttpoos/g/a/aaa/eeee 08864, 2023.
Oquab, M., Darcet, T.,, Moutakanni, T., Vo, H., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza, D., Massa, F., El- Nouby, A., Howes, R., Huang, P.-Y., Xu, H., Sharma, V., Li, S.-W., Galuba, W., Rabbat, M., Assran, M., Ballas, N., Synnaeve, G., Misra, L., Jegou, H., Mairal, J., Labatut, P., Joulin, A., and Bojanowski, P. Dinov2: Learning robust visual features without supervision, 2023.
Pearce, T., Rashid, T., Kanervisto, A., Bignell, D., Sun, M., Georgescu, R., Macua, S. V., Tan, S. 7.,, Momennejad, I., Hofmann, K., and Devlin, S. Imitating human behaviour with diffusion models, 2023. URL https: https:////aaat //arxiv. org/ /bs/2301.0
I., Furuta, H., Bharadhwaj. H., Walke, H., Fang, H., Ha, H., Mordatch, I., Radosavovic, I., Leal, I. Liang, J., Abou-Chakra, J., Kim, J., Drake, J., Peters, J. Schneider, J., Hsu, J., Vakil, J., Bohg, J., Bingham, J. Wu, J., Gao, J., Hu, J., Wu, J., Wu, J., Sun, J., Luo, J. Gu, J., Tan, J., Oh, J., Wu, J., Lu, J., Yang, J., Malik, J., Silvério, J., Hejna, J., Booher, J., Tompson, J., Yang, J., Salvador, J., Lim, J. J.. Han, J., Wang. K., Rao, K., Pertsch, K., Hausman, K., Go, K., Gopalakrishnan, K., Goldberg, K., Byme, K., Oslund, K., Kawaharazuka, K., Black, K., Lin, K., Zhang, K., Ehsani, K., Lekkala. K., Ellis, K., Rana, K., Srinivasan, K., Fang, K., Singh, K. P., Zeng. K.-H., Hatch, K., Hsu, K., Itti, L., Chen, L. Y., Pinto, L.., Fei-Fei, L.,, Tan, L.,, Fan, L. J., Ott, L., Lee, L.., Weihs, L., Chen, M., Lepert, M., Memmel, M., Tomizuka, M., Itkina, M., Castro, M. G., Spero, M., Du, M., Ahn, M., Yip, M. C., Zhang, M., Ding, M., Heo, M.. Srirama, M. K., Sharma, M., Kim, M. J., Kanazawa, N., Hansen, N., Heess, N., Joshi, N. J., Suenderhauf, N., Liu, N., Palo, N. D., Shafiullah, N. M. M., Mees, O., Kroemer, O., Bas- tani, O., Sanketi, P. R., Miller, P. T., Yin, P., Wohlhart, P., Xu, P., Fagan, P. D., Mitrano, P., Sermanet, P., Abbeel, P., Sundaresan, P., Chen, Q., Vuong, Q., Rafailov, R., Tian, R. Doshi, R. Mart Mn-Matt in, R., Baijal, R., Scalise, R., Hendrix, R., Lin, R., Qian, R., Zhang, R., Mendonca, R., Shah, R., Hoque, R., Julian, R., Bustamante, S., Kir- mani, S., Levine, S., Lin, S., Moore, S., Bahl, S., Dass S., Sonawani, S., Tulsiani, S., Song, S., Xu, S., Haldar, S., Karamcheti, S., Adebola, S., Guist, S., Nasiriany, S., Schaal, S., Welker, S., Tian, S., Ramamoorthy, S., Dasari, S., Belkhale, S., Park, S., Nair, S., Mirchandani, S., Osa, T., Gupta, T., Harada, T., Matsushima, T. Xiao, T., Kollar, T., Yu, T., Ding, T., Davchev, T., Zhao, T. Z., Armstrong, T., Darrell, T., Chung, T., Jain, V., Kumar, V., Vanhoucke, V., Zhan, W., Zhou, W., Burgard, W., Chen, X., Chen, X., Wang, X., Zhu, X., Geng, X., Liu, X., Liangwei, X., Li, X., Pang. Y., Lu, Y., Ma. Y. J., Kim, Y., Chebotar, Y., Zhou, Y., Zhu, Y., Wu, Y. Xu, Y., Wang, Y. Bisk, Y., Dou, Y., Cho, Y., Lee, Y., Cui, Y., Cao. Y., Wu, Y.-H., Tang, Y., Zhu, Y., Zhang, Y., Jiang, Y., Li, Y., Li, Y., Iwasawa, Y., Matsuo, Y., Ma. Z., Xu,
Peebles, W. and Xie, S. Scalable diffusion models with trans- formers, 2023. URL https: /ttp//aaxvvv.../.. abs/ 2212.09748.
Radosavovic, I., Shi, B., Fu, L.,, Goldberg, K., Darrell, T. and Malik, J. Robot learning with sensorimotor pre- training. In Conference on Robot Learning, pp. 683-693. PMLR, 2023.
Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep gen- erative models. In International conference on machine learning, pp. 1278 1286. PMLR, 2014.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models, 2022. URL https: aaxxiv.rgg../////.... abs abs/2112.1075
Schmeckpeper, K., Xie, A., Rybkin, O., Tian, S., Dani- ilidis, K., Levine, S., and Finn, C. Learning predictive models from observation and interaction. In European Conference on Computer Vision, pp. 708-725. Springer, 2020.
Schmidhuber, Reinforcement learning upside down: Don't 't predict rewards just map them to actions. arXiv preprint arXiv: 1912:v:112.0227.. 2019.
Shafuullah, N. M. M., Cui, 7. J., Altanzaya, A., and Pinto, L. Behavior transformers: Cloning k modes with one stone, 2022. URL https: httpst//axxiv..rrr http://arx..gg//ea..aaa/ 2200..1255
Sohl-Dickstein, J., Weiss, E. A., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics, 2015. URL https: / l/arxiv.org/abee

---

Latent Diffusion Planning for Imitation Learning
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-hased generative modeling through stochastic differential equations. arXiv preprint arXiv:2011. 2020.
Zhang, S., Cao. Z.., Sadigh, D., and Sui, Y. Confidence- aware imitation learning from demonstrations with vary- ing optimality, 2022. URL https://arxiv.ort abs / 2110.14754.
Torabi, F., Warnell, G., and Stone, P. Behavioral cloning from observation. arXiv preprint arXiv:1805.0195,, 2018.
Zhao, T. Z., Kumar, V., Levine, S., and Finn, C. Learn ing fine-grained bimanual manipulation with low-cost hardware, 2023. URL https: /atps://arxiii www.i.oog.......a...oooo..ooooot /abs/
Watter, M., Springenberg, J., Boedecker, J., and Riedmiller, M. Embed to control: A locally linear latent dynamics model for control from raw images. Advances neural information processing systems, 28, 2015.
Zhao, T. Z., Tompson, J., Driess, D., Florence, P. Ghasemipour, S. K. S., Finn, C.,, and Wahid, A. ALOHA unlcashed: A simple recipe for robot dexterity. In 8th An- nual Conference on Robot Learning 2024. URL https: / //openreview. net wwtp:/eettfooooeeeeeeee www.rmiidiggdaa... ikHH.
Wen, C., Lin, X., So, J., Chen, K., Dou, Q. Gao, Y., and Abbeel, P. Any-point trajectory modeling for policy learn- ing, 2024. URL https: //axxii////xviv///// org/ abs 2401. 00025.
Zhou, S., Du, Y., Chen, J., Li, Y., Yeung, D.-Y., and Gan, C. Robodreamer: Learning compositional world models for robot imagination, 2024. URL https: // /arxivv. org/ 2237771777. 12377.
Wu, P., Escontrela, A., Hafner, D., Abbeel, P., and Goldberg, K. Daydreamer: World models for physical robot learn- ing. In Conference on robot learning, pp. 2226-2222222222 2240. PMLR, 2023a.
Wu, P., Majumdar, A., Stone, K., Lin, Y., Mordatch, I., Abbeel, P., and Rajeswaran, A. Masked trajectory mod- els for prediction, representation, and control. In Inter- national Conference on Machine Learning, pp. 37607- 37623. PMLR, 2023b.
Yan, W., Zhang, Y., Abbeel, P., and Srinivas, A. Videogpt: Video generation using vq-vae and transformers. arXiv preprint 1rXiv:2104.00157 2021.
Yang, M., Du, Y., Ghasemipour, K., Tompson, J., Schuur- mans, D., and Abbeel, P. Learning interactive real-world simulators. arXiv preprint arXiv:23 10.06/ 14, 2023.
Yang. S., Walker, J., Parker-Holder, J., Du, Y., Bruce, J., Barreto, A., Abbeel, P., and Schuurmans, D. Video as the new language for real-world decision making, 2024. URL https: / aarxiv.v.vv.////iv//...///aaaavv org/
Young, S., Gandhi, D., Tulsiani, S., Gupta, A., Abbeel, P., and Pinto, L. Visual imitation made casy. In Conference on Robot Learning, pp. 2995-2005. PMLR, 2021.
Yu, T., Kumar, A., Chebotar, Y., Hausman, K., Finn, C.. and Levine. S. How leverage unlabeled data in offline reinforcement learning. In Chaudhuri, K., Jegelka, S., Song, L.,, Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the 39th International Conference on Ma- chine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 2561 25611-25660 1-635. PMLR, 17-23 Jul 2022. URL https:tpps... /tt.tttppe//pooeeeaaoooo http///pceeenngsssecee mlr press/ v1 wwtyy222.......o..o..o..///ool html.
12

---

Latent Diffusion Planning for Imitation Learning
A. Appendix
A.I. Additional Ablations
A.1.1. LDP WITH PRETRAINED EMBEDDINGS
We evaluate the effectiveness of planning and extracting actions from pretrained embeddings with a DINOv2 ablation. In this approach, we replace our VAE embeddings with DINOv2 embeddings, with no other changes to the LDP architecture.
The DINOv2 latent embedding is 384-dimensional. We found directly planning over the DINOv2 embeddings does lead to good learned behaviors (0% success), which we hypothesize is due to the challenges of planning over large latent space. Thus, as an alternative, we fix random projection matrix, reducing the 384 dim feature space to 16 dim, matching LDP. We choose this, because this is straightforward way of using frozen embeddings. It may be possible to plan over the large embeddings space with a much larger and complex model, or it may be possible to learn alternative ways to project DINOv2 embeddings to a lwwer-dimensionl space, but this may lead to fundamental changes in our method, so we leave more complicated approaches to use pretrained embeddings for future work. We include results on Robomimic tasks in Table 5.
Table DINOv2 Ablation. Planning frozen embeddings (DINOv2), when projected to lower dimensional space, has significantly lower success rate than planning over VAE embeddings. This suggests that the learned latent space is crucial to LDP's performance, and that compact, well-structure latent space leads to optimal performance.
Method
Latent Dim
Lift
Can
Square
LDP w/ DINOv2 (frozen)
384
0,00 +: 0.00
0.00 0.00
0.00 0.00
LDP w/ DINOv2 (frozen, randomly projected)
16
0.44 +: 0.24
0.03 = 0.01
0,01 4 0.01
LDP + Action-Free + Subopt
16
1.00 + 0.00
0.98 = 0.00
0.83 # 0.01
A.1.2. UNIPI FROM PRETRAINED MODEL
In our experimental results, we train our UniPi model from scratch on demonstration videos. To provide comparison to training UniPi from pre-trained weights, use the AVDC-THOR pretrained checkpoint (Ko et al., 2023). We finetune for an additional 50k steps at learning rate le-4.
In Table 6 and Table 7, we find that the pretrained UniPi model does not substantially improve performance, likely because the pretrained model quality is limited. For UniPi-OL, there appears to be slightly improvement when pretraining from scratch, possibly because without closed-loop planning, the pretrained initialization may be helpful.
Table 6. UniPi from Pretrained Checkpoint
Method
Lift
Can
Square
UniPi-OL (from scratch)
0.09 0.03
0.27 : 0.01
0.07 + 0.01
UniPi-OL (from pretrain)
0.13 0.05
0.31 + 0.01
0.07 + 0.01
UniPi-CL (from scratch)
0.12 0.02
0.30 1: 0.02
0.10 1 0.04
UniPi-CL (from pretrain)
0.13 0.05
0.23 + 0.01
0.08 + 0.02
Table 7. UniPi Action-Free from Pretrained Checkpoint
Method
Lift
Can
Square
UniPi-OL (from scratch)
0.11 0.03
0.25 + 0.03
0.05 + 0.03
UniPi-OL. (from pretrain)
0.10 0.08
0.26 + 0.00
0.08 + 0.02
UniPi-CL (from scratch)
0.14 0.02
0.32
+ 0.04
0.09 + 0.01
UniPi-CL (from pretrain)
0.11 0.03
0.25 0.01
0.10 1: 0.00
A.2. Implementation Details
Diffusion Policy We use a Jax reimplementation of the convolutional Diffusion Policy, which we verify can reproduce reported Robomimic benchmark results.
13

---

Latent Diffusion Planning for Imitation Learning
UniPi We use the open-source implementation of UniPi (Ko al., 2023). For UniPi-OL and UniPi-CL, predict 7 future frames. During training time, for UniPi-OL, the future frames evenly sampled from training demonstration. For UniPi-CL, the 7 future frames are the next consecutive frames.
The goal-conditioned behavior cloning agent is implemented as goal-conditioned Diffusion Policy (Chi et al., 2023) agent, and it is trained on chunks of 16. The inverse dynamics model is based off of IDQL. (Hansen- Estruchh et al., 2023), and shares the same architecture as the IDM used LDP.
We train the video prediction models for 200k gradient steps with batch size 16.
LDP The LDP VAE is adapted from Diffusion Transformer (Peebles & Xie, 2023). The planner is based directly off of the convolutional U-Net from Diffusion Policy (Chi et al., 2023), with modifications to plan across latent embeddings instead of et al., 2023).
action chunks. The IDM is based off of IDQL
Table Diffusion Policy
Architecture Hyperparameters DP and LDP
UniPi-OL GCBC
LDP ALOHA Cube
down_dims
[256, 512, 1024]
[256, 512, 1024]
[512, 1024, 2048)
n diffusion steps
100
100
100
site
256
256
256
Ir
le-4
le-4
le-4
steps
500k
500k
500k
Table IDM Architecture Hyperparameters
UniPi-CL IDM
LDP IDM
blocks
3
3 (Lift, Square, ALOHA Cube); 5 (Can)
n diffusion steps
100
100
batch_size
256
256
Ir
le-4
le-4
sttpss
500k
500k
Table 10. VAE Architecture Hyperparameters
VAE
block.out.channels
[128, 256, 256, 256, 256, 256] [DownEneoderBlook22)) x6
block.lock... types
up_block_typer
[UpDecoderBBloc22]] x6
Aatent-chamne
4
Latent Dim
(2, 4)
Lift KL Beta
le-5
Can KI. Beta
le-6
Square KL Beta
le-6
ALOHA Cube KL. Beta
le-7
n n-grad_steps
300k
A.3. Simulation Experiments
I. ENVIRONMENT
For all environments, we use an observation horizon of 1. For Robomimic tasks, we use the third-persee view image. For ALOHA Cube and Franka Lift, we use the wrist camera. All images are 64x64.
14

---

Latent Diffusion Planning for Imitation Learning
A.3.2. DATASET SIZES
We detail the design decisions for the number of expert, action-free, and suboptimal trajectories. For simulated robotics tasks, kept consistent the action-free and suboptimal trajectories: 100 action-free trajectories for robomimic tasks and 500 suboptimal trajectories for all tasks.
Number of Demonstrations Robomimic includes 200 proficient demonstrations per task, and ALOHA includes 50 demonstrations per task. To showcase the effectiveness of our method, we always take half the number of demonstrations- 100 for Robomimic tasks, and 25 for ALOHA tasks. This is a consistent protocol we use to select the number of demonstrations. However, Lift a particularly simple task, so in order to not saturate our baseline methods with excessive numbers of demos, reduce the number of demonstrations to 3. Without this adjustment, Lift results would not be useful discern the effect of different types of data, and furthermore, Lift results would all be saturated.
Number of Action-Free Trajectories We chose to use the remaining half of the demonstrations as action-free trajectories, Hence, ALOHA uses 25 action-free free trajectories, and Robomimic uses 100 action- oetion-free trajectories. Note again that Lift only used 3 expert demonstrations, but in order to keep our numbers as consistent as possible and avoid further confusion, we opted to use 100 action-free trajectories.
A.3.3. METHOD DETAILS
LDP VAE We train with additional action-free and suboptimal data. This VAE is shared across single task for all variants of LDP, and including for DP + Repr.
LDP + Subopt We train the inverse dynamics model with batches of 50% optimal, 50% suboptimal data.
LDP + Action-Free We train only the planner on additional action-free demos. The IDM is trained only on expert
demonstrations.
LDP Action-Free + Subopt We train only the planner on additional action-free demos. We train the inverse dynamies model with batches of 50% optimal, 50% suboptimal data.
LDP Hierarchical The IDM is parameterized as Conditional U-Net, like the planner, but the IDM is smaller and has down.dims=(256, 512].
LDP DINOv2 In this variant, we plan over pretrained DINOv2 (Oquab et al., 2023; Darcet et al., 2023) embeddings instead of VAE embeddings. We use the ViT-S variant and pad the 64x64 image to be 70x70, which matches the 14x14 patch size. The batch size is reduced to 128. but the remaining hyperparameters are identical LDP. We further project the 384-dimensional embedding 16-dimensional via random matrix, in order for easier planning. Results are presented in Appendix A.1.I
DP Repr We use the same VAE that LDP methods used.
DP PT + FT We pretrain for 300k steps and finetune with a lower learning rate (1e-5) for 200k steps.
DP-VPT We use an inverse dynamics model trained for 500k steps on curated data, as used in LDP, to relabel actions for the action-free data.
UniPi We train two goal-conditioned agents and two inverse dynamics models and report the average success. For UniPi-OL evaluations, we predetermine the number of steps for the GCBC to reach each image subgoal based on the demonstration lengths. For Lift, evaluation episode lengths are 60 steps; Can is 140 steps; and Square is 160 steps. This maximum horizon is also enforced for UniPi-CL evaluations, for consistency.
15


Reviews:
[{"review_id": "6bmhaMQ64b", "reviewer": "Reviewer_9mX6", "summary": "This paper presents Latent Diffusion Planning (LDP), an algorithm aimed at performing imitation learning with the presence of additional suboptimal and action-free demonstrations.\n\n----\n\nProblem Setting and Key Assumptions:\n- Vision-based imitation learning for table top manipulation\n- Aside from expert demos, assume access to (expert) action-free demos and suboptimal/failed data\n\n----\n\nThe main algorithm consists of three components:\n- A VAE to learn a low-dimensional embedding space for the visual observations\n- A diffusion-based planner that performs forward prediction in the frozen VAE embedding space\n- A diffusion-based inverse dynamics model that takes adjacent VAE latents to predict robot actions.\nThis design allows the method to learn from the three types of data as mentioned earlier. Specifically, the VAE is trained to enc/dec individual frames, and thus, can take all the data. The planner only need the optimal sequence of latent states. Hence, both expert and action-free demos can be used. Finally, the inverse dynamics can be trained on all $(o, a, o')$ tuples, regardless of whether they lead to task success or not.\n\n----\n\nThe authors evaluate LDP in 4 simulated tasks and a real robot task on tabletop manipulation. Because the assumption on the available data is new, they compare with Diffusion Policy and a few variants. They also compare with UniPi. Results support that in the low-data setting, LDP outperforms these relevant baselines by also utilizing the action-free demo and suboptimal data.\n\n----\n\n## Update after Rebuttal\n\nI reviewed other reviewers' comments and the authors' responses. I have raised my score from 2 to 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you for your detailed feedback on our project. We are happy that you found the paper easy to follow and relevant to learning from demonstrations. Please let us know whether there are any other concerns you have that prevent you from increasing your score.\n\n**Q1: Pretrained UniPi + Finetuning In-Domain**\n\nThe comparison with UniPi is limited as the original UniPi model, code and pretraining data are not released, per Appendix A.1. To provide a comparison, we use AVDC-THOR, which is the best available reproduction of UniPi. We include results on Robomimic tasks. The experiments compare how finetuning a pretrained UniPi model, instead of training a UniPi model from scratch, affects overall performance. We finetune for an additional 50k steps at learning rate 1e-4. \n\n| Method                   | Lift         | Can          | Square       |\n|--------------------------|--------------|--------------|--------------|\n| UniPi-OL (from scratch)  | 0.09 +- 0.03 | 0.27 +- 0.01 | 0.07 +- 0.01 |\n| UniPi-OL (from pretrain) | 0.13 +- 0.05 | 0.31 +- 0.01 | 0.07 +- 0.01 |\n| UniPi-CL (from scratch)  | 0.12 +- 0.02 | 0.30 +- 0.02 | 0.10 +- 0.04 |\n| UniPi-CL (from pretrain) | 0.13 +- 0.05 | 0.23 +- 0.01 | 0.08 +- 0.02 |\n\n| Method                                 | Lift         | Can          | Square       |\n|----------------------------------------|--------------|--------------|--------------|\n| UniPi-OL + Action-Free (from scratch)  | 0.11 +- 0.03 | 0.25 +- 0.03 | 0.05 +- 0.03 |\n| UniPi-OL + Action-Free (from pretrain) | 0.10 +- 0.08 | 0.26 +- 0.00 | 0.08 +- 0.02 |\n| UniPi-CL + Action-Free (from scratch)  | 0.14 +- 0.02 | 0.32 +- 0.04 | 0.09 +- 0.01 |\n| UniPi-CL + Action-Free (from pretrain) | 0.11 +- 0.03 | 0.25 +- 0.01 | 0.10 +- 0.00 |\n\nWe observe the pretrained AVDC/UniPi model does not substantially improve performance, likely because the pretrained model quality is limited. For UniPi-OL, there appears to be slightly improvement when pretraining from scratch, possibly because without closed-loop planning, the pretrained initialization may be helpful.\n\n**Q2: Action-Free Demonstrations**\n\nThank you for this suggestion. We were unable to set up the UMI hardware within the short rebuttal time, but we agree that a more scalable and realistic way of obtaining action-free expert data is through data collection tools such as UMI (Universal Manipulation Interface). \n\n**Q3: Collecting Demonstrations in the Real World**\n\nThe demonstrations were collected from one teleoperator. The average time for one demonstration was around 25 seconds, including resetting the cube. In total, for 82 trajectories, it would take around 34 minutes to collect all demos. \n\nFor suboptimal trajectories, the teleoperator merely supervises the policy to ensure safety, and does not need to actually teleoperate the robot at all. Each suboptimal trajectory takes around 75 seconds, including resets. For 85 trajectories, this would take around 1 hour 45 minutes. \n\nThe total time is longer than LANE from Zhao et. al., which collects limited demonstrations (requires teleoperator) and performs RL (does not require active teleoperation). It’s possible that our method may work with fewer demonstrations (less teleoperation time) or with less suboptimal data, but we did not sweep over these numbers or optimize for these metrics.\n\n**Q4: Additional Comments**\n\n“The tasks are rather simple with clean visuals.” -- We agree that the tasks are not extremely complex nor do they have complex backgrounds or visuals. However, we chose commonly used robotics benchmarks for this project and focused on showing the effect of action-free and suboptimal data. \n\n“The main weakness is that the algorithm still seems to require many demonstrations (and more suboptimal data) for rather simple tasks.” -- We agree that this paper doesn’t operate in a few-shot demonstration setting. However, the standard number of demonstrations provided by robomimic is 200 per task, and ALOHA is 50 per task. The robomimic benchmark is saturated with 200 demos, and we chose to restrict the number of demonstrations to evaluate our method in a lower demonstration regime. In order to learn from even fewer demonstrations for these tasks, it would likely be necessary to learn from prior data (e.g. from retrieval), or reinforcement learning, which we do not focus on.\n\n**Q5: Improved Experiments**\n\nPlease **see Reviewer LWMg Q2**. We include updated experiments where LDP leverages both action-free and suboptimal data. While BC-like methods work well in settings with a lot of optimal demonstration, LDP effectively leverages additional data sources, which we believe is a fundamental advantage. The results strongly outperform baselines now, by around 30% or more. \n\nIn addition, we have reran UniPi with identical hyperparameters to LDP, per Reviewer PyaC’s suggestion.", "claims_and_evidence": "The main claim is that LDP should be data efficient (for expert data) by leveraging sub-optimal data and action-free data. This is indeed verified with experiments. The main weakness is that the algorithm still seems to require many demonstrations (and more suboptimal data) for rather simple tasks.", "methods_and_evaluation_criteria": "Overall, the method and evaluation procedures are sound. However, the experiment results are not convincing enough because:\n- The comparison with a foundation model like UniPi feels somewhat out of place because I assume the authors train the model from scratch using only in-domain data. Fine-tuning a pre-trained model and allowing joint training on multiple datasets are more appropriate.\n- The tasks are rather simple with clean visuals.", "theoretical_claims": "No theoretical claims.", "experimental_designs_or_analyses": "The experiment procedure and selected baseline methods on imitation learning are fair. Regarding the real robot experiment, the authors provide action-free demonstration by removing actions from actual demonstrations. This is a rather awkward design choice. Instead, they should consider collecting demonstrations using a setup similar to the Universal Manipulation Interface paper. In that case, I wonder if the planner model can still reliably predict because the data distribution of action-free demo will be different from the expert demos.", "relation_to_broader_scientific_literature": "This work aims to improve the data efficiency of IL methods by utilizing data otherwise not useful. It is highly relevant to the field of LfD (Learning from Demonstrations).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The paper is easy to follow and includes clear visualizations of the model architecture, and comprehensive discussion of relative work. The experiments sections presents the research questions clearly followed by the corresponding studies.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "For the real robot experiment, how much time does it take the collect the demonstrations? Recent works have shown that similar tasks can be learned from 1 hour of data, including a few demonstrations and sparse-reward RL. (Accelerating Visual Sparse-Reward Learning with Latent Nearest-Demonstration-Guided Explorations​)", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ilEbW9qSyB", "reviewer": "Reviewer_PyaC", "summary": "The work proposes a novel approach for imitation learning that combines an inverse dynamics model (IDM) with a planner that proposes future goal states in latent space. The approach first trains a variational autoencoder (VAE) that encodes visual representations of states into a lower dimensional latent space. Using the embeddings from this encoder, a IDM model is trained that given a state representation and future state representation predicts the action that will lead to that future state. At the same time, a planner is trained that predicts a sequence of embeddings of future states given an embedding of a state. The main insight of this work is that the IDM model can be trained with suboptimal/ general data that contains states and actions within the given environment, and the planner can be trained with (ideally optimal) demonstrations of solving the desired task but no actions are required for training the planner. Both the IDM and planner models are implemented using diffusion models. The approach is evaluated in a series of simulated robotics tasks and a real world robotics experiment. Compared to a series of imitation learning and planning baselines and ablations, the proposed approach is found to be more effective at leveraging data that is partly suboptimal and/ or action-free.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you for your detailed feedback on our project. We are happy that you found the paper well-presented with clear experimental evaluations. Please let us know whether there are any other concerns you have that prevent you from increasing your score.\n\n**Q1: Would the authors be able to explain why these vastly different amounts of trajectories were chosen?**\n\nFor different robotics tasks, we tried to keep most of the parameters the same: 100 action-free trajectories for robomimic tasks and 500 suboptimal trajectories for all tasks.\n\n\\# of demonstrations: Robomimic includes 200 proficient demonstrations per task, and ALOHA includes 50 demonstrations per task. To showcase the effectiveness of our method, we always take **half** the number of demonstrations-- 100 for Robomimic tasks, and 25 for ALOHA tasks. This is a consistent protocol we use to select the number of demonstrations.\n\n\\# of action-free trajectories: we chose to use the remaining **half** of the demonstrations as action-free trajectories\nLift is a very simple task, so we reduce the number of demonstrations to 3 increase the complexity.\n\n**Q2: Ablating the VAE Component**\n\nTo address “directly planning in high-dimensional image-space,” we include the UniPi baseline, which learns a video prediction planner. UniPi-CL, in particular, mirrors LDP’s planner, in that UniPi-CL forecasts over dense states instead of subgoals. Let us know if that addresses your concern!\n\n**Q3.1 DP vs. LDP on ALOHA Cube**\n\nSpecifically for ALOHA Cube, we found that a larger planner was imperative for reasonable performance. We did not find this necessary for DP, and thus, we kept the smaller DP architecture. We have updated results where DP is trained with the same architecture size as all LDP variants (down_dims = [512, 1024, 2048]). We use batch_size=16 for DP, since the end-to-end training of the encoder requires much more GPU memory than LDP, which uses frozen embeddings. We couldn't run DP with a larger batch size or down_dims.\n\n| DP           | LDP          | LDP + Action-Free | LDP + Subopt | LDP + Action-Free + Subopt |\n|--------------|--------------|-------------------|--------------|----------------------------|\n| 0.32 +- 0.00 | 0.64 +- 0.04 | 0.70 +- 0.02      | 0.71 +- 0.03 | 0.97 +- 0.01               |\n\n\n**Q3.2 LDP vs. UniPI: IDM size for Can, and IDM train time**\n\nWe observed a bigger IDM is helpful for the Can task, likely because the scene is slightly more visually complex. As requested, we have now rerun Can experiments with UniPi to ensure the same IDM architecture.\n\nWe train the IDM model for longer in LDP because we use a smaller batch size. However, we agree that for consistency, rerun UniPi experiments with the same batch size and \\# gradient steps for LDP and UniPi IDM. In addition, we have retrained UniPi GCBC models with the same hyperparameters as LDP. **See Reviewer LWMg Q2.**\n\n**Q3.3 LDP 50% optimal and 50% suboptimal batches**\n\nEach batch consists of 50% optimal and 50% suboptimal trajectories.\n\n**Q3.4 IDM one expert demonstration**\n\nWe meant to write, “The IDM is trained only **ON** expert demonstrations.” The IDM is only trained on action-labelled data, which in this case, are from the expert demonstrations.\n\n**Q3.5 LDP Hierarchical IDM**\n\n(Appendix A.2) Hierarchical LDP’s IDM is a Conditional U-Net with down-dims [256, 512], which has 1.67e7 parameters. (Appendix A.1) Non-Hierarchical LDP uses an MLP ResNet based off of IDQL, which has 1.79e6 parameters. The Hierarchical LDP IDM has more parameters, while still underperforming our method.\n\n**Q4: Pretrained Encoders**\n\nLDP uses a task-specific VAE, because existing image encoders or VAEs (e.g. DINOv2, R3M, Stable Diffusion VAE) produce high-dimensional embeddings. Our VAE produces a much more compact latent space (16 dimension), enabling much faster training and inference, as well as interpretability through decoding image latents. \n\nTo address this approach, we swap out VAE embeddings with frozen DINOv2 embeddings. \n\nHowever, we found directly planning over the DINOv2 embeddings does not lead to good learned behaviors (0% success), which we hypothesize is due to the challenges of planning over a large latent space (384 dim). Thus, as an alternative, we fix a random projection matrix, reducing the 384 dim feature space to 16 dim, matching LDP. We choose this, because this is a straightforward way of using frozen embeddings. It may be possible to plan over the large embeddings space with a much larger and complex model, or it may be possible to learn alternative ways to project DINOv2 embeddings to a lower-dimensional space, but this may lead to fundamental changes in our method, so we don’t explore more complicated approaches to use pretrained embeddings. \n\n| Method | Lift         | Can          | Square       |\n|--------|--------------|--------------|--------------|\n| DINOv2 | 0.44 +- 0.24 | 0.03 +- 0.01 | 0.01 +- 0.01 |", "claims_and_evidence": "The claims made in this work are largely clearly presented and well supported through empirical evidence. However, there are several inconsistencies in the experimental evaluation that I would expect to be corrected or well justified since these might invalidate some of the findings of the provided experiments (see Experimental Design and Analyses, esp. 1. and 3.)", "methods_and_evaluation_criteria": "The work strongly relies on the latent space learned by the VAE to be expressive and representative of features that would be important for learning the dynamics and decision making within the task. Given the efficacy of pre-trained visual encoders for imitation learning (e.g. [1, 2, 3, 4]), I wonder whether this step of training a task-specific VAE is even necessary or if you could replace this with an off-the-shelf pre-trained encoder (and if necessary fine-tune on available data). Establishing that this work could work well with pre-trained visual encoders would further generalise the applicability of this approach and reduce training cost.\n\n[1] Nair, Suraj, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. \"R3m: A universal visual representation for robot manipulation.\" arXiv preprint arXiv:2203.12601 (2022).\n\n[2] Schäfer, Lukas, Logan Jones, Anssi Kanervisto, Yuhan Cao, Tabish Rashid, Raluca Georgescu, Dave Bignell, Siddhartha Sen, Andrea Treviño Gavito, and Sam Devlin. \"Visual encoders for data-efficient imitation learning in modern video games.\" arXiv preprint arXiv:2312.02312 (2023).\n\n[3] Shang, Jinghuan, Karl Schmeckpeper, Brandon B. May, Maria Vittoria Minniti, Tarik Kelestemur, David Watkins, and Laura Herlant. \"Theia: Distilling diverse vision foundation models for robot learning.\" arXiv preprint arXiv:2407.20179 (2024).\n\n[4] Yuan, Zhecheng, Zhengrong Xue, Bo Yuan, Xueqian Wang, Yi Wu, Yang Gao, and Huazhe Xu. \"Pre-trained image encoder for generalizable visual reinforcement learning.\" Advances in Neural Information Processing Systems 35 (2022): 13022-13037.", "theoretical_claims": "There are no theoretical contributions or proofs to check.", "experimental_designs_or_analyses": "The work uses vastly different amounts of demonstrations across all tasks in the evaluation\n- Can, Square: 100 demonstrations, 500 suboptimal trajectories, 100 action-free trajectories\n- Lift: 3 demonstrations, 500 suboptimal trajectories, 100 action-free trajectories\n- Transfer Cube: 25 demonstrations, 500 suboptimal trajectories, 25 action-free trajectories\nand the amount of suboptimal trajectories in particular is significantly larger than the main demonstrations provided in the Lift and Transfer Cube tasks. In the ALOHA Transfer Cube task, it is clear that leveraging the 500 suboptimal trajectories is necessary for good performance from the results in Table 2. In particular typical DP is unable to perform well in this task since it is only trained on 25 demonstrations but using suboptimal trajectories via reward conditioning (RC-DP) or pre-training (DP PT + FT) significantly improves upon DP and in particular in the latter case performs comparable to LDP.\n\n1. Would the authors be able to explain why these vastly different amounts of trajectories were chosen?\n2. The work provides several ablations on varying data being used across its experiments and compares to a large set of sensible baselines, but it does not ablate the VAE component in its experiments. I would expect that directly planning in high-dimensional image-space performs worse than the proposed approach in latent space, but the work provides no direct evidence for this claim.\n3. From the supplementary material, I see several differences between LDP and baseline algorithms or ablations. Would the authors be able to explain the following differences and provide fair comparisons in these tasks?\n\t1. According to Table 5, LDP trains a larger network for the ALOHA cube task compared to other baselines (esp. DP and DP-based algorithms). Why do you use larger networks for LDP and could you provide comparisons to DP at the same size of policy network to ensure that LDP does not outperform the baselines in this task due to its larger networks?\n\t2. According to Table 6, LDP trains a larger IDM model than the UniPi baselines in the Can task (5 vs 3 blocks) and trains its IDM model for longer in all tasks compared to UniPi (500k vs 200k gradient steps). Again, what is the reason for these discrepancies? For fair comparisons, I would expect models to be trained for similar amounts of steps and models to be of identical size where possible.\n\t3. According to Appendix section A.2, LDP + Subopt is trained with 50% optimal and 50% suboptimal data. Would the authors be able to clarify what they mean? Is each update batch constructed of optimal and suboptimal data in equal proportions, or do you subsample the dataset of optimal and suboptimal demonstrations to have an equal share of both?\n\t4. According to Appendix section A.2, LDP + Action-Free trains the main IDM model only on a single expert demonstration. Why do you not train the IDM model still on all available demonstrations as done for other algorithms?\n\t5. According to Appendix section A.2, LDP Hierarchical uses a smaller IDM model compared to LDP. This renders the comparison in Table 4 as unfair and not convincing anymore to state that dense forecasting is an important contribution and part of LDP.", "relation_to_broader_scientific_literature": "The work does a good job at relating to alternative diffusion-based and IDM-based imitation learning algorithms. In particular, it refers to important literature that leverages action-free or suboptimal data but clarifies the distinctions to LDP.", "essential_references_not_discussed": "I am not aware of any essential references that require further discussion.", "other_strengths_and_weaknesses": "I would like to commend the authors on a overall well constructed and clearly motivated evaluation. In particular, all the baselines serve specific purposes and allow to identify the importance of different components of LDP.", "other_comments_or_suggestions": "No other comments or suggestions.", "questions_for_authors": "1. The supplementary material (A.2, Table 5 and 6) reveal several inconsistencies across LDP and its baselines or ablations in several experiments. I would expect the author to correct these inconsistencies and/ or provide convincing justifications for them. Otherwise, it is unclear whether the empirical findings are due to the proposed algorithms or differences in hyperparameters. (see 3. in Experimental Design and Analyses for more details)\n   I otherwise consider this a strong paper and will increase my score if the authors are able to provide convincing justification or corrections for these inconsistencies.\n2. Would the authors be able to explain why the varying amounts of trajectories used to train LDP and baseline algorithms across tasks? (see Experimental Design and Analyses for more details)\n3. Would the authors be able to provide an ablation that does LDP with its planner and IDM model directly operating on images rather than the VAE latent space? I would expect such an ablation to perform worse but the work currently provides no evidence for such claims.\n\n**The score has been updated in response to the author rebuttal**", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b0YIjcdlCi", "reviewer": "Reviewer_zcGe", "summary": "- This paper ultimately aims to do some form of imitation learning in robotic settings\n- It does this with a modular approach, using: 1) a 'planner' to predict sequences of observations from those provided by an expert demonstrator. 2) an IDM predicting actions from past and future observations. \n- At inference, the planner is used to generate a sequence, which the IDM converts to actions, which can be executed.\n- Both models operate on observations that are compressed by a beta-VAE. Both are trained with diffusion as the loss.\n- The paper emphasizes the types of data that can be used to train each model. The planner can also be trained on non-labelled (but still expert) data, the IDM can be trained sub-optimal (but still labelled) data.\n- Experiments in several simulated robotics tasks and one real robotic task show slight improvements over baselines.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you for your detailed feedback on our project. We provide the requested experiments and address the comments and questions in detail below. Please let us know whether there are any other concerns you have that prevent you from increasing your score.\n\n**Q1: Improvement in Experiments**\n\nWe agree the experimental results in our submitted draft show that there is around a 10% improvement for LDP. However, we have now run updated experiments where LDP leverages both action-free and suboptimal data. While BC-like methods work well in settings with a lot of optimal demonstration, LDP effectively leverages additional data sources, which we believe is a fundamental advantage. The results strongly outperform baselines now, by around 30% or more. **See Reviewer LWMgQ2**\n\n**Q2: Novelty**\n\nThe novelty of this paper is to propose a simple and scalable algorithm to address learning from heterogeneous data sources like suboptimal and action-free data. Our main contribution is in showing good performance with a model-based method, which will enable future work to build on it. Further, we provide a comprehensive evaluation of alternative methods for learning from heterogeneous data, and establish a novel finding that our planning-based is particularly suited for this setting.\n\nIn terms of comparison to UniPi, we agree that the method is fairly similar. However, we find that LDP strongly outperforms UniPi due to its method, with the main difference being the much lower-dimensional latent space. Due to this change in method, there is a large improvement in performance. We attach updated results below and highlighted in red in our PDF.\n\nIn terms of how to shape the space, we add additional experiments using pretrained DINOv2 embeddings, which investigates how pretrained embeddings compare to LDP’s latent space. In these experiments, we directly swap out VAE embeddings with frozen DINOv2 embeddings. \n\nHowever, we found directly planning over the DINOv2 embeddings does not lead to good learned behaviors (0% success rate even for easy tasks), which we hypothesize is due to the challenges of planning over a large latent space (384 dimensional). Thus, as an alternative, we fix a random projection matrix that reduces the 384 dimensional feature space to 16 dimensions, matching LDP’s latent space. We find that LDP strongly outperforms using pretrained latent embeddings. **See Reviewer PyaC Q4**\n\n\n**Q3: Data Size in Experiments**\n\nWe agree that using 1,000 or 10,000 trajectories would be interesting. However, most robotics benchmarks and tasks typically use 50-300 demonstrations (Diffusion Policy, Robomimic, Action Chunking with Transformers, etc.). For example. Robomimic, a popular simulated imitation learning benchmark, includes datasets with tasks that have 200 demonstrations. Simulated ALOHA tasks use only 50 demonstrations. It is typically difficult to find single-task imitation learning datasets with more than a few hundred trajectories, due to the difficulty of collecting large-scale expert demonstrations.\n\nLarger datasets are often multi-task or language-conditioned, and they are often used to train large multi-task or language-conditioned policies. (Open X-Embodiment, DROID, Octo, OpenVLA). These datasets are expensive to train on, especially for video models. We do not explore these datasets, due to the lack of good pretrained models, and because video models are much more expensive to train. This is why we are starting with small datasets, but this is an important direction for future work. \n\n\n**Q4: Pretrained Video Models**\n\nWe agree that LDP cannot use pretrained video models. However, for our UniPi baseline, we compare finetuning a pretrained model vs. training from scratch. To provide a comparison, we use a model with weights pretrained on THOR, and from scratch. In these results, we find that initializing from a pretrained model does not actually improve performance. **See Reviewer 9mX6 Q1**\n\nThe power of pretrained video models may lie in its ability to generalize and extrapolate to new tasks and scenes, often through language conditioning, which we leave to future work. Generalization is important for robot learning, but we find that for single-task imitation learning, which typically use smaller datasets of 50-300 demonstrations, using pretrained video models is not essential. \n\n**Q5: Outdated Beta-VAE and DDPM**\n\nWe agree that there are advances to beta-VAE (VQ-VAE, VQ-GAN, etc.) and DDPM (DDIM, Consistency Models, etc.). However, we choose both of these implementations due to their simplicity. For our planner, fast inference is not crucial to our contribution, and hence, we don’t use faster samplers like DDIM. However, we agree that improvements to both beta-VAE and DDPM can lead to additional improvements and scalability, which we leave to future work.\n\n**Q6: Citations**\n\nThank you for those references. We will include them in the draft.", "claims_and_evidence": "See strengths/weakenesses.", "methods_and_evaluation_criteria": "Fine.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "- Number of trajectories was quite small in all cases (100's of trajectories), which limits the impact of the work.\n- Error bars overlapping for real results in Table 3", "relation_to_broader_scientific_literature": "See strengths/weakenesses.", "essential_references_not_discussed": "Diffusion policy is one example of diffusion for imitation learning, but there are other relevant papers missed.\n\n- Diffusion policies as an expressive policy class for offline reinforcement learning\n- Imitating Human Behaviour with Diffusion Models\n- ...\n\nAlso,\n- DP-VPT is presented as a baseline, but the Tobari paper might be cited here (it is elsewhere) and used in the name, since it came well before the VPT paper.", "other_strengths_and_weaknesses": "Strengths\n- Real robotic experiments.\n- Main idea is sensible.\n- Minor improvements in experiments.\n- Nice justification about different types of data for different models.\n\nWeaknesses\n- My main criticism of the paper is wrt the novelty. The method feels like a repeat of unipi (and probably other works) with the modification that things are done in a vae's latent space. This just doesn't feel innovative in itself to justify acceptance. Given this is a key difference, I'd expect to see a deep investigation of how to shape this space, and analysis of speedups etc. But this is lacking.\n- Improvement in experiments is minor compared to baselines (around 10%). Given the variability in results that can be caused by implementation details of baselines, I'm hesitant to believe that latent diffusion planning is really delivering some substantial gain here.\n- The data size is quite small for all experiments. It'd be better to see if these techniques hold in larger (1000's or 10,000's trajectories)\n- While doing things in a custom latent space might be better from a speed perspective, it means pre-trained video generation models cannot be used, which I would expect might offer a good initialization for learning from the small datasets used in this work.\n- In general the implementations for various components are a little outdated -- beta-vae and ddpm.", "other_comments_or_suggestions": "NA", "questions_for_authors": "See weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "o9TW9vRQVa", "reviewer": "Reviewer_LWMg", "summary": "The paper proposes latent diffusion planning (LDP), a method for imitation learning featuring 3 components: 1) A variational autoencoder, mapping images to a latent spaces 2) A latent diffusion planner, which generates a sequence of latent states that the policy should visit 3) An inverse dynamics model, also leveraging diffusion, which associates an action to a latent transition. The interesting property of this framework is that it can use suboptimal demonstrations (with actions) to refine the inverse dynamics model, and unlabeled expert videos to improve the latent planner. Experiments show that this setup better makes use of unlabeled or suboptimal demonstrations than previous methods (on average).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you for your detailed feedback on our project. We are happy that you found the paper to be well-written and clear. We address your questions and comments below. Please let us know whether there are any other concerns you have that prevent you from increasing your score.\n\n**Q1: Choice of labeled / unlabeled / suboptimal trajectories are needed for optimal performance**\n\nFor different robotics tasks, we tried to keep most of the parameters the same: 100 action-free trajectories for robomimic tasks and 500 suboptimal trajectories for all tasks. However, we chose:\n\n\\# of demonstrations: Robomimic, by default, includes 200 proficient demonstrations per task, and ALOHA includes 50 demonstrations per task. To showcase the effectiveness of our method when the number of demonstrations is low, we chose to take **half** the number of demonstrations-- 100 for Robomimic tasks, and 25 for ALOHA tasks. However, Lift is a very simple task, and even with around 10 demonstrations, behavior cloning policies can achieve very high accuracy. Thus, specifically due to the simplicity of the Lift task, we chose 3 demonstrations to adhere to the low demonstration setting.\n\n\\# of action-free trajectories: we chose to use the remaining **half** of the demonstrations as action-free trajectories. This meant 100 demonstrations for Robomimic tasks, and 25 demonstrations for ALOHA. Since Lift only used 3 demonstrations, we considered using only 3 action-free trajectories. However, because we wanted to try and remain consistent with the number of trajectories, we chose to use 100 action-free demonstrations.\n\n**Q2: Does using action-free and suboptimal trajectories at the same time further improve the results?**\nYes, using action-free and suboptimal trajectories further improves results and strongly exceeds our baselines. We have included updated results. Here is a how LDP + Action-Free + Suboptimal compares against just LDP without the additional trajectories:\n\nIn addition, we have reran UniPi with identical hyperparameters to LDP, per Reviewer PyaC’s suggestion.\n\n| Method                     | Lift         | Can          | Square       | ALOHA Cube   | Average |\n|----------------------------|--------------|--------------|--------------|--------------|---------|\n| DP                         | 0.60 +- 0.00 | 0.63 +- 0.01 | 0.48 +- 0.00 | 0.32 +- 0.00 | 0.51    |\n| DP-VPT                     | 0.69 +- 0.01 | 0.75 +- 0.01 | 0.48 +- 0.04 | 0.45 +- 0.03 | 0.59    |\n| UniPi-OL + Action-Free     | 0.09 +- 0.05 | 0.23 +- 0.03 | 0.07 +- 0.03 | 0.02 +- 0.00 | 0.11    |\n| UniPi-CL + Action-Free     | 0.14 +- 0.02 | 0.32 +- 0.04 | 0.09 +- 0.01 | 0.17 +- 0.03 | 0.18    |\n| LDP                        | 0.69 +- 0.03 | 0.70 +- 0.02 | 0.46 +- 0.00 | 0.64 +- 0.04 | 0.65    |\n| LDP + Action-Free          | 0.67 +- 0.01 | 0.78 +- 0.04 | 0.47 +- 0.03 | 0.70 +- 0.02 | 0.66    |\n| LDP + Action-Free + Subopt | 1.00 +- 0.00 | 0.98 +- 0.00 | 0.83 +- 0.01 | 0.97 +- 0.01 | 0.95    |\n\n| Method                     | Lift         | Can          | Square       | ALOHA Cube   | Average |\n|----------------------------|--------------|--------------|--------------|--------------|---------|\n| DP                         | 0.60 +- 0.00 | 0.63 +- 0.01 | 0.48 +- 0.00 | 0.32 +- 0.00 | 0.51    |\n| RC-DP                      | 0.40 +- 0.04 | 0.73 +- 0.03 | 0.66 +- 0.02 | 0.60 +- 0.04 | 0.60    |\n| DP+Repr                    | 0.66 +- 0.04 | 0.61 +- 0.01 | 0.44 +- 0.02 | 0.25 +- 0.03 | 0.49    |\n| DP PT + FT                 | 0.52 +- 0.02 | 0.67 +- 0.01 | 0.57 +- 0.03 | 0.78 +- 0.00 | 0.64    |\n| UniPi-OL                   | 0.12 +- 0.06 | 0.28 +- 0.02 | 0.07 +- 0.01 | 0.00 +- 0.00 | 0.12    |\n| UniPi-CL                   | 0.12 +- 0.02 | 0.30 +- 0.02 | 0.10 +- 0.04 | 0.15 +- 0.07 | 0.17    |\n| LDP                        | 0.69 +- 0.03 | 0.70 +- 0.02 | 0.46 +- 0.00 | 0.64 +- 0.04 | 0.65    |\n| LDP + Subopt               | 0.84 +- 0.06 | 0.68 +- 0.02 | 0.55 +- 0.03 | 0.71 +- 0.03 | 0.70    |\n| LDP + Action-Free + Subopt | 1.00 +- 0.00 | 0.98 +- 0.00 | 0.83 +- 0.01 | 0.97 +- 0.01 | 0.95    |\n\n**Q3: Training the inverse dynamics on suboptimal trajectories...**\n\nWe believe this does help, because LDP + Subopt > LDP, and LDP + Action-Free + Subopt > LDP + Action-Free. In both of these cases, the difference is the suboptimal data for the IDM, which improves policy rollouts. One metric we can use to quantitatively measure this is action MSE for a validation dataset, but we found that this is noisy and not extremely correlated with the policy performance, which is true for many behavior cloning methods. Thus, we don’t find action MSE metrics for the IDM to be particularly insightful, and instead refer to policy performance.\n\n**Q4: Are the action-free data always from expert policies? To my understanding, suboptimal action-free trajectories cannot be used because they would decrease the performance of the planner.**\n\nYes, action-free data assumes that optimal behavior.", "claims_and_evidence": "The main claims of the papers are 1) LDP outperforms previous methods thanks to its better usage of unlabeled and suboptimal data 2) LDP can be applied to real robots, where collecting data with actions can be costly.\n\nThe results listed in Table 1 and 2 are convincing. LDP is competitive with the best baselines in all the considered tasks, often outperforming them. The authors make the effort of fairly comparing LDP with other methods by providing them with the same data LDP has access to. For example, relabeling action-free data with an inverse dynamics model is a strong alternative to LDP (which in fact performs comparably to LDP in Lift and Square), but overall LDP is stronger.\n\nI would suggest removing LDP + Subopt from the first table, or to merge the 2, because I found it confusing at a first read (to my understanding, the baselines in table 1 do not use suboptimal data).\n\nExperiments using Franka also show an improvement over DP, especially using action-free data.\n\nOne possible issue in using LDP is that often the performance improvement is unpredictable - in some cases using action-free or suboptimal trajectories leads to a large improvement, sometimes even to a small decrease in performance. Understanding in which situations suboptimal or action-free data are beneficial would improve the applicability of the method.", "methods_and_evaluation_criteria": "Yes, the benchmarks correctly support the claims. Additional experiments which would help understanding the role of action-free and suboptimal data could evaluate their role more precisely. For example, the authors could show a graph of the performance (maybe in one or in a subset of tasks) as a function of the amount of suboptimal / action-free data, to better understand whether their impact saturates at a certain point, and how many labeled optimal examples are necessary for the method to be effective.", "theoretical_claims": "The paper makes no theoretical claims.", "experimental_designs_or_analyses": "I do not see particular issues with the experimental design. My only concern is the lack of justification for the choice of labeled / unlabeled / suboptimal subsets, which might have a relevant impact on the results.", "relation_to_broader_scientific_literature": "The method is compared to strong baselines and therefore its relevance in the context of the broader literature is clear.", "essential_references_not_discussed": "The background discussion is broad.", "other_strengths_and_weaknesses": "Strenghts:\n* Clarity of exposition\n* Good comparison with baselines\n* Elegant architecture\n* Convincing results\n\nWeaknesses:\n* Lack of evaluation of how many labeled / unlabeled / suboptimal trajectories are needed for optimal performance", "other_comments_or_suggestions": "108: which enabling\n252 right column - the paragraph on the number of labeled trajectories and suboptimal / action free is difficult to follow, a small table would be clearer\n294: remove use or collect", "questions_for_authors": "- Does using action-free and suboptimal trajectories at the same time further improve the results?\n- Does training the inverse dynamics on suboptimal trajectories, that will likely never be encountered during deployment (the planner will try to follow the optimal trajectories) actually improve action regression?\n- Are the action-free data always from expert policies? To my understanding, suboptimal action-free trajectories cannot be used because they would decrease the performance of the planner.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces Latent Diffusion Planning (LDP), a modular imitation learning framework that separates planning and control by operating in a learned latent space. It combines a planner trained with action-free demonstrations and an inverse dynamics model trained with suboptimal, action-labeled data. Both modules are trained using diffusion-based objectives and rely on a variational autoencoder (VAE) to encode image-based observations into compact latent representations. The key insight is that decoupling planning from action prediction allows each module to exploit broader, weaker forms of supervision. Experimental results on simulated and real-world robotic manipulation tasks demonstrate improved performance and data efficiency compared to existing imitation learning methods, especially in settings with limited expert demonstrations.

The paper offers a creative modular solution for leveraging imperfect demonstration data. The use of a latent space to enable diffusion-based planning and control is intuitive and allows the effective reuse of suboptimal and action-free demonstrations. Empirical results are generally favorable, and real-world experiments strengthen the practical relevance. However, the novelty is somewhat incremental, given similarities to prior works like UniPi, and the improvements over baselines are modest. Experimental inconsistencies—such as varying model sizes, training steps, and dataset sizes across methods—raise concerns about the fairness of some comparisons. Additionally, the lack of ablation for the VAE and limited use of larger-scale datasets or modern pre-trained visual encoders limits the generality and robustness of the findings. However, most of the concerns were addressed by the author's rebuttal, leading to the increased scores.

In sum, I recommend accepting this paper.
Decision: Accept (spotlight poster)