- The paper introduces a discrepancy-aware co-training framework that jointly optimizes policy parameters and per-sample weights to maximize target-domain generalization.
- It employs k-nearest neighbor estimates in the latent policy space to assess instance-level discrepancies without relying on explicit feature alignment.
- Experimental results on robotic manipulation tasks show enhanced success rates and improved out-of-distribution generalization in both sim-to-sim and sim-to-real settings.
BEACON: Theory-Grounded Cross-Domain Co-Training for Generative Robot Policies
Introduction and Theoretical Rationale
BEACON introduces a principled discrepancy-aware framework for training high-dimensional generative visuomotor robot policies when large numbers of source demonstrations and limited (expensive) target demonstrations are available. The motivation stems from the empirical degradation of behavior cloning (BC) policies under domain shift, including both visual and physical discrepancies (e.g., sim-to-real gaps in appearance and dynamics). Prior solutions, such as domain randomization, explicit feature alignment, or fixed-ratio data mixing, lack theoretical grounding in maximizing target-domain generalization.
BEACON operationalizes Best-Effort Adaptation (BEA) [awasthi2023best], reframing cross-domain co-training as an instance-level importance reweighting problem. Under this formulation, one jointly optimizes 1) policy parameters and 2) per-sample weights on source (and target) data, minimizing a bound on target-domain risk that encapsulates both empirical fit and distributional discrepancy (Figure 1).
Figure 1: BEACON's cross-domain policy co-training: joint optimization of a diffusion-based policy and per-sample importance weights to minimize a bound on target generalization error.
Method: Discrepancy-Aware Joint Learning with Diffusion Policies
BEACON's objective extends the BEA theoretical framework to high-dimensional sequence prediction via the following instantiation:
- The policy is a diffusion-based parameterization πθ​ over action sequences, trained via behavior cloning with per-example loss.
- Source and target samples receive individualized weights. The objective penalizes empirical error, source-target discrepancy (measured per-sample in latent space), capacity via weight-norm regularization, deviation from the target-only empirical distribution, and total source weight.
- Crucially, the discrepancy term is estimated at the instance level (not domain-level), with a scalable implementation using k-NN distances in policy encoder embedding space.
Joint optimization over θ and the weight vector is carried out via stochastic alternating updates: policy updates using weighted losses, and periodic weight updates driven by projected subgradient steps. Multi-source extension is realized by introducing an additional optimization over source-domain weights, effectively learning to arbitrate among multiple heterogeneous source datasets by their measured relevance for target generalization.
Discrepancy Estimation and Policy Embedding Space
A practical challenge is the estimation of a meaningful discrepancy term in high-dimensional observation-action spaces. Contrary to prior work, BEACON does not assume direct feature-level alignment targets. Instead, instance-level discrepancies di​ are estimated via k-NN distances in the (co-adapting) visual encoder space. This approach co-evolves sample relevance with the state of the policy, preferentially emphasizing source trajectories that are close to the target distribution in embedding geometry. Two alternative discrepancy estimators are also considered: localized labeled discrepancy maximization and a binary domain classifier in latent space. Empirically, k-NN is shown to provide the most robust and interpretable weighting.
Experimental Results: Robust Generalization and Multi-Source Transfer
BEACON demonstrates significant advantages over fixed data mixing and explicit feature alignment baselines (including MMD and UOT) in extensive experiments across three robotic manipulation benchmarks (block stacking, mug cleanup, threading), under both sim-to-sim and sim-to-real settings. The challenging nature of these tasks is reflected in their requirements for precise, long-horizon action sequencing under significant observation and embodiment shift (Figure 2).
Figure 2: Experimental setup illustrating the induced domain gaps (texture, viewpoint, real) between source and target datasets.
Numerical results indicate that BEACON achieves superior average success rates, with pronounced improvement in both in-distribution and out-of-distribution (OOD) evaluations (see main experimental tables in the paper). When extended to multi-source settings (MS-BEACON), the framework further leverages heterogeneous simulation domains, combining their strengths by learned weighting and mitigating negative transfer.
Feature-level analysis via UMAP projections shows that, despite the absence of any explicit alignment loss, source and target demonstrations become aligned in policy latent space as a consequence of BEACON's discrepancy-aware optimization (Figure 3).
Figure 3: UMAP visualization of block stacking latent features; source and target encode to overlapping regions as an emergent property of BEACON.
Ablation studies confirm that BEACON's selection mechanism is robust to different discrepancy estimators and weight-update schedules, while compute-efficient and insensitive to moderate hyperparameter variation.
Generalization to Out-of-Distribution States
The framework's robustness is most clearly demonstrated under OOD evaluation, where the target test states differ spatially from those observed in limited target demonstrations (Figure 4). BEACON is able to leverage the relevant portions of the source dataset to provide coverage and generalization, outperforming feature matching methods that are more susceptible to distributional collapse or misalignment under state shift.
Figure 4: Placement regions for OOD evaluation: BEACON enables strong generalization to red regions excluded from both source and target demonstration scenes.
Policy Weight Dynamics and Data Efficiency
Analysis of the optimization process reveals that BEACON dynamically suppresses harmful or mismatched source samples, amplifies target samples within their admissible range, and learns a distinct, data-driven selection among available demonstrations. Moreover, increasing the quantity of target demonstrations compounds rather than obviates the benefits of cross-domain source reweighting, thereby addressing practical scenarios in which target data can be incrementally expanded.
Implications and Future Directions
BEACON establishes the practical efficacy of theoretically driven, discrepancy-aware instance reweighting for robot policy transfer when operating under severe data imbalance and domain shift. Rather than relying on heuristic mixing or feature-level global alignment, it offers a robust, scalable joint optimization solution that provably anchors source influence to generalization-relevant examples in the latent policy space.
The framework is directly extensible to multi-source adaptation, heterogeneous data pools, and (with suitable future work) other forms of high-capacity generative policies. Notably, it situates importance reweighting—rather than feature alignment—as the fundamental primitive for cross-domain robot policy learning.
Directions for future research include adaptation to highly dynamic tasks where both observation and action domain gaps are amplified, automatic estimation of optimal sum-budget hyperparameters, and application to egocentric datasets and diverse real-world robot fleets.
Conclusion
BEACON provides a rigorously grounded, empirically validated method for cross-domain co-training of generative robot policies, centering policy learning on maximizing target-domain generalization through adaptive, per-sample importance weighting. This approach yields substantial gains in task success, OOD generalization, and robustness over co-training and alignment-based baselines, and suggests that advances in domain adaptation for robot learning should prioritize discrepancy-based instance selection over heuristic representation matching.
Figure 5: UMAP visualization of latent features for mug cleanup; BEACON yields substantial source-target overlap without explicit alignment.