PaperlayerView exact revision on arXiv

Source-linked AI summary

Goal-Conditioned Imitation Learning using Score-based Diffusion Policies

Moritz Reuss, Maximilian Li, Xiaogang Jia, Rudolf Lioutikov

arXiv:2304.02532v2cs.LGcs.RO

TL;DR

Learning goal-conditioned policies from uncurated play data is difficult because demonstrations are multimodal and conventional approaches may require task labels, hierarchy, or clustering. BESO uses a score-based diffusion policy with classifier-free guidance, achieving effective benchmark behavior with three denoising steps while also supporting goal-independent policies.

  • Problem

    Learning from play requires policies that handle multimodal demonstrations while avoiding costly labeled trajectories and limited predefined-task coverage.

  • Method

    BESO models goal-conditioned action distributions with score-based diffusion, using fast inference and classifier-free guidance to learn goal-dependent and goal-independent policies.

  • Results

    BESO consistently outperforms several state-of-the-art GCIL methods on challenging benchmarks while generating actions in 3 denoising steps.

  • Takeaways & Limitations

    Score-based diffusion provides an expressive policy representation for multimodal behavior generation from reward-free play data without hierarchical structures or additional clustering.

Abstract

from arXiv · show

We propose a new policy representation based on score-based diffusion models (SDMs). We apply our new policy representation in the domain of Goal-Conditioned Imitation Learning (GCIL) to learn general-purpose goal-specified policies from large uncurated datasets without rewards. Our new goal-conditioned policy architecture "$\textbf{BE}$havior generation with $\textbf{S}$c$\textbf{O}$re-based Diffusion Policies" (BESO) leverages a generative, score-based diffusion model as its policy. BESO decouples the learning of the score model from the inference sampling process, and, hence allows for fast sampling strategies to generate goal-specified behavior in just 3 denoising steps, compared to 30+ steps of other diffusion based policies. Furthermore, BESO is highly expressive and can effectively capture multi-modality present in the solution space of the play data. Unlike previous methods such as Latent Plans or C-Bet, BESO does not rely on complex hierarchical policies or additional clustering for effective goal-conditioned behavior learning. Finally, we show how BESO can even be used to learn a goal-independent policy from play-data using classifier-free guidance. To the best of our knowledge this is the first work that a) represents a behavior policy based on such a decoupled SDM b) learns an SDM based policy in the domain of GCIL and c) provides a way to simultaneously learn a goal-dependent and a goal-independent policy from play-data. We evaluate BESO through detailed simulation and show that it consistently outperforms several state-of-the-art goal-conditioned imitation learning methods on challenging benchmarks. We additionally provide extensive ablation studies and experiments to demonstrate the effectiveness of our method for goal-conditioned behavior generation. Demonstrations and Code are available at https://intuitive-robots.github.io/beso-website/

I. INTRODUCTION

BESO addresses the difficulty of learning general-purpose goal-conditioned behavior from uncurated play data, whose multimodality challenges conventional imitation-learning policies. It uses score-based diffusion policies to model expressive goal-conditioned behavior, enable three-step action generation, and jointly learn goal-dependent and goal-independent policies.

  • Motivation: Goal-conditioned imitation learning seeks versatile agents from offline data, but commonly depends on predefined tasks and costly labeled, segmented expert trajectories.Such methods also generalize poorly to new scenes and different tasks.
  • Motivation: Play data is unsegmented and multimodal: one task may have several solutions, while different tasks may share similar behaviors.Policies must preserve multiple solutions, distinguish related behaviors, and support long-horizon skills.
  • BESO: BESO uses score-based diffusion models as a policy representation for goal-conditioned behavior generation from reward-free, offline datasets.The model learns to reverse a noise-adding diffusion process and generate samples iteratively from noise.
  • BESO: BESO captures multimodal action distributions without hierarchical policies or additional clustering, and consistently outperforms C-BeT and Latent Motor Plans on challenging GCIL benchmarks.The paper attributes effective conditioned behavior generation to the expressiveness of score-based diffusion models.
  • Classifier-free guidance: Classifier-free guidance enables BESO to learn goal-dependent and goal-independent policies simultaneously, usable together or independently at test time.The method uses a single training objective without additional rewards.
  • Efficient generation: 3 denoising steps let BESO generate actions faster than diffusion-based policies requiring 30+ steps, using decoupled training and inference with fast numerical solvers.The approach also uses a Transformer architecture with preconditioning for step-based action synthesis.

III. PROBLEM FORMULATION AND METHOD

The method converts unstructured play trajectories into sequence-goal training examples and learns a policy over the resulting dataset. Because multiple trajectories can reach the same goal, the policy must represent multimodal behavior.

  • A. Problem Formulation: GCIL aims to learn a general-purpose goal-conditioned policy from uncurated, task-agnostic play trajectories.The formulation starts from unstructured trajectories rather than predefined demonstrations.
  • A. Problem Formulation: Each trajectory is split into tuples containing a sub-trajectory sequence and one or more later goal states.The sequence and goal state are paired as a training example.
  • A. Problem Formulation: A goal state must come from the same trajectory as its sequence and begin after that sequence ends.This ordering defines the behavior considered relevant for reaching the selected goal.
  • A. Problem Formulation: The dataset may contain overlapping sequences, and the final play dataset is formed by combining the per-trajectory example sets.Relative indices identify state-action positions within each sequence.
  • A. Problem Formulation: State-action pairs preceding a goal are treated as optimal behavior for reaching that goal, and the policy maximizes their dataset log-likelihood.This defines the behavioral learning target used by the formulation.
  • A. Problem Formulation: Several trajectories may lead to the same goal state, so successful optimization requires a policy capable of encoding multimodal behavior.The multimodality arises directly from the demonstrations.

B. Score-based Diffusion Policies

BESO learns a conditional action distribution by diffusing play-data actions into noise and modeling the reverse process with a score function. Its probability-flow ODE removes stochastic injections, allowing deterministic generation from a tractable Gaussian prior.

  • B. Score-based Diffusion Policies: BESO defines a continuous diffusion process whose initial action distribution is the play-data policy and whose later distributions add progressively more Gaussian noise.The process spans t in [0, T], ending in a tractable prior.
  • B. Score-based Diffusion Policies: The diffusion process is specified by a stochastic differential equation, with the score function guiding the reverse transformation toward the data distribution.The score is ∇_a log p_t(a|s, g), while the Wiener process represents infinitesimal Gaussian noise.
  • B. Score-based Diffusion Policies: Noise levels are scheduled through σ_t and β(t), and each timestep defines a corresponding noisy marginal action distribution.BESO uses σ_t(t) = t and models Gaussian perturbations of play-data actions.
  • B. Score-based Diffusion Policies: The terminal distribution is an unstructured Gaussian prior that contains no information about the play-data distribution.Sampling begins from this prior before reverse denoising.
  • B. Score-based Diffusion Policies: BESO uses the probability-flow ODE, which preserves the SDE’s marginal distributions while eliminating additional random-noise injections.Setting β(t) = 0 recovers the probability-flow ODE.
  • B. Score-based Diffusion Policies: The negative score defines the vector field used to move samples through the current marginal distribution during deterministic generation.Its scale depends on the noise level and its rate of change.
  • B. Score-based Diffusion Policies: BESO training requires a play dataset, a score model, and a log-logistic noise distribution within an iterative training loop.The supplied training algorithm lists these as its principal inputs and training components.

C. Diffusion Training

BESO trains a noise-conditioned score model to approximate the score across action distributions corrupted at different noise levels, then generates actions by reversing the diffusion process. Sampling begins from Gaussian noise and iteratively denoises it with a fast DDIM solver.

  • Score-model training: The denoising network Dθ is trained to match the score function across marginal action distributions pt(a|s, g).The score estimate is required for numerically approximating the reverse ODE.
  • Score-model training: Gaussian noise is added to actions, and denoising score matching minimizes the difference between the network output and original actions.
  • Noise-level training: Noise levels are sampled from a truncated log-logistic training distribution over {σmin, σmax}, with losses weighted by α(σt).The stated distribution uses location parameter α and scale parameter β.
  • Action generation: Action generation samples aT from a Gaussian prior and substitutes the learned score model into the reverse ODE or SDE.
  • Action generation: BESO uses DDIM for fast, deterministic sampling, and ablations find that three denoising steps generate high-accuracy actions.More inference steps marginally improve performance while significantly slowing sampling.

IV. GOAL-GUIDED SCORE-BASED DIFFUSION POLICIES

BESO provides conditioned and classifier-free guided score-based diffusion policies for synthesizing goal-directed actions. Its architecture addresses broad noise ranges with preconditioning and uses Transformer-based temporal modeling, while the DDIM sampler supports efficient generation.

  • Conditioned policy: C-BESO directly learns a goal-and-state-conditioned action distribution that captures multiple play-data solutions while supporting long-term goals.
  • Guided policy: CFG-BESO combines BESO with Classifier-Free Guidance to synthesize goal-conditioned behavior.
  • Guided policy: Goal dropout trains an implicit goal-independent policy within the same goal-conditioned diffusion model.The training procedure applies a dropout rate of 0.1 to the goal.
  • Guided policy: The guidance factor λ balances the goal-conditioned and goal-independent gradients during denoising.The paper states that λ commonly ranges from 2 to 7.5 in diffusion literature.
  • Model architecture: Additional skip connections and preconditioning layers help the score model handle noise levels σt ∈ {0.001, 40}.
  • Model architecture: BESO uses a causally masked Transformer inner model to learn temporal relations between observations and actions.States, noise, and noisy actions are embedded into a shared linear representation.

V. EVALUATION

BESO is evaluated against multiple goal-conditioned imitation-learning baselines, including unimodal, hierarchical, energy-based, clustered-transformer, and diffusion policies. Experiments use repeated rollouts on kitchen and block-push tasks, with fewer seeds and baselines for computationally expensive CALVIN training.

  • Evaluation goals: The evaluation asks whether BESO is competitive, which components enable fast sampling, and whether Classifier-Free Guidance supports goal-conditional synthesis.
  • Baselines: The compared baselines include GCBC, RIL, LMP, C-IBC, C-BeT, and Diffusion-X.These span unimodal, hierarchical, energy-based, clustered-transformer, and DDPM-based policies.
  • Results: C-BESO and CFG-BESO consistently outperformed all baselines on conditioned Block-Push and Relay Kitchen despite using only 3 inference steps.
  • Results: CX-Diff with 3 inference steps achieved 2.74(±0.26) on Relay-Kitchen.
  • Evaluation protocol: Kitchen and block-push evaluations used 10 seeds and 100 rollouts each, while CALVIN used 3 seeds and fewer baselines because training was computationally costly.

B. Simulation Experiments

BESO was evaluated against baseline methods across three simulation benchmarks using challenging goal-conditioned settings. It consistently performed strongly, including on unseen goal states and multi-task executions, while requiring only three denoising steps.

  • Relay-Kitchen and Block-Push: BESO consistently outperformed competing methods on Relay-Kitchen and Block-Push across 10 seeds, with low variance indicating robustness.Diffusion policies performed strongly overall, while C-BeT was comparable on Block-Push.
  • Relay-Kitchen and Block-Push: 3 denoising steps enabled BESO to outperform baselines while CX-Diff achieved only 2.74(±0.26) average result in Kitchen when limited to three steps.BESO used three denoising steps on both environments, compared with 50(20) + 8 steps for CX-Diff.
  • Relay-Kitchen and Block-Push: 0.012 seconds per action prediction made BESO over 10 times faster than CX-Diffusion’s 0.15 seconds on a modern desktop PC.
  • CALVIN: BESO achieved the best overall performance on 13 difficult CALVIN single tasks conditioned on unseen goal images without nearby end-effector guidance.The setting required models to infer environmental changes and execute tasks without relying on the end-effector position.
  • CALVIN: BESO and CFG-BESO outperformed other models on two tasks from one goal image, although CFG-BESO performed slightly worse.The experiment used goal images with the end-effector positioned away from both tasks.

D. BESO design choices

The ablations examine conditioning, sampling, stochasticity, and classifier-free guidance as design choices for BESO. They support sequential conditioning, fast deterministic generation, and CFG’s ability to recover goal-independent behavior, while exposing guidance trade-offs.

  • Conditioning Method: 0.93 to 0.91 on Block-Push and 3.76 to 3.4 on Kitchen marked the performance drop from sequential conditioning to FiLM conditioning.FiLM also required additional MLP models and increased overall model capacity, so BESO used sequential conditioning.
  • Classifier-Free Guidance: CFG provided effective goal-conditioning for step-based action generation, with slightly lower average results but equal average reward than standard goal-conditioning.CFG-BESO also learned effectively in image-based CALVIN.
  • Classifier-Free Guidance: λ = 0 recovered an unconditional policy that produced random rollouts with high rewards and low results.The low average result indicates the policy ignored the goal state and acted from the current state alone.
  • Classifier-Free Guidance: λ = 1.25 was empirically best for most tested environments, whereas higher values reduced performance in high-dimensional action spaces and indicated instability.The authors hypothesize that goal guidance is most crucial when selecting which task to solve during rollouts.
  • Sampling Algorithm: 3 denoising steps supported fast deterministic action generation with BESO’s DDIM-based sampling design.BESO numerically approximates the reverse ODE from Gaussian-prior samples and iteratively denoises them.
  • Design Summary: BESO’s score-based representation learned task-agnostic behavior from offline reward-free play data without hierarchical structures or additional clustering.The conclusion also reports simultaneous learning of goal-dependent and goal-independent policies through CFG.

A. BESO Hyperparameters

BESO’s hyperparameters cover normalization, noise scheduling, optimization, and reverse-diffusion discretization. The reported recommendations favor consistent diffusion settings while tuning task-dependent transformer choices.

  • General settings: Transformer dropout rates require task-specific tuning, whereas general diffusion hyperparameters remain consistent across tasks.
  • Normalization: Actions are optimally diffused over [−1, 1] with a noise range of 0.005, 1; larger action ranges require higher noise values.For [−3, 3] actions, the recommended noise range is {0.4, 40}.
  • Training noise: BESO uses the log-logistic noise distribution LogLogistic(α = 0.5, β = 0.5) because prior log-normal recommendations were not optimal for action diffusion.
  • Optimization: Experiments use Adam or AdamW with a standard learning rate of 1e−4, alongside Exponential Moving Average weight optimization.
  • Time steps: Exponential time steps are most effective on average, while linear and Karras schedules can perform comparably on individual tasks.
  • Recommendations: The recommended starting configuration combines noise range {0.005, 1}, exponential time steps, and the DDIM solver.Other samplers, including Euler Ancestral, and linear time steps may improve performance on particular tasks.

B. Sampler Ablation

The sampler ablation compares ODE and SDE-based solvers for step-based action generation across environments. DDIM provides the best overall deterministic performance, and additional inference steps do not reliably improve averages.

  • Solver set: The study evaluates first-order Euler and second-order Heun solvers, alongside DPM-family samplers and stochastic ancestral variants.
  • Deterministic samplers: DDIM achieves the best overall performance among the evaluated deterministic solvers across kitchen and block-push comparisons.The comparison averages 10 seeds with 100 rollouts per environment.
  • Inference steps: Increasing inference steps does not significantly improve average performance and can reduce some samplers’ average results.

C. Baselines Implementation

The baseline implementations use matched architecture and training settings where possible, while retaining method-specific configurations and evaluation procedures. The appendix also specifies BESO’s deterministic and stochastic Euler sampling algorithms.

  • Architecture and training: All diffusion models use the same transformer backbone, while C-BeT uses its recommended parameters.MLP baselines use four layers with 512 neurons and ReLU activations.
  • GCBC: GCBC concatenates the goal with the state and feeds both into a four-layer MLP with dropout rate 0.1.
  • GC-IBC: GC-IBC uses InfoNCE with energy regularization and Wasserstein Gradient loss; λ = 0.005 improved training stability in experiments.Its results were sensitive to initial random seeds.
  • Sequence and latent-plan baselines: C-BeT uses recommended parameters, while Latent Motor Plans uses extensive hyperparameter sweeps and HULC uses prior-work settings on CALVIN.
  • Diffusion-X: Diffusion-X uses 50 inference steps plus 10 fine-tuning steps on kitchen, and 20 inference steps plus 8 fine-tuning steps on block-push.
  • BESO samplers: The deterministic Euler sampler starts from Gaussian action noise and iteratively updates actions using the score-denoising model and noise scheduler.
  • BESO samplers: The stochastic Euler sampler adds ancestral noise scheduling to the iterative action updates.
Loading 2304.02532v2…