Source-linked AI summary
Memorize, Adapt, Ignore: Diagnosing Robot Learning Mechanisms under Training Data Variation
Ke Zhang, Danica J. Sutherland, Chao Liu
TL;DR
Robot-learning practitioners lack clear diagnostics for how training-data variation shapes policy learning, despite relying on it to improve robustness. The paper uses representation kernels, primarily the eNTK, to distinguish memorization, adaptation, and invariance; it reports that variation shapes these regimes across model classes and that the diagnostics track hardware performance.
Problem
Training-data variation is a key lever for robot-policy robustness, but its mechanisms are poorly understood and factor-selection decisions often rely on intuition or costly trial and error.
Source: §1 · p. 1Method
The paper applies representation-kernel statistics, primarily the eNTK, to diagnose whether policies memorize conditions, adapt across relevant variation, or ignore irrelevant factors.
Source: §1 · p. 2Results
Training-data design shapes learning regimes across model classes: relevant variation can shift policies toward adaptation, irrelevant-factor variation promotes invariance, and independent variation improves generalization over confounded variation at the same data budget.
Source: §5 · p. 10Takeaways & Limitations
Kernel-based SNR and factor sensitivity offer practical signals for assessing OOD robustness and shortcut learning alongside task-level evaluation.
Source: §5 · p. 10The framework requires probe sets around known factors and cannot discover unknown variation sources; kernel magnitudes are not directly comparable, so conclusions use qualitative trends and within-kernel comparisons.
Source: §5 · p. 10
Abstract
from arXiv · showhide
Training data variation, whether through designing a domain randomization (DR) scheme in simulation or curating demonstrations for imitation learning, is a primary lever for improving the robustness of robotic manipulation policies. Yet its underlying mechanisms remain poorly understood, and practitioners typically select randomization parameters through expensive trial and error. We investigate these mechanisms through a series of case studies, randomizing object size, color, and type as well as scene lighting and linguistic prompts across settings including pick-and-place RL in ManiSkill and fine-tuning of vision-language-action (VLA) models on LIBERO and RoboTwin. We examine both model behavior and internal representations, using the empirical neural tangent kernel (NTK) as our primary diagnostic tool. We show that the NTK distinguishes a shift in the internal learning mechanism from \textit{memorizing} different situations with insufficient variation (e.g.\ learning what to do for a large cube, and what to do for a small cube) to \textit{adapting} to the situation at hand with sufficient variation. An NTK-based signal-to-noise ratio also helps distinguish when policies have learned to \emph{ignore} task-irrelevant factors (e.g.\ treating blue and red cubes identically, instead of learning a blue sub-policy and a red sub-policy). We use these diagnostics to develop practical guidance for designing DR schemes, selecting models, and detecting shortcut learning. We further compare different kinds of representations and validate our findings with real-world hardware experiments using ACT-based imitation learning.
1 Introduction
Training-data variation can improve robot-policy robustness, but practitioners still choose what and how much to vary through intuition or costly trial and error. The paper introduces representation-kernel diagnostics to distinguish how policies learn from that variation.
1 Introduction
- Training-data variation is a primary lever for policy robustness, while decisions about which factors to vary, how much, and how to compose them remain largely intuition-driven or costly.Real-world rollouts, hardware data, and expert demonstrations are expensive or scarce.Source: §1 · p. 1
- Task success rates cannot distinguish fragile memorization, shortcut reliance, and adaptive generalization, although these regimes call for different practitioner interventions.This limits the ability to interpret and optimize policies before real-world deployment.Source: §1 · pp. 1–2
- The paper uses representation-kernel statistics, primarily the eNTK, to track internal representations and diagnose memorization, adaptation, or disregard of task-irrelevant variation.Last-hidden-layer and output kernels are more computationally efficient alternatives, but each has characteristic blind spots.Source: §1 · p. 2
2 Background and Related Work
Representation kernels capture similarities between inputs through chosen network features, while data-variation research has largely assessed generalization through outcomes. This paper uses kernel metrics to diagnose whether policies memorize, adapt to, or ignore variation.
2.1 Representation Kernels
- A representation kernel computes input similarity through a feature map; activations, outputs, and parameter gradients expose different aspects of a network.The eNTK reflects current representations and how training on one input propagates to another; unlike activation-based prior work on robustness, this paper diagnoses responses to training-data variation.Source: §2.1 · p. 2
- Restricting the eNTK Jacobian to selected LoRA modules sketches those layers’ kernel contributions more cheaply and can diagnose the existing representation without performing LoRA fine-tuning.Its tangent-feature similarities indicate how an update on one data point would affect another.Source: §2.1 · p. 2
2.2 Data Variation and Generalization in Robot Learning
- Domain randomization varies simulated environments, while imitation-learning variation comes from diverse demonstrations; prior work studies their relation to generalization across tasks and factors.Findings include that environment and object diversity can matter more than demonstration count, but not every kind of diversity helps.Source: §2.2 · p. 3
- Prior studies often measure generalization through success rates, showing whether policies generalize but not whether they memorize, adapt, or ignore a factor.This paper diagnoses those mechanisms with representation-kernel metrics.Source: §2.2 · p. 3
3 Methodology
The methodology varies selected training factors, constructs controlled probes, and applies kernel-based diagnostics. In the square-color study, increasing variation shifts the representation from treating colors as isolated cases toward adaptation, with OOD gains and growing relative sensitivity to task-relevant state.
3 Methodology
- The three-stage pipeline controls one or two variation factors, builds probes that vary a target factor while holding other inputs fixed, then computes representation kernels and extracts diagnostic statistics.Jointly randomized factors can be confounded or independent, with independent sampling covering all combinations.Source: §3 · p. 3
- In the square-size warm-up, models are trained with Ncolor values from 1 to 27 and evaluated using 27 color probes to inspect their kernel matrices.Ten distinct color subsets are sampled at each Ncolor level.Source: §3 · pp. 3–4
- With few training colors, kernel matrices are block-structured and treat colors as isolated cases; with more colors, this memorize regime gives way to an adapt regime.Effective rank first increases with more colors, then decreases as the model treats colors similarly.Source: §3 · p. 4
- OOD success surges around the point where kernel effective rank flattens, while in-distribution success is near-perfect across nearly all color-randomization levels.The evaluation averages results across 10 random color subsets.Source: §3 · p. 4
- PCA probes show that increasing color variation brings same-state, different-color representations closer while preserving state-dependent structure.The comparison uses models trained with Ncolor = 4 and Ncolor = 27.Source: §3 · p. 4
- SNR compares representation distances for task-relevant versus task-irrelevant factors, while factor-specific CKA measures alignment between the kernel and a factor-label kernel.In the color experiment, SNR rises with Ncolor, indicating greater relative sensitivity to square size than color; state and color CKA provide a factor-specific view of the transition.Source: §3 · p. 5
4 Experimental Analyses
Across robot-learning settings, representation-kernel diagnostics distinguish memorization from adaptation and identify when policies suppress task-irrelevant visual variation. The experiments also show that factor-conditioned signals track out-of-distribution robustness and expose complementary blind spots across kernel representations.
4 Experimental Analyses
- Insufficient domain randomization produces block-diagonal kernels associated with memorization, while sufficient randomization smooths the structure toward adaptation.The random_color setting is given as an example of the adapt-style regime.Source: §4 · pp. 5–6
- Sufficient background and object randomization yields an ignore regime that collapses irrelevant visual variance and leads to OOD policy robustness.Source: §4 · p. 6
- For ACT peg insertion, no-DR training has the lowest validation loss but worst OOD success, while SNRs track OOD success and joint-probe SNR orders all four conditions by joint OOD success.Single-factor variation mainly improves generalization to that factor; joint variation is needed for joint shifts.Source: §4 · p. 6, Table 1 · p. 6
- For OpenVLA, OpenVLA-OFT, and RIPT-VLA, state-to-lighting SNR ordering matches empirical OOD lighting robustness on LIBERO-plus.Source: §4 · p. 6, Figure 3 · p. 7
- After a third RoboTwin task breaks the visual-language correlation, scene-to-language FSR decreases, indicating reduced relative scene sensitivity; matched-step controls rule out longer training alone.The decrease occurs at all 10 step-matched checkpoints relative to a control that remains at 15.0.Source: §4 · p. 7, Figure 4 · p. 7
- At equal data budget, independent object-texture variation lowers action error on novel categories, while the conditions are indistinguishable in-distribution.Source: §4 · p. 8, Figure 5 · p. 8
- Kernel representations have complementary blind spots: output kernels miss shortcut reliance when instructions imply the same action, while last-hidden kernels miss arm-specific structure introduced in final heads.The eNTK captures the bimanual decoupling trend more weakly than the output kernel, whose SNR rises from about 2 to 48 for both arms.Source: §4 · pp. 8–9, Table 2 · p. 9
- Lightness DR raises hardware training SNR from 2.5 to 5.5; bright-light success rises from 1/10 to 8/10 and dark-light success from 1/10 to 6/10.In-distribution success rises from 6/10 to 10/10.Source: §4 · p. 9, Figure 6 · p. 9
5 Conclusion
The framework uses representation kernels to distinguish memorization, adaptation, and ignoring task-irrelevant variation. Training-data design shapes these regimes, though the framework requires probes around known factors and supports only within-kernel comparisons.
5 Conclusion
- Kernel-based SNR and factor sensitivity distinguish memorizing conditions, adapting across relevant variation, and ignoring irrelevant variation without relying solely on task-level evaluation.The framework is presented as a practical diagnostic for data design and model analysis.Source: §5 · p. 10
- Increasing relevant variation can shift policies from memorization toward adaptation, while irrelevant-factor variation promotes invariance across model classes and hardware settings.Independent variation also improves generalization over confounded variation at the same data budget.Source: §5 · p. 10
- The framework requires probe sets around known factors, so it cannot discover unknown variation; kernel metric magnitudes also cannot be compared directly across kernels.Conclusions therefore rely on qualitative trends and within-kernel comparisons.Source: §5 · p. 10
A.1 Hardware Evaluation Details
The hardware evaluation tests an ACT imitation-learning policy on cube pick-and-place under training-matched and out-of-distribution lighting. Lightness domain randomization improves performance in all tested conditions and substantially raises success under both lighting shifts.
A.1 Hardware Evaluation Details
- The evaluation used an ACT imitation-learning policy on a hardware cube pick-and-place task across Base, Bright, and Dark lighting conditions.Base matched training; Bright and Dark introduced out-of-distribution illumination changes.Source: §A.1 · p. 13, Figure 7 · p. 13
- With Lightness DR, success was 10/10 in Base, 8/10 in Bright, and 6/10 in Dark, versus 6/10 in Base and 1/10 in each OOD condition without it.Each condition was evaluated over 10 trials; Base matches training, while Bright and Dark are OOD lighting variations.Source: §A.1 · p. 13, Figure 7 · p. 13
A.2 eNTK vs. Activation Similarity as a Predictor of Transfer
The comparison tests whether input similarity predicts how training on one input changes the output on another. Across checkpoints, the eNTK closely tracks this transfer, while activation similarity remains weak even for the best post-hoc-selected layer.
A.2 eNTK vs. Activation Similarity as a Predictor of Transfer
- The experiment measures output changes on held-out probes after a single SGD step on a probe input, then correlates those changes with candidate similarity measures.It uses approximately 1.7k input-pair comparisons per checkpoint.Source: §A.2 · p. 14
- Activation similarity is cosine similarity between inputs’ representations at each network layer, with the best-correlating layer selected post hoc.The MLP head is also reported separately.Source: §A.2 · p. 14
- The results support eNTK as a predictor of cross-input training effects because gradient descent transfer depends on parameter-gradient alignment, not representation alignment.Similar representations can still be weakly coupled through shared parameters.Source: §A.2 · p. 14
- ρ = 0.97 ± 0.06 for eNTK, compared with ρ = 0.22 ± 0.28 for the best post-hoc activation layer and ρ = 0.09 ± 0.52 for the MLP head.These Spearman correlations compare similarity with measured output change after one SGD step, across 27 checkpoints.Source: Table 4 · p. 14, §A.2 · p. 14
Related papers
This paper cites
- OpenVLA: An Open-Source Vision-Language-Action Model2024
- Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World2017
- Neural Tangent Kernel: Convergence and Generalization in Neural Networks2018
- What Matters in Learning from Offline Human Demonstrations for Robot Manipulation2021
- Solving Rubik's Cube with a Robot Hand2019
- On Lazy Training in Differentiable Programming2018