PaperlayerView exact revision on arXiv

Source-linked AI summary

Action Shaping: Policies Absorb What They Can Express

Yanjun Chen, Jinghan Wang, Xiaoyu Shen, Wenjie Li, Wei Zhang

arXiv:2609.32752v2cs.AIcs.LG

TL;DR

Action shaping adds a training-time action offset but lacks reward shaping’s guarantee that the offset can be removed without changing the optimal policy. The paper tests when a trainable policy absorbs such an offset, and finds absorption when the base head can reproduce it exactly; the offset’s amplitude indicates what removal will cost.

  • Problem

    Unlike reward shaping, action shaping has no theorem guaranteeing that a training-time action offset can be removed without changing the policy, so a co-trained correction may need to remain at deployment.

    Source: §1 · p. 1

  • Method

    The paper studies a zero-initialized, scalar-gated action head co-trained with an actor and defines absorption by exact reproduction through the trainable base head.

    Source: §1 · p. 1, §2.3 · p. 4

  • Results

    On 20 tasks, removing the head costs almost nothing when the offset is absorbed; exact reproduction, rather than capacity, distinguishes absorbed from unabsorbed heads.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    The offset’s amplitude during training indicates the cost of dropping the head, providing a check before removal.

    Source: §1 · p. 2

    Expressibility is relative to the base head’s function class on the shared representation; a trainable trunk’s capacity to approximate an offset is not sufficient.

    Source: §2.3 · p. 4

Abstract

from arXiv · show

Reward shaping has a theorem: a potential-based term can be removed without changing the optimal policy. The same practice on the action channel, an offset added in training and dropped at deployment, has no theorem. Nothing cancels an action offset, so the correction is kept at deployment or removed without a guarantee. We call it action shaping and state its principle. A trainable policy absorbs an offset its own output layer can reproduce exactly, which is what we mean by express; what is absorbed can be removed with the return intact. Its minimal instance is a zero-initialized linear head behind a learnable gate, added to an actor that trains through a learned action-value function, with no penalty or schedule. The gate rises and then falls on its own, for deterministic and stochastic actors alike, and on 20 tasks removing the head costs almost nothing. The condition is exact reproduction, not capacity: a nonlinear head with more parameters is not absorbed, and in a paired control, one linear path added to a nonlinear base head restores absorption. Exact reproduction gives the loss a flat direction that gradient noise drifts along, and the offset's amplitude indicates, before removal, what dropping the head will cost. Action shaping thus gains the counterpart of the shaping theorem, a condition for absorption, together with the mechanism behind it and a diagnostic that reads it. Policies absorb what they can express, and only that.

1 INTRODUCTION

Action shaping adds a training-time action offset without reward shaping’s invariance guarantee. The paper finds that a trainable base absorbs the offset when its output layer can reproduce it exactly, and uses the offset’s amplitude to indicate removal cost.

1 INTRODUCTION

  • A zero-initialized, scalar-gated linear head is absorbed by a co-trained TD3 or SAC actor: its gate rises, then falls without a schedule or penalty, and removing it costs almost nothing.The head adds a second linear layer on the actor’s existing representation.Source: §1 · p. 1
  • Absorption requires exact reproduction by the base head’s own weights, not capacity: a nonlinear head is not absorbed, while a linear path restores absorption on a nonlinear base head.A frozen base or a head the base cannot reproduce is not absorbed.Source: §1 · p. 2
  • Exact reproduction creates a flat loss direction, and gradient noise moves parameters along it toward a closed gate; under label-noise SGD, the predicted gate decay is exponential.The head’s weights decay more slowly than the gate, by a factor of the head’s size.Source: §1 · p. 2
  • The offset’s amplitude predicts removal cost: across environments, runs retaining less than a tenth of peak amplitude lose nothing in interquartile mean when the head is removed, while others lose 0.03 normalized return.Within a run, removal is costly near the amplitude peak and small at either end.Source: §1 · p. 2
  • TD3 gains 0.077 in interquartile mean over 20 environments, and the shaped SAC actor with an ensemble critic exceeds its host at every tested update ratio.On 26 Atari games, the actor deployed without the head also outscores the unshaped host.Source: §1 · pp. 2–3

2 ACTION SHAPING AND THE AMPLITUDE DIAGNOSTIC

The section specifies a gated action-shaping head, measures its offset amplitude and removal cost, and defines expressibility through exact reproduction by the base head. Its experimental arms vary trainability, gradient flow, and the functions available to each head.

2.1 THE SHAPED ACTOR

  • The shaped actor adds a second head on the trunk representation and a learnable scalar gate; actor-loss gradients update both heads, the trunk, and gate, with no penalty or schedule, and deployment drops the shaping head.The head is zero-initialized, with a small random gate so its gradient is nonzero; the offset is δ and its amplitude is ∥δ∥.Source: §2.1 · p. 3

2.2 WHAT WE MEASURE

  • The study tracks mean offset amplitude and the return difference between shaped and base-only actors; the terminal residual is the offset’s retained fraction of peak amplitude over the last 20 evaluations.A peak amplitude never above 10^-3 marks a non-engaged run; r ≤ 0.10 marks completion, while rank correlations remain 0.565–0.671 under four residual conventions.Source: §2.2 · p. 3
  • The gate and the offset it scales distinguish a branch that was absorbed, one that remains active, and one that never engaged; head-weight norms are not the diagnostic.A gate returning to zero indicates the base alone computes what the shaped actor computed.Source: §2.2 · p. 3

2.3 SEVEN ARMS

  • Seven experimental arms vary representation sharing, base trainability, and head function classes; expressibility means the base head can reproduce the shaped output exactly on the shared representation.A linear offset is expressible for a linear base head, but not for an MLP head or a linear offset on an MLP base head unless it has its own linear path; the linear head’s absorbed state is known in closed form.Source: §2.3 · p. 4

3 WHY THE OFFSET HELPS AND WHY IT SHOULD DECAY

An action offset helps to first order when it aligns on average with the critic’s action gradient. When the base can reproduce the offset exactly, a flat direction lets gradient noise drive absorption; inexpressible offsets have no such path to removal.

3.1 FIRST-ORDER IMPROVEMENT AND THE CRITIC

  • An offset helps to first order exactly when it aligns on average with the critic’s action gradient.For a deterministic policy, this term is the deterministic policy gradient with respect to the offset scale at α = 0.Source: §3.1 · p. 4
  • The offset’s first-order effect fades as the base policy approaches optimality, explaining why the gate stops rising and later decreases.A Q-critic supplies the actor’s gate gradient as a pathwise gradient; under PPO, it is carried only through a score-function estimate over sampled actions.Source: §3.1 · p. 4

3.2 EXPRESSIBILITY CREATES A FLAT LINE

  • For a linear head, shifting weight between the base and shaping heads leaves the loss unchanged, forming a flat line whose flattest split has α + λ = 0.Along this line, the trace of the Hessian with respect to the shaping-head weights scales with (α + λ)^2.Source: §3.2 · p. 5
  • Gradient noise can drift along that flat line toward α = 0; in the stated fixed-trunk, squared-error setting, both the gate and head tend to zero.The projected trace-gradient dynamics make the gate decay exponentially whenever the head’s output power is below m.Source: §3.2 · p. 5
  • The gate can close while most head output remains, so residual head norm alone does not show absorption is incomplete.In the paper’s idealized dynamics, the head decays per covariance eigen-direction, more slowly than the gate on the gate-closing timescale.Source: §3.2 · p. 5
  • If the base cannot reproduce the offset exactly, no flat direction leads to α = 0; freezing the base also prevents the split from changing.Detaching the shaping head’s gradient from the trunk does not change whether the gate returns to zero.Source: §3.2 · p. 5

4 EXPERIMENTS

Across continuous-control and Atari evaluations, action shaping can improve returns while often leaving little removal cost when the base absorbs the head. The experiments show that absorption tracks the amplitude trajectory and requires an expressible offset and a trainable base.

4 EXPERIMENTS

  • The continuous-control evaluation trains three shaped algorithms for 1M steps across 20 environments, while mechanism experiments use six environments under TD3 and SAC.Scores are normalized per environment and aggregated by IQM; mechanism arms use five seeds.Source: §4 · p. 6

4.1 WHAT THE HEAD ADDS AND WHAT THE BASE RETAINS

  • TD3’s IQM rises by +0.077 [0.040, 0.116] over 20 environments, while removal costs are +0.006 [0.002, 0.012] under TD3 and +0.001 [−0.001, 0.003] under SAC.The head also raises mean return in 14 of 20 environments; intervals for per-environment removal cost contain zero in 15 of 20.Source: §4.1 · p. 6

4.2 THE ARC: RISE, THEN A DECLINE THAT DEPENDS ON THE CRITIC

  • In all 200 TD3 and SAC runs across 20 environments, the head engages; on the mechanism suite its amplitude rises and then falls, with SAC ending at or below 11% of peak in every environment.The arc completes within 1M steps in 7 environments under TD3 and 13 under SAC; a ten-member ensemble critic closes Hopper’s TD3 gate, reducing its terminal-to-peak ratio from 21% to 0.2%.Source: §4.2 · pp. 6–7, Figure 3 · p. 7

4.3 REMOVAL COST ALONG THE AMPLITUDE TRAJECTORY

  • When the amplitude arc completes, removal costs are near zero; when it does not, IQM costs are +0.028 under TD3 and +0.030 under SAC.Across environments, removal cost correlates with residual at 0.61 under TD3 and 0.64 under SAC; within runs, cost is largest at or just after the amplitude peak and indistinguishable from zero at the end.Source: §4.3 · p. 7, Table 1 · p. 7

4.4 WHAT ABSORPTION REQUIRES: SEVEN ARMS

  • Across twelve cells, the shared arm has residual 0.10 and removal cost +0.003; detaching head-to-trunk gradients leaves residual 0.06 and cost +0.0004.A separate head trunk raises residual to 0.55 under TD3 and 0.68 under SAC; freezing the base leaves residuals of 0.68 and 0.77, respectively.Source: §4.4 · p. 8, Figure 4 · p. 8
  • A one-hidden-layer shaping head remains unabsorbed at widths 16, 64, and 256, with removal costs between +0.44 and +0.78 and no completed arcs in 36 cells.The shared arm completes eight of twelve cells, showing that greater width alone does not restore absorption.Source: §4.4 · p. 8
  • On a nonlinear base head, adding one zero-initialized linear path lowers median residual from 0.79 to 0.012 and brings removal cost to −0.0002 [−0.0007, 0.0004].Without that path, the head’s residual is 0.60 under TD3 and 0.82 under SAC, with removal costs of +0.38 and +0.35.Source: §4.4 · p. 8

4.5 DISCRETE ACTIONS

  • On 26 Atari games, removing the discrete-action head raises human-normalized IQM by 0.013 [0.003, 0.023] over the shaped actor.The base actor alone scores 0.183 [0.157, 0.216], versus 0.170 for the shaped actor and 0.154 for the unshaped host.Source: §4.5 · p. 8

5 RELATED WORK

The paper relates action shaping to reward-shaping invariance, branch folding, gated networks, training-only components, and internalization. Its distinction is that an expressible branch on a trainable base can be absorbed without a designer-set schedule.

Reward shaping and the action channel

  • Unlike potential-based reward shaping, action offsets change the state distribution, so their removability depends on training dynamics rather than an invariance theorem.The paper frames its contribution as a condition for absorption and its mechanism.Source: §5 · p. 9

Folding a branch into a layer

  • Structural re-parameterization folds trained branches into a layer, but this paper finds training reaches that merge only when exact reproduction is possible (Ding et al., 2021; Hu et al., 2022).A base head that only approximates a branch retains it; adding a linear path that restores the merge restores absorption.Source: §5 · p. 9

Scalar-gated branches

  • Unlike ReZero and LayerScale gates that grow and stay open, the paper’s gate closes without penalty after opening (Bachlechner et al., 2021; Touvron et al., 2021).The paper distinguishes this sequence from gates that close through dead gradients and from explicit sparsity penalties.Source: §5 · p. 9

Training-only components

  • Prior training-only components are generally removed by designer-set schedules, whereas an expressible branch on a trainable base can drift toward absorption without one.Residual policy learning instead keeps its correction on a fixed base (Silver et al., 2018; Johannink et al., 2019).Source: §5 · p. 9

Internalization and removability

  • Compared with internalization defined as absorbing a procedure into the same hypothesis class, this paper adds a trainable-base condition, a flat-line mechanism, and an amplitude diagnostic (Tsilivis et al., 2026).Checkpoint interventions also show removability depends on the training trajectory, as concurrent work finds for an annealable prior (Yang et al., 2026).Source: §5 · p. 9

6 DISCUSSION

Removal deploys the trained base actor with the head deleted, while absorption means training closes the distance to the linear head’s closed-form merge without performing that merge. The discussion frames expressibility and trainability as the conditions for predicting whether a branch will be absorbed.

6 DISCUSSION

  • For the linear head, removal uses the base actor as trained, while the fold Wµ + αWh is a reference point for measuring absorption rather than a merge performed at deployment.A branch can be removable without absorption if it never engages or collapses under strong noise; where no fold exists, the amplitude reading is the applicable one.Source: §6 · p. 9
  • At three million steps, the width-64 MLP head remains unabsorbed, with median residual 0.73 and removal cost +0.91.The paired linear control separates inexpressibility from insufficient training budget.Source: §6 · p. 9
  • In five of six agent pairs, batch size 64 completes the arc before batch size 1024, matching Proposition 3’s predicted direction.The agents use Adam and a bootstrapped critic.Source: §6 · p. 10
  • The condition reduces removability to whether the trainable base head can reproduce the branch exactly; amplitude below one-tenth of peak is the proposed removal cue.An inexpressible branch or a branch on a fixed base stays; the discussion also notes the amplitude should be read in action space when the output map saturates.Source: §6 · p. 10

7 CONCLUSION

The conclusion presents exact reproduction by a trainable base head as the condition for absorbing an action-channel branch. It connects that condition to a flat loss direction, a gate-based diagnostic, and low removal costs across evaluated tasks.

7 CONCLUSION

  • Absorption requires exact reproduction by a trainable base head, not merely greater capacity; adding one linear path to a nonlinear base restores absorption.The paired control supports the distinction between exact reproduction and parameter count.Source: §7 · p. 10
  • Expressibility creates a flat direction in the loss, along which gradient noise moves parameters and closes the gate without a schedule.For deterministic and stochastic actors, the gate rises and then falls on its own.Source: §7 · p. 10
  • Across 20 tasks, removing the absorbed head costs almost nothing; with default hyperparameters and no per-task tuning, TD3’s IQM rises by 0.077.On 26 Atari games, the actor deployed without the head outscores its unshaped host.Source: §7 · p. 10
  • The condition is proposed for co-trained action-channel branches, including branches from controllers, priors, or residual policies, and extends to discrete actions through the softmax case.The paper’s arms show absorption only for branches the base head can reproduce.Source: §7 · p. 10

A PROOFS

The proofs relate the return change from an action offset to a first-order action-value term, then establish when that signal vanishes and how gradient noise moves parameters along an equivalent-loss line.

A PROOFS

  • The performance-difference argument and Taylor expansion give a first-order return change proportional to the action-value gradient projected onto the offset, with an O(α^2) remainder.For smooth output maps, the offset is transformed by the output-map Jacobian; for discrete softmax policies, the state-distribution shift is also O(α), leaving a second-order cross-term.Source: A PROOFS · p. 15

A.1 SELF-DIMINISHING COROLLARY

  • At an optimal base policy, the shaping signal vanishes for deterministic and softmax policies, and the full SAC actor-loss gradient with respect to the head vanishes at the soft optimum.As the base approaches optimality, the signal diminishes under the stated continuity assumptions.Source: §A.1 · p. 16

A.2 FLAT LINE

  • With an expressible offset, changing the base and shaping-head weights along a line preserves their merged readout, shaped action, and loss.The proof also shows the Hessian trace is least on that line when α + λ = 0.Source: §A.2 · pp. 16–17

A.3 DRIFT ALONG THE LINE

  • Under label-noise stochastic gradient descent on the zero-loss manifold, the projected flow drives an expressible linear head and gate to zero, with the gate turning over when output power is high.The gate moves before the head shrinks; an inexpressible head that the target needs and a frozen base have no zero-loss point with α = 0.Source: §A.3 · pp. 17–18

B MINIMAL MODEL OF ABSORPTION

The minimal model tests whether noise drives absorption in a gated linear shaping head and whether absorption depends on exact expressibility. The results distinguish noise levels, base-head configurations, and optimizer behavior.

Model

  • The minimal-model runs use a two-hidden-layer trainable trunk, linear base and shaping heads, an exact fixed critic, and three noise conditions.Each configuration runs 100K steps with three seeds; Table 2 reports terminal residual, completion count, and retained head norm.Source: B MINIMAL MODEL OF ABSORPTION · p. 18, Table 2 · p. 18

What the model shows

  • With target noise, the shared linear head returns to zero in every seed by about 58K steps, while minibatch noise leaves it incomplete and no noise produces no return.At N2 the residual is 0.0014; at N1 it is 0.70 and still falling at the 100K-step budget, with 96% of head norm retained.Source: B MINIMAL MODEL OF ABSORPTION · p. 18, Figure 5 · p. 20
  • The linear head on an MLP base remains unabsorbed at every noise level, while a linear path restoring exact reproduction enables absorption in the paired control.The MLP head on a linear base can instead collapse under target noise, a different mechanism from absorption along an expressible flat line.Source: B MINIMAL MODEL OF ABSORPTION · pp. 18–19

The optimizer

  • Plain SGD leaves the zero-initialized head disengaged, with peak amplitude below 6 × 10^-4, whereas Adam produces peaks of 0.15 to 0.62.The account attributes the difference to Adam’s per-parameter normalization converting small consistent gradients into learning-rate-sized steps.Source: B MINIMAL MODEL OF ABSORPTION · p. 19

Drift check

  • In a longer noisy drift test, the gate decays exponentially while the head retains 92% to 96% of its norm after three million steps.The measured gate-decay rates are −2.8 × 10^-6 and −6.2 × 10^-6 per step at σ = 0.3 and σ = 0.45, respectively.Source: B MINIMAL MODEL OF ABSORPTION · p. 19

C A NONLINEAR BASE HEAD, WITH AND WITHOUT A LINEAR PATH

Paired controls test whether a linear path lets a nonlinear base head reproduce a linear offset exactly, and whether a nonlinear shaping head remains unabsorbed with a longer training budget.

C A NONLINEAR BASE HEAD, WITH AND WITHOUT A LINEAR PATH

  • The paired arms use the same function at initialization and differ in whether the base head can exactly reproduce the shaped output.The added linear path permits an exact parameter change; a one-hidden-layer base head without that path cannot reproduce the affine addition exactly.Source: C A NONLINEAR BASE HEAD, WITH AND… · p. 20
  • With a linear path, the median residual across twelve paired cells is 0.012 and nine cells complete; without it, the median is 0.786 and none complete.Removal cost is −0.0002 [−0.0007, 0.0004] with the path and +0.296 [0.240, 0.360] without it.Source: C A NONLINEAR BASE HEAD, WITH AND… · p. 20, Table 3 · p. 21

D THE MLP SHAPING HEAD AT THREE MILLION STEPS

  • At three million steps, the width-64 MLP shaping head has median residual 0.73 across twelve environment-algorithm cells, completes in none, and has removal cost +0.91 [0.86, 0.95].The median residual falls by only 0.02 from one to three million steps within the same runs.Source: D THE MLP SHAPING HEAD AT THREE… · p. 21, Table 4 · p. 21

E THREE MILLION STEPS FOR THE UNFINISHED TD3 ARCS

Longer runs lower the residual in all six tested unfinished TD3 arcs, while a multiplicative gain head remains largely unabsorbed across most tested cells. The batch-size experiment also finds earlier gate closure at smaller batch size in five of six pairs.

E THREE MILLION STEPS FOR THE UNFINISHED TD3 ARCS

  • At three million steps, residuals fell in all six extended TD3 environments, but removing the head still cost +0.49 normalized return on InvertedDoublePendulum.A residual of 0.056 there still corresponded to a mean action-space offset of 0.18 after squashing.Source: E THREE MILLION STEPS FOR THE… · p. 22, Table 5 · p. 22

F BATCH SIZE AND THE CLOSING TIME OF THE GATE

G THE SAME HEAD COMPOSED MULTIPLICATIVELY

  • The multiplicative gain head’s residual exceeded the shared-arm control in 11 of 12 cells, and removal cost exceeded control by more than 0.1 in four.Its state-dependent gain makes the offset inexpressible by the linear base head, so the gate generally stays open.Source: G THE SAME HEAD COMPOSED… · pp. 23–24

H EVALUATION PROTOCOL AND RESIDUAL CONVENTIONS

The evaluation protocol measures removal cost with paired evaluations of the same checkpoint and environment seeds, then summarizes normalized scores and amplitude trajectories using specified aggregation conventions.

H EVALUATION PROTOCOL AND RESIDUAL CONVENTIONS

  • Each checkpoint was evaluated with and without the shaping head on the same environment seeds, making removal cost a paired return difference within a run.Terminal scores average the last 20 evaluations; checkpoints are selected from the offset trajectory and evaluated for 10 deterministic episodes.Source: H EVALUATION PROTOCOL AND RESIDUAL… · pp. 24–25
  • Scores are min-max normalized per environment using shared random and best anchors, preserving the sign and order of within-environment differences.Aggregate scores use IQM with 95% stratified bootstrap intervals; removal-cost IQM is computed from paired per-run differences.Source: H EVALUATION PROTOCOL AND RESIDUAL… · p. 25
  • Across four residual conventions, the residual-removal-cost rank correlation stayed between 0.565 and 0.671, with p < 0.01 in every case.The number of environments counted as completed varied from 7 to 10 under TD3 and 13 to 14 under SAC.Source: H EVALUATION PROTOCOL AND RESIDUAL… · p. 25, Table 8 · p. 25

I PER-ENVIRONMENT RESULTS

The shaped TD3 actor outperforms its host on the 20-environment aggregate, while per-environment removal intervals often include zero. The amplitude decomposition shows that gate decline, rather than head-output decline, accounts for most of the amplitude fall.

I PER-ENVIRONMENT RESULTS

  • Removal-cost intervals contained zero in 15 of 20 environments when shaped TD3 and SAC actors were pooled.Figure 7 shows per-environment removal cost as the mean over pooled runs with 95% bootstrap confidence intervals.Source: I PER-ENVIRONMENT RESULTS · p. 25, Figure 7 · p. 26

J COMPARISON WITH THE HOST ALGORITHM

  • The shaped TD3 actor’s IQM exceeded TD3 by +0.077 [0.040, 0.116] over 20 environments, with higher mean return in 14.The comparison used default hyperparameters and no per-task setting.Source: J COMPARISON WITH THE HOST ALGORITHM · p. 27

K AMPLITUDE DECOMPOSITION

  • The median decline in |α| was 98.14%, while the head-output norm remained above half its peak in 10 of 12 cells.The amplitude factorizes as ∥δ∥ = |α| · ∥h∥, with the gate closing almost completely while the head retains most of its scale.Source: K AMPLITUDE DECOMPOSITION · p. 27

L UPDATE-TO-DATA RATIO

The shaped SAC actor with an ensemble critic exceeds its host at every tested update-to-data ratio, while a tenfold change in the initial gate produces mixed return changes. Atari results report both head-removal performance and where the gate engages.

L UPDATE-TO-DATA RATIO

  • Under SAC-REDQ, the shaped actor’s aggregate IQM exceeded its host by 0.029, 0.038, and 0.064 across the three tested update-to-data ratios.The experiment used an ensemble critic with N = 10 networks and M = 2 sampled per update.Source: L UPDATE-TO-DATA RATIO · p. 27, Table 10 · p. 29

M SENSITIVITY TO ε

  • A tenfold increase in ε raised terminal return in four of eight panels and lowered it in four, by at most 11.4%.The compared values were ε = 0.01 and ε = 0.1 across four environments under both algorithms.Source: M SENSITIVITY TO ε · p. 27, Figure 9 · p. 30

N ATARI

  • On 26 Atari games, the base actor alone had IQM 0.183 [0.157, 0.216], above the shaped actor’s 0.170 and the host’s 0.154.The gate engaged in seven games; three completed the arc within one million steps, while four remained active.Source: N ATARI · pp. 27–28

O EXTENDED RELATED WORK

  • The related-work discussion distinguishes this paper’s training-time absorption from hand-folded linear branches, scheduled removal, and residual policies with frozen bases.It also situates the proposed mechanism among accounts of drift near manifolds of minimizers.Source: O EXTENDED RELATED WORK · pp. 31–32

P HYPERPARAMETERS

Related papers

Loading 2609.32752v2…