Technical Summary: Language-Critique Imitation Learning from Suboptimal Demonstrations

Problem Statement

Imitation learning (IL) from suboptimal demonstrations aims to learn robust policies by leveraging abundant, imperfect trajectories alongside limited expert data. However, existing methods typically rely on compressed scalar supervision signals—such as confidence estimates, discriminator scores, importance weights, or rewards—to distinguish useful behaviors from noise. The authors argue that these scalar signals are inherently limited because they cannot explicitly express intermediate reasoning regarding task progress, specific failure modes, or corrective actions. Consequently, they often fail to capture the structured information necessary for learning in long-horizon or multimodal tasks where success depends on subtle, stage-specific decisions.

Methodology

The paper proposes Language-Critique Imitation Learning, an offline framework that replaces scalar signals with natural language as a structured supervision signal. The method consists of three core components:

1. Structured Language Label Generation

To avoid the variability and cost of human annotation or free-form Large Language Model (LLM) generation, the authors introduce a deterministic Language Label Generator ($\mu_g$). For any state-action pair $(s, a)$, $\mu_g$ extracts task-specific heuristics (e.g., object distance, gripper status) and routes them to three specific selectors to generate a composite language label $l$:

• Task Progress ($<T>$): Describes the current semantic stage of the task (e.g., "approaching object," "moving toward target").
• Action Optimality ($<A>$): Evaluates whether the action is beneficial, ineffective, or harmful for the current context (e.g., "The action is good," "The action is bad").
• Movement Guidance ($<M>$): Provides fine-grained corrective instructions (e.g., "Move to the right gently," "Pitch toward the bottom").

This structured approach ensures the labels are consistent, reproducible, and directly relevant to the control task without relying on privileged environmental information.

2. Differentiable LLM-Captioner ($\mu_\phi$)

To enable end-to-end policy optimization, the discrete language labels generated by $\mu_g$ are distilled into a differentiable LLM-Captioner ($\mu_\phi$). This model, built upon a pretrained LLM (specifically SmolLM2) and an MLP projector, maps state-action pairs to the structured language label distribution.

• Training: $\mu_\phi$ is fine-tuned on the general dataset (containing both expert and suboptimal data) using a token-level cross-entropy loss to mimic $\mu_g$.
• Role: Once trained, $\mu_\phi$ acts as a frozen differentiable critic that can evaluate the likelihood of a policy's action $\hat{a}$ producing the expert's language label $l^E$.

3. Language-Critique Loss ($L_{LC}$)

The core learning objective is the Language-Critique Loss, which directly trains the policy $\pi_\theta$ using the structured signals from $\mu_\phi$ without reducing them to scalars.

• Mechanism: The loss penalizes policy actions that make the expert language label less likely than the expert action itself. It is formulated as a clipped difference between the policy's cross-entropy loss and the expert's cross-entropy loss:
  $$L_{LC}(\pi_\theta, \pi_E) = [L_{CE}(\pi_\theta) - \text{sg}(L_{CE}(\pi_E))]_+$$
  where $[x]_+ = \max(x, 0)$ and $\text{sg}$ denotes stop-gradient.
• Instantiations:
  • LC-BC: Applied to standard feedforward Behavior Cloning (BC).
  • LC-DP: Applied to Diffusion Policies. Since diffusion models predict noise rather than actions, the loss is applied to the one-step reconstructed action $\hat{a}_0$. A timestep-dependent reweighting factor $\omega_k$ is introduced to balance the contribution of different diffusion steps, ensuring the loss aligns with the standard diffusion objective.

Key Contributions

1. Framework: Introduction of a language-critique framework that utilizes natural language as a structured, non-scalar supervision signal for imitation learning from suboptimal demonstrations.
2. Algorithm: Development of LC-BC and LC-DP, which integrate language-guided supervision into behavior cloning and diffusion policies, respectively.
3. Theoretical Guarantee: A theoretical proof showing that minimizing the language-critique gap upper-bounds the expert-policy performance gap under standard assumptions (linear feature realizability and language sufficiency).
4. Empirical Validation: Comprehensive evaluation across eight diverse continuous-control tasks (navigation, manipulation, gameplay) demonstrating consistent outperformance over strong baselines.

Experimental Results

The authors evaluated their methods on eight tasks: MAZE, PARKING, SWEEP, BOX-CLOSE, BLOCKPUSH, PEGINSERT, HAMMER, and RELOCATE.

• Performance vs. Baselines: LC-BC and LC-DP consistently outperformed or matched strong imitation learning baselines (e.g., DemoDICE, DWBC, ILID) and offline RL methods (e.g., CQL, TD3+BC, Decision Transformer, Diffusion Policy).
  • Notable gains were observed in multimodal (BLOCKPUSH) and precision-control (PEGINSERT, HAMMER) tasks. For instance, on BLOCKPUSH, LC-BC improved success rates from 34.4% (standard BC) to 47.2%.
  • In the precision task PEGINSERT, LC-DP maintained a clear advantage over standard Diffusion Policy, while LC-BC showed modest improvements, suggesting diffusion policies better integrate high-level language guidance through iterative denoising.
• Mixed-Quality Data: The method remained effective even when the expert dataset was not cleanly separated, using language-based heuristics to filter suboptimal data from a mixed dataset.
• Ablation Studies:
  • Structured vs. Scalar: Methods using language labels significantly outperformed variants that collapsed labels into scalar rewards (LC-BC-reward) or discrete categorical classes (LC-BC-classifier), confirming the value of preserving natural language expressiveness.
  • Component Importance: All three components ($<T>, <A>, <M>$) were found to be complementary; removing any degraded performance.
  • Label Source: The structured generator $\mu_g$ proved superior to open-ended VLM-generated labels (e.g., from o4-mini) due to better semantic consistency and feature alignment, particularly for fine-grained movement guidance.

Significance and Claims

The paper claims that natural language serves as a powerful and structured form of supervision that is superior to compressed scalar signals for learning from suboptimal data. By explicitly describing task progress, failure modes, and corrective actions, language labels provide the structured information necessary to guide policies in complex, long-horizon, and multimodal environments.

The authors emphasize that their approach does not merely use language as a descriptive aid but as a direct, differentiable supervision signal that enables robust policy learning without requiring additional environment interaction or expert data. The theoretical results further ground this empirical success, showing that the proposed objective controls the performance gap to the expert policy. The work suggests that language can effectively bridge the gap between limited expert data and abundant suboptimal data, offering a scalable path toward robust imitation learning in continuous control.