Source-linked AI summary
CoDrift: Compositional Drifting for Offline Reinforcement Learning
Xiewei Ni, Ruofeng Mei, Xiangyu Xu
TL;DR
Offline RL must balance staying within the support of a fixed dataset with selecting high-value actions. CoDrift represents these objectives as composable action-space fields and absorbs their composition into a one-step stochastic generator. Across 73 OGBench and D4RL tasks in offline and offline-to-online settings, it achieves the best average rank among compared methods.
Problem
Offline RL requires policies to remain compatible with fixed-dataset support while preferentially selecting high-value actions.
Method
CoDrift composes conditional behavioral, marginal behavioral, and critic-guided value fields into a unified field for a one-step stochastic actor.
Results
CoDrift achieves the best overall average rank across OGBench and D4RL tasks in offline and offline-to-online settings.
Takeaways & Limitations
Compositional drifting provides a competitive alternative to Gaussian-, diffusion-, and flow-based policy classes while retaining one-step stochastic generation at deployment.
Abstract
from arXiv · showhide
Offline reinforcement learning is intrinsically multi-objective: a policy must remain compatible with the behavioral support of a fixed dataset while preferentially selecting high-value actions. We recast these objectives in a common form by viewing each as an action-space motion field that specifies how generated actions should move. This perspective enables heterogeneous learning objectives to be combined directly through field composition. Inspired by drifting models, we propose CoDrift, a compositional framework for one-step generative policy learning. CoDrift combines three objective-level fields into a unified policy field. The conditional field preserves state-dependent behavioral structure, while the marginal field pools actions across states to provide a more stable generative signal in the single-positive-sample regime of continuous-control offline RL. The value field moves generated actions toward higher-value regions. The composed field is absorbed into a stochastic generator that produces an action with a single forward pass at deployment. We evaluate CoDrift on 73 tasks from OGBench and D4RL in both offline and offline-to-online settings. CoDrift compares favorably with state-of-the-art methods and achieves the best average rank in both settings.
1 Introduction
CoDrift frames offline RL as composing action-space fields that balance behavioral support with higher-value actions. It combines conditional and marginal behavioral fields with a value field in a one-step stochastic generative policy.
- Motivation: Offline RL must balance compatibility with dataset support against selecting high-value actions.Actions far outside the data distribution may yield unreliable value estimates, while behavior cloning alone is insufficient.
- Core principle: CoDrift represents heterogeneous learning objectives as additive action-space displacement fields.Behavioral fields move generated actions toward supported regions, while value fields move them toward higher predicted return.
- CoDrift: CoDrift combines conditional, marginal, and value fields into a single generative policy field.The conditional field preserves state-dependent support, the marginal field captures global action structure, and the value field uses critic action gradients.
- Evaluation: CoDrift provides expressive stochastic generation with a single forward pass at deployment.The framework is evaluated across 73 OGBench and D4RL tasks in offline and offline-to-online settings.
- CoDrift: Marginal drifting addresses the single-positive-sample problem by pooling actions across states for a more stable training signal.This is particularly useful in continuous-control offline RL, where each state is essentially paired with one observed action.
- Evaluation: CoDrift achieves the best overall average rank across the full offline benchmark suite and in both evaluated settings.It performs particularly strongly on OGBench while remaining competitive on D4RL.
2 Preliminaries
This section frames offline RL as balancing behavioral fidelity against value maximization, then introduces drifting models and CoDrift’s field-based alternative. Drifting fields move generated samples toward positive references and away from generated negatives, with training regressing a one-step generator toward displaced targets.
- 2.1 Offline RL: Offline RL maximizes return from a fixed dataset while constraining policies to remain near the unknown behavior distribution.Distribution shift makes critic extrapolation outside the dataset support unreliable.
- 2.1 Offline RL: Behavior-regularized actor–critic methods balance value maximization with an explicit penalty toward the offline behavior distribution.The regularization coefficient controls the trade-off between behavioral fidelity and value maximization.
- 2.1 Offline RL: CoDrift represents behavioral fidelity and value maximization as action-space displacement fields, then composes them into a joint policy field.This realizes the same balancing principle as behavior-regularized methods without combining objectives directly at the loss level.
- 2.2 Drifting Models: Drifting models use a one-step generator whose samples are shifted during training rather than transported through multiple inference-time states.The generator maps prior noise to samples, while training constructs displacement-based targets.
- 2.2 Drifting Models: The drifting field combines kernel-normalized attraction toward positive references with repulsion from generated samples to discourage mode collapse.The Laplace kernel weights nearby references more heavily, with τ controlling effective neighborhood size.
- 2.2 Drifting Models: Training forms each target as x + Vp,q(x), treats it as fixed, and regresses the generator toward that displaced sample using minibatch approximations.The minibatch field excludes each query particle from its own negative reference set.
3 Method
CoDrift represents behavioral matching and value improvement as action-space displacement fields, then composes them into a one-step stochastic policy generator. Its conditional and marginal behavioral fields provide complementary state-dependent and population-level signals, while the value field directs actions toward higher predicted return.
- 3.4 Compositional Drifting: CoDrift represents conditional behavioral matching, marginal behavioral matching, and value maximization as three action-space fields composed into one policy field.The conditional and marginal fields preserve behavioral structure at complementary statistical scales, while the value field pushes actions toward higher-value regions.
- 3.1 Conditional Behavioral Field: A one-step stochastic generator produces each policy action from a state and one noise draw, requiring a single forward pass at deployment.The generator uses tanh to keep actions within [−1, 1]^d.
- 3.1 Conditional Behavioral Field: The conditional behavioral field matches the offline conditional action distribution by attracting generated actions toward the action observed at the corresponding state and repelling same-state generated actions.The conditional field treats the dataset conditional distribution as the target and the policy conditional distribution as the generated distribution.
- 3.1 Conditional Behavioral Field: Because each exact continuous-control state is typically paired with one observed action, conditional drifting has a high-variance one-sample attractive estimate.This limitation motivates adding a field that pools distributional information across states.
- 3.2 Marginal Behavioral Field: Marginal drifting pools positive actions across states and generated actions across states, removing state–action pairing to provide a more stable, lower-variance generative signal.The marginal field uses B positive samples and complements the conditional field's preservation of state–action correspondence.
- 3.3 Value Field: The value field is the action gradient of the critic, specifying the local action-space direction of increasing predicted return.Unlike the behavioral fields, it is not associated with a reference distribution, but it shares their displacement-field representation.
- 3.4 Compositional Drifting: Nonnegative weights ηc, ηm, and ηv control the relative contributions of conditional matching, marginal matching, and value maximization in the composed field.The generator is trained toward a displaced target formed by adding the composed field to each generated action.
4 Experiments
CoDrift is evaluated across diverse OGBench and D4RL tasks in offline and offline-to-online settings, including comparisons with representative policy classes and an ablation of marginal drifting. It achieves the best overall average rank in both settings, while marginal drifting improves performance over conditional drifting alone.
- Experimental Setup: CoDrift is evaluated on continuous-control benchmarks spanning 73 tasks from OGBench and D4RL in offline and offline-to-online settings.The evaluation follows existing benchmark suites and protocols, including periodic evaluation without selecting the best-performing checkpoint.
- Experimental Setup: The comparison covers eleven representative Gaussian-, diffusion-, and flow-based offline RL methods.Baselines include behavior cloning, value-based and diffusion policies, and flow-based policies.
- Offline RL Performance: CoDrift achieves the best overall average rank across the offline benchmark suite, performing particularly strongly on OGBench while remaining competitive on D4RL.Lower average rank is better, and the pixel-based row is excluded from ranking because several baselines are unavailable.
- Offline-to-Online Performance: CoDrift attains the best overall average rank across 15 offline-to-online tasks and remains effective after transitioning to online policy improvement.Fine-tuning adds newly collected transitions to the replay buffer and uses the same approach as the offline phase, without a separate online-stage objective.
- Marginal Drifting Ablation: Adding marginal drifting improves performance on all four representative OGBench tasks, with gains ranging from 6 to 8 percentage points and averaging 6.8 points.Marginal drifting pools actions across states, producing a larger batch-level positive sample that complements conditional state–action supervision.
5 Related Work
Related work frames CoDrift as a generative, behavior-regularized offline RL method that represents behavioral fidelity and value maximization as composable action-space displacement fields. It builds on expressive generative policy classes while using conditional and marginal behavior constraints.
- Offline RL and Behavior Regularization: Offline RL methods address extrapolation errors through value regularization, in-sample learning, or policy regularization toward the behavior distribution.Policy regularization may use explicit penalties or weighted regression within actor–critic formulations.
- Offline RL and Behavior Regularization: CoDrift follows behavior regularization while representing behavioral fidelity and value maximization as action-space displacement fields.Its behavioral constraints operate at both conditional and marginal levels and are composed with a critic-guided value field.
- Generative Policies: Generative policy research spans autoregressive, adversarial, diffusion, and flow-matching paradigms for modeling complex and multimodal action distributions.These paradigms provide the broader policy-class context for CoDrift's generative formulation.
6 Concluding Remarks
CoDrift represents heterogeneous offline RL objectives as composable action-space displacement fields and uses their sum to define a one-step stochastic actor. Across offline and offline-to-online tasks, it achieves the best overall average rank and motivates objective compositionality as a broader design perspective.
- CoDrift: CoDrift combines conditional behavioral, marginal behavioral, and critic-guided value fields into a joint field for offline RL.The conditional field preserves state-dependent behavior, the marginal field captures population-level action structure, and the value field promotes higher-return actions.
- CoDrift: The joint field directly defines the regression target of a one-step stochastic actor.This supports native single-forward-pass stochastic generation at deployment.
- Empirical Results: CoDrift achieves the best overall average rank across broad OGBench and D4RL evaluations in both offline and offline-to-online settings.The paper presents compositional drifting as a competitive alternative to Gaussian-, diffusion-, and flow-based policy classes.
- Broader Perspective: Objective compositionality allows different requirements to be translated into action-space displacements, additively composed, and absorbed into one generative policy.The authors propose this perspective as a foundation for incorporating additional objectives and constraints.