When cutting policy memory gradients affects learning and control accuracy

When Does Backpropagating Through Policy Memory Matter? Physical Credit, Optimizer Updates, and Observability

Machine LearningRobotics

Summary

This paper investigates how cutting certain backward connections in neural networks with memory affects learning and control tasks. The authors test this by changing how gradients flow through stored memory but keep the same forward computations. They find that cutting gradient paths through memory sometimes drastically reduces learning signal strength, which can affect the quality of control policies, like those controlling a drone's movement. The impact depends on factors like how noise impacts the system and how the optimizer processes updates. The paper suggests that measuring these effects requires comparing actual optimizer updates during training, not just raw gradients.

What this means in practice

  • For robotics engineers: Optimize training strategies for control policies by deciding when to allow gradients through memory for better trajectory and tracking results.
  • For ai system builders: Improve training of recurrent and memory-augmented networks by choosing backward graph truncations that balance computational cost and learning performance.

Authors

Xingjian Li, Yi Han, Jianhua Z. Huang

Abstract

Policies with memory can learn along two backward paths: through the physical states their actions produce and through the representations they store. Transformer-XL and truncated backpropagation through time cut the second path at stored history while keeping its values. We ask when this cut matters. Holding the forward computation fixed and varying only derivative edges, we measure parameter gradients, the updates the optimizer applies, and continued training in a Transformer vessel-trajectory model and a quadrotor tracking policy. In the vessel model, detaching the key-value cache shrank the gradient to about a tenth of its norm, with little rotation, when gradients flowed through all earlier physical states, but barely changed it under one-step physical credit. In this strongly clipped regime the optimizer, not the gradient, set how far updates differed: global-norm clipping removed most of the gradient difference between memory-cut graphs, whereas AdamW turned a 2% gradient difference between two placements of the cut into update differences of up to 31% at the step where the placement was switched. In a quadrotor trained from initialization with 0.20 m/s velocity noise, removing memory raised tracking error by 43% and cutting memory gradients raised it by 32%; at low noise the cut's mean cost exceeded the value of memory. Two-step truncation segments gave no measurable gain, although with hidden velocity a two-step window captured most of the value of memory; eight-step segments removed half to three quarters of the cost. Switching the cut on only for the last fifth of training understated its cost about threefold at 0.20-0.30 m/s, but not at low noise or with hidden velocity. These results suggest measuring the cost of a memory cut by training with it from initialization, and comparing backward graphs by the updates the optimizer applies rather than by raw gradients.