The policy gradient that leaks an action—until it vanishes

Picture a household robot learning from camera frames. Its device keeps the images local and sends gradients to a coordinator. A gradient is the vector of tiny parameter changes requested by one training example. It feels safer than uploading pixels, but it was computed from those pixels—and can retain their structure.

Bhujel and colleagues study that leak for embodied reinforcement learning (RL), where an agent acts inside an environment. Their attacker, TRACE, receives an ordered, per-step stream of policy gradients and reconstructs both RGB frames and actions. The coordinator is honest-but-curious: it follows the protocol but knows the policy architecture, checkpoint, action space, and learning rule, and owns auxiliary trajectories from the same task family.

Let π_k be the policy probability of action k, A theadvantage (how much better the sampled action looked than the baseline), and h the non-negative hidden representation feeding the policy head. Summing gradient column k gives σ_k = A‖h‖₁(π_k − 1[k = a]). If A is positive, the sampled action is the smallest column sum; if A is negative, it is the largest.

def recover_action(probabilities, action, advantage, hidden_l1):
    columns = [
        advantage * hidden_l1 * (p - int(k == action))
        for k, p in enumerate(probabilities)
    ]
    if advantage > 0:
        recovered = min(range(len(columns)), key=columns.__getitem__)
    else:
        recovered = max(range(len(columns)), key=columns.__getitem__)
    others = [p for k, p in enumerate(probabilities) if k != action]
    gap = abs(advantage) * hidden_l1 * (
        1 - probabilities[action] + min(others)
    )
    return columns, recovered, gap

With three or more actions, the paper also proves a sign-free rule: choose the largest absolute column sum. With exactly two actions, the magnitudes tie, so the sign of A must come from another signal such as the value head. The experiments use five actions.

The theorem's separation is Δ = |A|‖h‖₁(1 − π_a + min_k≠a π_k). As the policy becomes deterministic, π_a approaches one and every alternative approaches zero. Both Δ and the entire policy-gradient block vanish. The same formula that exposes the action also states exactly when that evidence disappears.

TRACE encodes each gradient, interleaves it with causal context tokens, and uses a causal transformer to decode images and actions autoregressively. “Causal” means frame t may use gradients up to t, never future ones. On 100 held-out eight-step trajectories from AI2-THOR, it beats single-frame and optimization-based attacks under the paper's adapted RL protocol.

Mild compression barely changes the attack. Keeping only the largest 10% of PPO gradient entries leaves 99.9% action accuracy, and eight-bit quantization does the same. Coarse quantization, strong noise, and differentially private SGD (DP-SGD) do much more—but the paper does not measure how much navigation quality they sacrifice.

The appendix also tries gradients averaged over four or eight steps. Once order is lost, the action target becomes a histogram, not a sequence, and image quality is scored only on the last frame. For four-step windows, a from-scratch attacker reports 17.03 dB PSNR and total-variation distance 0.296; at eight steps, 16.03 dB and 0.268. Those are single runs under a changed target, not evidence that ordinary local training is broadly invertible.

First, sweep policy entropy and advantage magnitude while measuring action error against Δ; that directly tests the theorem's predicted failure surface. Then separate policy, value, and shared-backbone gradients to learn which block leaks pixels after the action block goes quiet. Finally, replace ordered per-step gradients with realistic local updates: shuffled batches, multiple epochs, clipping, secure aggregation, and an explicit control-utility budget.

The engineering lesson is not “never share gradients.” It is that gradients need a threat model. A neural networkcan expose behavior through a few signed column sums and perception through thousands of correlated coordinates. Audit the exact artifact the server receives, not the fact that raw data stayed on device.