Locally Constrained Representations in Reinforcement Learning

TL;DR

This paper addresses the challenge that RL-trained representations can overfit to evolving value targets by introducing Locally Constrained Representations (LCR), an auxiliary loss that enforces linear predictability of a state’s latent φ_T from neighboring states’ latents. The method defines a neighborhood window of size K and learns nonnegative weights W to form a linear predictor W φ_nearest(T) that approximates φ_T, integrating this soft constraint with the main RL objective. Empirical results across MiniGrid, MuJoCo, Robosuite, and Atari show improved performance and robustness, especially in continuous control tasks where environment dynamics are well-behaved locally, with ablations illuminating hyperparameter sensitivities. The approach demonstrates that decoupling representation learning from the RL loss via local linearity constraints can yield more stable, generalizable representations and faster learning in diverse environments. Overall, LCR offers a practical, broadly applicable technique to boost RL performance by embedding local dynamics into the representation learning process.

Abstract

The success of Reinforcement Learning (RL) heavily relies on the ability to learn robust representations from the observations of the environment. In most cases, the representations learned purely by the reinforcement learning loss can differ vastly across states depending on how the value functions change. However, the representations learned need not be very specific to the task at hand. Relying only on the RL objective may yield representations that vary greatly across successive time steps. In addition, since the RL loss has a changing target, the representations learned would depend on how good the current values/policies are. Thus, disentangling the representations from the main task would allow them to focus not only on the task-specific features but also the environment dynamics. To this end, we propose locally constrained representations, where an auxiliary loss forces the state representations to be predictable by the representations of the neighboring states. This encourages the representations to be driven not only by the value/policy learning but also by an additional loss that constrains the representations from over-fitting to the value loss. We evaluate the proposed method on several known benchmarks and observe strong performance. Especially in continuous control tasks, our experiments show a significant performance improvement.

Figures (12)

• Figure 1: (a) We show the proposed training formalism of the Locally Constrained Representations Algorithm. We consider the time-step $T$, where at state $s_T$, the agent takes actions $a_T$. For value-based RL methods, we have Mean Squared TD Error (MSTDE) that learns the state values, $Q$. In addition, we have the loss, which is a Mean Squared Error (MSE) loss between the latent state $\Phi_T$ and a linear combination of the neighboring latent states, $\Phi_{T-2}$,$\Phi_{T-1}$,$\Phi_{T+1}$ and $\Phi_{T+2}$. This loss encourages the current representation $\Phi_T$ to be closer to this linear combination. (b) An example is shown for a sequence length of 5, where the current processed state is shown in red and the neighboring states considered are in blue. The loss calculated is fed directly to the Loss in (a)
• Figure 2: Performance of DQN (red) and DQN with LCR (blue) on the MiniGrid Environments. Both algorithms were trained for 10 runs with LCR using a sequence length of 11, 100 gradient steps, and a batch size of 5000. The detailed hyperparameters are mentioned in the appendix.
• Figure 3: tSNE plots of the state representations obtained by 20 random trajectories of the respective environments. constrains the state representations by encouraging linear predictability with respect to its neighboring representations.
• Figure 4: Training curves on 6 Mujoco environments using SAC with and without LCR across 10 runs. This figure highlights the impact of in these domains with well-defined physics. The poor performance of ATC in these settings highlights that constraining representations based on proximity is not a good strategy and adding linear predictability is much stronger.
• Figure 5: Training curves on 6 Robosuite environments with two robots Sawyer and Panda. We see that adding to the SAC consistently improves performance. All the curves are averaged across 10 independent seeds with the shaded portion representing the standard error.