Temporal Gradient Inversion for Private Trajectory Reconstruction in Embodied Reinforcement Learning

AI summary

Overview

Research area: Privacy and security in distributed/embodied reinforcement learning, specifically gradient inversion attacks on actor-critic policy gradients (A2C and PPO).

Technical level: Advanced. The paper combines gradient-inversion attack design, causal transformer architectures, and information-theoretic analysis (conditional mutual information, Bayes-optimal reconstruction error).

Scope: The paper proposes and evaluates TRACE, an amortized temporal gradient-inversion attack that reconstructs sequences of private RGB observations and discrete actions from per-step policy gradients transmitted by embodied RL agents.

What This Paper Is About

Embodied RL agents (robots, navigation agents) keep raw camera data on-device and send only policy gradients to a server, which is often assumed to protect privacy. This paper shows that the temporal ordering of those gradients leaks far more than single-frame attacks capture: consecutive frames of the same physical scene are correlated, and the policy-head gradient contains a closed-form signature of the action taken. The goal is to build a fast, learned attacker that autoregressively reconstructs the agent's full observation-action trajectory from the ordered gradient stream.

Key Contributions

1. TRACE attack framework. An amortized trajectory-level inversion method with three parts — a gradient encoder that compresses the high-dimensional update into a low-dimensional latent, a causal transformer that propagates trajectory context, and a decoder that jointly reconstructs observations and predicts actions in milliseconds.

2. Theoretical foundation for actor-critic gradient leakage. Proposition 4.1 shows that A2C and PPO produce the same communicated gradient at the broadcast checkpoint θ₀, decomposing it into policy, value, and entropy gradient terms. Theorem 4.2 proves closed-form, learning-free discrete-action recovery from the column sums of the policy-head gradient, exact when entropy regularization is zero (c_e = 0) and still exact under the stated condition when c_e > 0.

3. Information-theoretic characterization of temporal gain. Definition 4.4 defines temporal-context gain as the conditional mutual information ΔI_t = I(s_t; g_{1:t−1} | g_t) = H(s_t | g_t) − H(s_t | g_{1:t}) ≥ 0, and Theorem 4.5 shows via the Law of Total Variance that MSE*(g_{1:t}) ≤ MSE*(g_t), i.e., adding past gradients cannot hurt reconstruction.

4. Broad empirical evaluation. Comparison against optimization-based (DLG, IG) and learning-based (LtI) baselines, few-shot adaptation to held-out scenes, four gradient defenses, plus additional tests on recurrent, residual, and compact transformer victims, multi-modal inputs, and larger discrete action spaces.

Main Findings

• TRACE beats all baselines on reconstruction. Under PPO victims, TRACE reaches MSE 0.014 ± 0.007, PSNR 18.771 ± 1.778, SSIM 0.627 ± 0.105, LPIPS 0.362 ± 0.070, with 100.0% action accuracy. Under A2C: MSE 0.014 ± 0.005, PSNR 18.873 ± 1.391, SSIM 0.583 ± 0.061, LPIPS 0.358 ± 0.060, 100.0% action accuracy.

• Optimization attacks fail badly in this setting. DLG reaches 5.474 ± 0.826 dB PSNR (PPO) and 6.106 ± 1.544 dB (A2C), with action accuracies of 16.6% and 51.0%. IG reaches 8.727 ± 2.791 dB (PPO) and 9.511 ± 0.810 dB (A2C), with 31.5% and 32.5% action accuracy.

• The learning-based baseline is strong; temporal context still adds a clear gain. LtI achieves 16.790 ± 1.708 dB (PPO) and 17.652 ± 1.366 dB (A2C) with 100.0% action accuracy. TRACE reduces MSE from 0.023 to 0.014 (PPO) and 0.018 to 0.014 (A2C), gains roughly 2 dB and 1 dB PSNR, and nearly halves LPIPS (0.671 to 0.362 for PPO; 0.639 to 0.358 for A2C).

• Speed. TRACE runs at 4.5 ms per frame (PPO) and 3.0 ms (A2C), compared with tens of seconds for DLG (31.0 s PPO, 76.1 s A2C) and IG (39.3 s PPO, 39.6 s A2C). LtI is faster still at 1.2 ms (PPO) and 0.2 ms (A2C).

• Single-frame TRACE is weaker than LtI. At T = 1, TRACE achieves 12.85 dB (Table 4), underperforming LtI's single-frame result — evidence that TRACE's advantage comes from temporal conditioning rather than a stronger per-frame extractor.

• Zero-shot transfer works, and small amounts of target data help a lot. With p = 0% target trajectories, TRACE achieves MSE 0.014 ± 0.007, PSNR 18.8 ± 1.8, SSIM 0.627 ± 0.105, LPIPS 0.362 ± 0.070, and 100.0% action accuracy. At 10% it improves to 20.9 ± 1.8 dB PSNR and LPIPS 0.273 ± 0.077; at 50% it reaches MSE 0.005 ± 0.002, PSNR 23.8 ± 2.1, SSIM 0.740 ± 0.092, LPIPS 0.161 ± 0.061, still with 100.0% action accuracy.

• Most defenses barely help; heavy quantization does degrade the attack. With no defense, TRACE achieves 18.9 ± 3.7 dB (PPO) and 18.4 ± 1.5 dB (A2C). 8-bit quantization (18.5 ± 3.9 dB PPO; 18.4 ± 1.5 dB A2C), 90% and 50% pruning (identical to no defense), 10% pruning (18.8 ± 3.8 dB PPO; 18.0 ± 1.7 dB A2C), and Gaussian noise σ = 0.001 (18.8 ± 3.8 dB PPO; 18.4 ± 1.5 dB A2C) leave reconstruction largely intact. 4-bit quantization drops PSNR to 12.4 ± 3.2 dB (PPO) and 13.3 ± 2.7 dB (A2C) with action accuracy 37.9% and 50.2%; 2-bit drops to 11.3 ± 2.8 dB (PPO) and 11.5 ± 1.8 dB (A2C) with 19.4% and 41.4% action accuracy.

• Action recovery has a theoretical limit. Remark 4.3 shows that for a fully deterministic policy (π_{a_t} → 1), every column sum of the policy-head gradient vanishes, Δ_t = 0, and the action becomes unrecoverable from that gradient block alone.

• Temporal gain is guaranteed in the tractable case. Theorem 4.5(i) establishes non-negativity; Theorem 4.5(ii) gives an exact identity under the Markov victim policy and un-normalized one-step advantage/return estimators of Lemma D.3. GAE with λ > 0 and rollout-level normalization violate those conditions.

Methodology in Plain English

The researchers adopt a passive, honest-but-curious server threat model: the server knows the model architecture, the current checkpoint θ₀, the action space, and the learning rule, and it holds auxiliary trajectories from the same task family for offline training. At inference it sees only the ordered per-step gradients.

Attack pipeline. A gradient encoder compresses the raw gradient g_t into a compact latent z_t; an image encoder turns observations (or a learned start token during warmup) into context embeddings e_t. These are interleaved into a sequence [z₁, e₀, z₂, e₁, …, z_T, e_{T−1}] and fed to a causal transformer with a lower-triangular mask, so each gradient token only attends to the past. Gated residual connections with learnable scalars keep residual branches initially down-weighted. A residual decoder then splits into an image head (transposed convolutions with sigmoidal output) and an action head (a two-layer MLP producing a categorical distribution over actions).

Training. An initial teacher-forced warmup parallelizes the forward pass and prevents compounding errors, but creates exposure bias. Two fixes are used: scheduled sampling, where the next image-context token comes from ground truth or the model's own prediction with epoch-wise probability β_n that ramps linearly from β_start to β_end after a warmup of τ epochs over N total epochs; and an auxiliary short-horizon rollout loss that unrolls R autoregressive steps, supervises them, and feeds predicted frames back as detached context. The total objective is a weighted sum: α₁·MSE + α₂·L1 + α₃·action + α₄·perceptual + α₅·rollout.

Evaluation. Victim policies are trained on point-goal navigation in the AI2-THOR simulator. The victim sees RGB frames downsampled to 84 × 84 and chooses among five discrete actions: MoveAhead (0.25 m grid step), RotateLeft (15°), RotateRight (15°), LookDown (15°), LookUp (15°). The default victim is a shared-CNN actor-critic with separate policy and value heads, trained with PPO or A2C using GAE advantages and the same critic, entropy regularization, and gradient-norm bound. All methods are evaluated on 100 held-out trajectories with T = 8, and results are reported as mean ± one standard deviation per image. Baselines (DLG, IG, LtI) are adapted from supervised classification by swapping in a simplified actor-critic surrogate of policy cross-entropy plus value regression with fixed coefficients.

Why This Matters

This work reframes gradient leakage in distributed RL as a trajectory-inversion problem rather than a collection of independent image inversions. It shows that temporal structure, which existing single-frame attacks treat as noise, is itself a side channel — and that a learned attacker can exploit it to near-real-time operation while a theory-derived action rule recovers discrete actions with no learning at all. The defense results suggest that standard perturbation and pruning defenses are largely ineffective against sequence-level attacks.

Real-world applications:

• Robot fleets and shared navigation. Distributed policy learning across robots could leak time-ordered camera views revealing room layout, objects, and where the agent moved, looked, and lingered.
• Activity profiling and premise mapping. Recovered trajectories enable inference about routines and physical spaces from what is nominally privacy-preserving gradient sharing.
• Persistent monitoring. A curious aggregation server could track an agent over time from a stream of per-step gradients rather than a single snapshot.
• Privacy auditing for collaborative RL deployments. The attack provides a concrete stress test for whether a given gradient-sharing protocol truly protects embodied trajectories.

Industry relevance. Any operator running collaborative or federated policy learning for embodied agents — robotics, autonomous navigation, smart-facility automation — needs to know that per-step gradient sharing may not preserve observation privacy, and that lightweight defenses like pruning or small Gaussian noise do not close the gap.

Future Directions

• Sequence-aware privacy mechanisms. The paper's defense experiments indicate that protecting temporal gradient streams may require mechanisms that model trajectories as sequences, not independent updates — a direct call for new defense designs.
• Extending the theory beyond the tractable case. Lemma D.3's conditions exclude GAE with λ > 0 and rollout-level normalization, which the paper says introduce cross-temporal coupling; the scaling behavior under practical PPO and A2C estimators is only discussed informally and validated empirically.
• Attacks on multi-step client updates. The paper focuses on the per-step gradient setting and notes that multi-step updates Δθ = θ₀ − θ_client combine gradients along a changing optimization path and "may reduce the temporal structure available to the attacker" — an open question.
• Broader victim and modality coverage. The paper already extends to recurrent, residual, and compact transformer victims, multi-modal inputs, and larger discrete action spaces, and also examines four- and eight-step gradient averaging and auxiliary-data scaling; further generalization to other domains remains untested here.

Target Audience

Researchers and practitioners in ML privacy and security, especially those working on gradient inversion, federated and distributed reinforcement learning, and embodied AI. It is also relevant to engineers deploying collaborative policy learning who need to reason about what ordered gradient streams reveal, and to theorists interested in information-theoretic characterizations of temporal leakage. Readers should be comfortable with actor-critic RL (A2C, PPO, GAE), transformer architectures, and basic information theory; the mathematical derivations will be difficult for beginners.

Authors’ abstract

Distributed learning in embodied reinforcement-learning agents offers a degree of privacy by retaining raw sensor data on-device and transmitting only policy gradients to the server. Yet temporal structure can amplify this leakage beyond single-frame attacks. We introduce Temporal Reconstruction Attack on Consecutive Encodings (TRACE), an amortized temporal gradient-inversion attack that autoregressively reconstructs the sequence of private observation-action trajectories from per-step policy-learning gradients. The attack exploits two structural signals ignored by prior single-frame methods: (i) cross-time correlation between successive embodied gradients, which we formalize via a conditional mutual-information bound, and (ii) closed-form action recovery from policy-head gradient structure, which we prove exact when standard entropy regularization is sufficiently small. On held-out embodied scenes, TRACE reaches $18.8$ dB PSNR with near-perfect action recovery at $3$-$4.5$ ms per reconstructed frame, dominating the learning-based baseline across all reconstruction metrics and exceeding optimization attacks while running orders of magnitude faster. Further evaluation demonstrates TRACE's broader applicability across recurrent, residual, and compact transformer victim architectures, multi-modal inputs, and larger discrete action spaces. Defense experiments suggest that protecting temporal gradient streams may require sequence-aware privacy mechanisms.