DAIL: Distribution Aligned Imitation Learning

• DAIL is a framework in imitation learning that matches the full distribution of expert behaviors using statistical divergence measures.
• It employs methods like reverse KL, Wasserstein distance, and adversarial losses to align policies across control, state-only, and language-based tasks.
• DAIL demonstrates empirical improvements in RLHF, LLM alignment, and continuous control with efficient policy updates and robust theoretical guarantees.

Distribution Aligned Imitation Learning (DAIL) refers to a class of imitation learning (IL) algorithms whose objective is to match the entire distribution of behaviors, returns, or responses generated by an expert, rather than focusing solely on first-moment statistics such as expected reward. DAIL approaches span both classical control, reinforcement learning, preference modeling, off-policy and state-only settings, and large pre-trained sequence models. Core to all DAIL formulations is the minimization of a statistical divergence between the learner’s induced distribution and the expert’s, utilizing variational, adversarial, or density-ratio approaches depending on the setting.

1. Theoretical Foundations and Objective Specification

The prototypical DAIL objective seeks to align the learner policy distribution $\pi_\theta$ with the expert distribution (or a surrogate, such as the distribution of human-chosen responses in preference learning). The class of objectives is parameterized by a divergence $D(\cdot\|\cdot)$:

• Reverse-KL alignment (response imitation):

$$L_{\rm DAIL}(\theta) = D_{\rm KL}\Bigl(\pi_\theta(\cdot|x) \;\Big\|\; \pi_{\rm chosen}(\cdot|x)\Bigr) = \mathbb{E}_{y\sim\pi_\theta}\left[\log\frac{\pi_\theta(y|x)}{\pi_{\rm chosen}(y|x)}\right]$$

as applied in Direct Imitation Learning (DIL) for aligning LLMs with preference data (Xiao et al., 7 Mar 2025).

• Forward-KL/state occupancy:

$$D_{\rm KL}(d^E \| d^\pi),\quad \text{where}\ d^\pi\ \text{is the policy-induced discounted occupancy}$$

and $d^E$ is the expert's occupancy. ValueDICE and related methods optimize off-policy surrogates of this form (Kostrikov et al., 2019).

• Distributional alignment via return matching:

$\min_\pi\ W_p(P_E(R), P_\pi(R))$

where $W_p$ is the Wasserstein-$p$ distance, $P_E$ and $P_\pi$ are expert and policy return distributions, and the policy class may be non-Markovian to match higher-order moments (Lazzati et al., 15 Sep 2025).

The divergence of choice (KL, Jensen–Shannon, Wasserstein, etc.) and which distribution is placed in the first/second argument induce fundamentally different optimization properties and algorithmic structures.

2. Algorithmic Realizations and Density Ratio Estimation

2.1. Direct Imitation Learning (DIL)

DIL (Xiao et al., 7 Mar 2025) formulates preference-based alignment over sequences as a reverse-KL imitation learning problem, seeking to directly match $\pi_\theta$ to the unknown $\pi_{\rm chosen}$ (distribution over human-preferred completions). The key insight is that by minimizing the reverse KL, the learner increases the likelihood of generating examples sampled from the distribution of chosen responses.

The algorithm estimates the density ratio $$r^*(x, y) = \pi_{\rm chosen}(y|x)/\pi_{\rm ref}(y|x)$$ from pairwise preference data using Bregman divergence minimization (LSIF, UKL, BCE), then updates the policy in closed-form, subsuming DPO, PPO-RLHF, and SFT as special cases.

Key Steps:

• Preference data $\mathcal{D} = \{(x, y_w, y_l)\}$ is used to learn $r_\phi$ by minimizing:

$$L(\phi) = \mathbb{E}_{(x, y_w, y_l)}\left[\frac{1}{2} r_\phi(x, y_l)^2 - r_\phi(x, y_w)\right]$$

• The optimal policy update is:

$$\log\frac{\pi^*(y|x)}{\pi_{\rm ref}(y|x)} = \log r^*(x, y)$$

• Single-loop fine-tuning avoids any inner reinforcement learning subroutines.

2.2. Off-Policy Occupancy Matching and ValueDICE

For RL/control, ValueDICE (Kostrikov et al., 2019) makes use of the Donsker–Varadhan dual of the KL divergence to recast on-policy divergence minimization as a fully off-policy min-max game:

$$J_{\rm DICE}(\pi, \nu) := \log \mathbb{E}_{d^E}\left[e^{\nu(s,a) - B^\pi \nu(s,a)}\right] - (1-\gamma)\mathbb{E}_{s_0 \sim p_0, a_0 \sim \pi}[ \nu(s_0, a_0)]$$

where $B^\pi$ is the zero-reward Bellman operator. The algorithm alternates between optimizing $\nu$ and the policy $\pi$, with all expectations computed off-policy from expert and replay data.

2.3. State-only and Distributional Matching

When only state trajectories are available (LfO), DAIL minimizes the per-step KL divergence between transition densities (Boborzi et al., 2022):

$$J_{\rm DAIL}(\theta) = \sum_{t=0}^{T-1} \mathbb{E}_{(s_t,s_{t+1}) \sim \mu^{\pi_\theta}}\left[ \log \mu^{\pi_\theta}(s_{t+1}|s_t) - \log \mu^E(s_{t+1}|s_t) \right]$$

This decomposes to a soft-actor-critic objective with a reward based on three density models (expert, forward, inverse-dynamics), supporting robust and interpretable convergence.

2.4. Return Distribution Matching and Risk Sensitivity

DAIL extends to strict risk-sensitive settings (Lazzati et al., 15 Sep 2025), seeking not only expected return but full alignment of the return distribution under Wasserstein distance. This necessitates non-Markovian policy classes (policies depending on cumulative reward), with efficient algorithms for both offline and model-based settings:

• RS-BC for unknown transitions (behavior cloning in augmented space)
• RS-KT for known transitions (linear programming over occupancy measures)

Consistency and finite-sample bounds are established for both settings.

2.5. Adversarial and Aggregated Data Alignment

Adversarial frameworks (AILAD (Woillemont et al., 2023); meta-learned DAIL (Chirra et al., 1 Oct 2025)) operate by training a discriminator on pathwise/aggregate metrics or occupancy ratios, using adversarial losses that minimize Jensen–Shannon or learned divergences between the metrics/occupancies of the learner and the expert.

Empirically, AILAD matches or surpasses benchmarks on metric-alignment for heterogeneous or partially observed expert data.

3. Unification with RLHF, Preference Learning, and LLM Alignment

DAIL’s theoretical formalism unifies RL from Human Feedback (RLHF) and modern LLM alignment pipelines under a distribution matching lens (Xiao et al., 7 Mar 2025):

• The two-stage process of RLHF—learning a reward/potential function by forward-KL to chosen data and then updating the policy via (reverse-KL) RL—is precisely an instance of imitation learning plus distillation.
• DPO and other modern preference learning objectives correspond to particular choices of density ratio estimator and loss.
• DAIL provides a single-loop, closed-form alignment alternative that is computationally attractive and empirically superior on LLM benchmarks (e.g., UltraFeedback, Anthropic-HH).

4. Empirical Performance and Practical Considerations

DAIL algorithms achieve or exceed state-of-the-art on a wide spectrum of benchmarks:

• LLM alignment: DIL (DAIL) achieves 1–3 absolute point improvements on BBH, MATH, GSM8K, ARC over DPO, SimPO, IPO (Xiao et al., 7 Mar 2025). On human-preference tasks, DIL attains >75% win-rate vs. SFT and >60% vs. human-chosen completions.
• Continuous control: ValueDICE and state-only DAIL algorithms match expert return with far fewer trajectories than adversarial or behavioral cloning baselines, especially in low-data or offline settings (Kostrikov et al., 2019, Boborzi et al., 2022).
• Risk-sensitive/tabular: Return distribution matching yields lower Wasserstein distances to expert returns vs. occupancy matching or TV-based imitation, with improved sample complexity in model-based settings (Lazzati et al., 15 Sep 2025).
• Language-conditioned RL: Distributional Aligned Learning (DAIL) combining value-distribution estimation and semantic alignment resolves instruction ambiguity, outperforming GCBC, IQL, and CQL on BabyAI and ALFRED by up to 10% (Xie et al., 22 Oct 2025).
• Heterogeneous/aggregated data: AILAD reduces metric JSD to expert pool by 10–30% over CARMI and ablations (Woillemont et al., 2023).
• Algorithmic discovery: Meta-learned reward assignment functions surpass all human-designed adversarial IL baselines in Brax/Minatar domains (Chirra et al., 1 Oct 2025).

Practical recommendations include using a strong, length-normalized reference policy, careful selection/tuning of divergence (LSIF, UKL, BCE), and relying on interpretable, model-derived convergence metrics instead of adversarial discriminator loss.

5. Extensions: Reasoning Models and Language-Based Learning

In the context of LLMs and high-level reasoning, DAIL methods enable sample-efficient learning even when expert data is limited ($<1000$ solutions) and typically out-of-distribution:

• Reasoning Trace Synthesis: DAIL generates in-distribution, student-style reasoning traces from didactic expert solutions using mixed-policy rollouts with a privileged-student model (Mendes et al., 2 Feb 2026).
• Contrastive Learning: Token-level contrastive objectives distill correct reasoning paths while penalizing rationalization artifacts, outperforming both NLL and RLVR on hard benchmarks.
• Efficiency and Generalization: DAIL models achieve 2x–4x greater reasoning efficiency (solving accuracy per token), maintain in/out-of-domain generalization, and avoid catastrophic forgetting.
• Limitations: DAIL requires a minimal model competence threshold, task-dependent waypoint extraction, and sufficiently long trace rollouts for difficult tasks.

6. Connections, Variants, and Theoretical Guarantees

DAIL’s mathematical underpinnings guarantee, under mild assumptions, that the global optimum of the formulated divergence objective recovers the expert distribution (whether over states, returns, or final outputs).

• Minimization of KL, JSD, or Wasserstein divergence yields statistical consistency in the limit of infinite data and expressivity.
• Finite-sample complexity is tightly characterized for tabular and model-based settings (Lazzati et al., 15 Sep 2025).
• DAIL provides a framework for analyzing and deriving many existing and future imitation learning and alignment algorithms, with special cases corresponding to SFT, RLHF with KL regularization, DPO, adversarial IRL, and distributional RL for language and control.

7. Summary Table

Setting	DAIL Variant (Paper)	Divergence/Metric	Empirical Domain
Preference-based alignment	DIL (Xiao et al., 7 Mar 2025)	Reverse KL ($\pi_\theta \|\pi_{chosen}$)	LLMs, RLHF
Off-policy control	ValueDICE (Kostrikov et al., 2019)	Forward KL (occupancy)	MuJoCo, RL-control
State-only observation	SOIL-TDM (Boborzi et al., 2022)	KL (state transitions)	PyBullet, LfO
Return distribution matching	RS-BC, RS-KT (Lazzati et al., 15 Sep 2025)	Wasserstein-1 ($W_1$)	Tabular RL, risk-sensitive IL
Language-conditioned RL	DAIL (Xie et al., 22 Oct 2025)	Value-distribution (C51) + MI	BabyAI, ALFRED
Aggregated metric matching	AILAD (Woillemont et al., 2023)	JSD (empirical metrics)	Procedural games
Meta-learned adversarial IL	DAIL (Chirra et al., 1 Oct 2025)	LLM-discovered RA (Wasserstein)	Brax, Minatar
Reasoning traces, LLMs	DAIL (Mendes et al., 2 Feb 2026)	Token-level contrastive KL	Math/Proof generation

The breadth and depth of DAIL’s formulations and empirical validations underscore its centrality to contemporary imitation learning, large-model alignment, and distribution matching in both classical and language-based RL domains.