Context-Smoothed Pre-training (CSP)

• CSP is a pre-training technique that injects forward-diffusion noise into policy inputs to broaden action coverage and enable a smooth shift from imitation to exploration.
• It replaces standard behavior cloning with learning over corrupted contexts and an explicit noise-level variable to control the balance between faithful imitation and exploratory behavior.
• The method mitigates support collapse and improves downstream reinforcement learning by expanding the action distribution, enhancing sample efficiency and out-of-distribution generalization.

Context-Smoothed Pre-training (CSP) is a pre-training procedure for robot policies that broadens action coverage by injecting forward-diffusion noise into policy inputs during supervised imitation learning. Introduced as the pre-training component of the TMRL framework, CSP replaces standard behavior cloning on clean context $c$ with training over corrupted contexts $\tilde c$ and an explicit noise-level variable, thereby creating a controllable continuum between faithful imitation and exploratory behavior. Its stated purpose is to bridge behavior-cloning pre-training and reinforcement-learning fine-tuning, especially in settings where narrow conditional action distributions impede downstream exploration (Hong et al., 12 May 2026).

1. Conceptual basis and motivation

The motivation for CSP is the observation that ordinary behavior cloning often has collapsed support. When a policy is trained only on clean demonstration context, it learns a narrow conditional distribution $p(a\mid c)$ that can perform well on-distribution yet assign near-zero probability to alternative actions in unfamiliar or sparse regions of context space. In downstream reinforcement learning, that narrowness becomes a bottleneck: if the initial policy never proposes useful actions in out-of-distribution states, online optimization receives little or no reward signal (Hong et al., 12 May 2026).

CSP addresses this problem by corrupting the context before action prediction. Rather than fitting only the exact demonstrated conditioning, it trains the policy over a family of corrupted contexts and conditions the model on the amount of corruption. The paper describes the induced behavior as a continuum between exploitation and exploration: at low noise, the policy recovers faithful imitation; at high noise, the context is blurred enough that the action distribution moves toward a broader marginal-like distribution $p(a)$. The key intuition is context aliasing: when nearby contexts are made partially indistinguishable by noise, the policy can borrow coherent actions from neighboring contexts and thereby expand coverage without resorting to unstructured Gaussian action noise (Hong et al., 12 May 2026).

This distinction from direct action-space perturbation is central. The framework does not inject randomness into actions after a sharply trained policy has been learned. Instead, it alters the conditioning structure during pre-training so that broader action support is itself learned as a function of context ambiguity. The paper therefore characterizes CSP not as ordinary regularized behavior cloning, but as pre-training a steerable policy over a noise axis (Hong et al., 12 May 2026).

2. Mathematical formulation

Standard behavior cloning is written as

$$\min_\theta \; \mathbb{E}_{(c,a)\sim D}\big[\ell(\theta; c,a)\big],$$

and in the simplest case

$\ell(\theta;c,a) = -\log p_\theta(a\mid c).$

CSP replaces clean conditioning with corrupted conditioning. If $\tilde c$ is sampled from a corruption kernel $q_o(\cdot\mid c)$ at noise level $o$, the training objective becomes

$$\min_\theta \; \mathbb{E}_{(c,a)\sim D}\;\mathbb{E}_{\tilde c \sim q_o(\cdot\mid c)}\big[\ell(\theta; a,\tilde c,o)\big].$$

The corresponding context-smoothed policy is

$\tilde c$0

In the concrete implementation, the corruption process is a forward diffusion over context. Let $\tilde c$1, with variance schedule $\tilde c$2,

$\tilde c$3

The closed-form forward corruption distribution is

$\tilde c$4

CSP then trains a policy $\tilde c$5 over all noise levels,

$\tilde c$6

The algorithmic structure remains close to imitation learning. The paper summarizes it as: sample a demonstration pair $\tilde c$7; sample a noise level $\tilde c$8; sample corrupted context $\tilde c$9; and update $p(a\mid c)$0 to minimize $p(a\mid c)$1. For generative control policies, $p(a\mid c)$2 is typically a denoising score-matching or flow-matching objective rather than plain log-likelihood (Hong et al., 12 May 2026).

3. Action coverage, overlap, and theoretical interpretation

The paper’s theoretical motivation is framed through Gaussian context smoothing,

$p(a\mid c)$3

Within this model, smoothing increases overlap between policies at different contexts. Theorem 1 states

$p(a\mid c)$4

or equivalently

$p(a\mid c)$5

The interpretation given in the paper is precise. Larger $p(a\mid c)$6 produces more overlap and therefore more coverage; closer contexts overlap more than distant ones; and the induced expansion is structured rather than random. A corollary further states that, under a local identifiability condition, sufficient smoothing increases overlap relative to the unsmoothed policy and mitigates support collapse from behavior cloning (Hong et al., 12 May 2026).

This theoretical picture explains why CSP is meant to improve reinforcement-learning fine-tuning. The policy does not merely become noisier. It becomes more likely to propose coherent action sequences that were observed in related contexts, which is a different exploration mechanism from temporally incoherent perturbation in action space. The paper explicitly contrasts this with Gaussian action noise, arguing that direct action noise often causes dithering and temporal incoherence, whereas CSP broadens support by borrowing behaviors from neighboring demonstrated states (Hong et al., 12 May 2026).

4. Role within TMRL and relation to standard pre-training

CSP is the pre-training stage of a two-stage framework. TMRL, or Timestep-Modulated Reinforcement Learning, is the subsequent RL fine-tuning method that exploits the steerability learned by CSP. The distinction is functional rather than cosmetic: CSP determines which behaviors are available at different noise levels, whereas TMRL learns when to select which noise level to maximize reward (Hong et al., 12 May 2026).

Component	Regime	Function
BC	Offline supervised	Learns narrow $p(a\mid c)$7
CSP	Offline supervised	Learns $p(a\mid c)$8 over a noise axis
TMRL	Online RL	Chooses $p(a\mid c)$9 and $p(a)$0 during execution

In TMRL, a high-level actor $p(a)$1 outputs both latent steering noise and smoothing level: $p(a)$2 The base CSP policy is frozen during this stage. In some experiments the framework uses timestep-aware critics $p(a)$3, and the RL optimizer is SAC or RLPD depending on the setting (Hong et al., 12 May 2026).

The paper also emphasizes that CSP is input-agnostic. The context $p(a)$4 can be state and goal information in state-based policies, a point cloud plus proprioceptive state in point-cloud manipulation, or a VLM embedding feeding the action expert in image-based VLA policies. In the reported OGBench implementation, the context-smoothed policy uses a Diffusion Transformer backbone adapted from prior work; it conditions on the noised context, action diffusion timestep, context timestep $p(a)$5, and any non-noised observations such as proprioception, with the conditioning injected through Adaptive LayerNorm. The corruption schedule is a discrete linear beta schedule with $p(a)$6, $p(a)$7, and $p(a)$8, while training samples a continuous $p(a)$9 and maps it to a discrete diffusion timestep for corruption (Hong et al., 12 May 2026).

5. Empirical behavior and observed advantages

The empirical claims tied specifically to CSP concern coverage and out-of-distribution generalization. On OGBench tasks, CSP achieves higher Success@K than standard behavior cloning and PostBC. The paper highlights the cube environment as an extreme case: behavior cloning and PostBC have essentially zero success across $$\min_\theta \; \mathbb{E}_{(c,a)\sim D}\big[\ell(\theta; c,a)\big],$$0, while CSP still succeeds. This is used as a proxy for action coverage: if at least one of $$\min_\theta \; \mathbb{E}_{(c,a)\sim D}\big[\ell(\theta; c,a)\big],$$1 base rollouts succeeds from an out-of-distribution state, the base policy is covering relevant demonstrator behaviors (Hong et al., 12 May 2026).

The paper also reports that CSP trained with noisy contexts can better approximate unseen conditioning values than standard conditional diffusion or behavior cloning. As the context timestep increases, the induced distributions smoothly interpolate from conditional to marginal behavior. This interpolative behavior is central to the framework’s interpretation: CSP is meant to preserve imitation fidelity at one end of the noise axis while enabling broader, reward-discovering rollouts at the other end (Hong et al., 12 May 2026).

A further comparison is made against a classifier-free-guidance baseline, denoted TMRL-CFG. That baseline interpolates between conditional and unconditional models using a coefficient $$\min_\theta \; \mathbb{E}_{(c,a)\sim D}\big[\ell(\theta; c,a)\big],$$2, but the paper reports it is worse than CSP+TMRL because CFG still relies on a conditional model that may already be out-of-support on out-of-distribution contexts. By contrast, CSP explicitly brings such contexts back toward the training manifold through context corruption (Hong et al., 12 May 2026).

At the system level, the combined CSP+TMRL framework is presented as improving RL fine-tuning sample efficiency across arbitrary policy inputs, including states, 3D point clouds, and image-based VLA policies. The abstract further states that TMRL enables successful real-world fine-tuning on complex manipulation tasks in under one hour, which situates CSP as the pre-training mechanism that makes such adaptive exploration feasible (Hong et al., 12 May 2026).

6. Terminological scope and related context-aware pre-training lines

The acronym “CSP” is not unique across machine learning, and this is important for terminological precision. In robot learning, “Context-Smoothed Pre-training” denotes the diffusion-based context-corruption method described above (Hong et al., 12 May 2026). In 3D point cloud learning, however, “CSP” in “PointCSP” stands for “cross-sample semantic propagation,” a different framework in which samples within a batch are serialized into continuous input and processed by a state-space model to enforce semantic consistency across scenes (Yu et al., 3 May 2026). The shared acronym should therefore not be taken to imply a shared algorithm.

Several adjacent papers nevertheless reinforce a broader context-governance perspective on pre-training. In speech self-supervision, “Bigger is not Always Better: The Effect of Context Size on Speech Pre-Training” shows that phoneme discriminability peaks at around $$\min_\theta \; \mathbb{E}_{(c,a)\sim D}\big[\ell(\theta; c,a)\big],$$3 ms of preceding context and degrades substantially beyond around $$\min_\theta \; \mathbb{E}_{(c,a)\sim D}\big[\ell(\theta; c,a)\big],$$4 ms, indicating that wider context can be harmful when the target representation is local and phonetic (Robertson et al., 2023). In dense retrieval, “Query-as-context Pre-training for Dense Passage Retrieval” replaces same-document passage pairs with passage-query pairs, arguing that context relevance matters more than raw context proximity and improving both in-domain retrieval and BEIR zero-shot performance (Wu et al., 2022).

Other papers are explicitly described as not using CSP by name while still being “CSP-like” in spirit. “A Framework for Quantifying How Pre-Training and Context Benefit In-Context Learning” formalizes how pre-training and constructed context jointly shift predictions toward the query-task distribution, with performance governed by context length and KL divergence (Song et al., 26 Oct 2025). “MEG-XL: Data-Efficient Brain-to-Text via Long-Context Pre-Training” argues the opposite of the short-context speech result in its own domain: pre-training with $$\min_\theta \; \mathbb{E}_{(c,a)\sim D}\big[\ell(\theta; c,a)\big],$$5 minutes of MEG context per sample improves low-data word decoding, and the gains from longer pre-training context show diminishing returns only after about $$\min_\theta \; \mathbb{E}_{(c,a)\sim D}\big[\ell(\theta; c,a)\big],$$6 seconds (Jayalath et al., 2 Feb 2026). This contrast suggests that context is not a universally monotone resource. Its effective use depends on modality, timescale, and downstream objective.

A common misconception is therefore that CSP or CSP-like pre-training simply means “use more context.” The broader literature summarized in these papers does not support that simplification. In some settings, such as the robot-policy case, smoothing context creates a controllable exploration–exploitation continuum; in others, such as speech, limiting context improves representation quality; in still others, such as MEG, long contiguous context is beneficial. What unifies these directions is not a single context length or corruption rule, but the treatment of context as a first-order design variable in pre-training rather than a fixed background assumption (Hong et al., 12 May 2026).