NASDAQ: Normalized Observation Space Dynamics-Augmented Q-Learning

2026-06-19

AI Summary

This work addresses the performance limitations of observation-prediction reinforcement learning under low-dimensional observations, where imbalanced reconstruction losses across dimensions hinder learning. To mitigate this issue, the authors propose an observation normalization method tailored for online reinforcement learning that unifies the handling of both low-dimensional states and high-dimensional image inputs. Dynamic modeling is performed in the normalized observation space, integrating the primary Q-learning objective with two auxiliary tasks: prediction of the next normalized observation and short-term value estimation. This joint training strategy effectively balances reconstruction gradients and enhances sample efficiency. Empirical results demonstrate that the proposed framework matches or surpasses existing model-based and self-predictive reinforcement learning baselines across multiple benchmark tasks, while significantly reducing training time.

Abstract

Augmenting model-free reinforcement learning (RL) with representations learned through observation dynamics prediction (observation-predictive RL) can improve sample efficiency and performance, with minor modifications and limited additional computation. However, this approach still struggles in challenging tasks with low-dimensional observations. In this paper, we identify a key factor behind this problem: unbalanced reconstruction losses across observation dimensions, where dimensions with larger value ranges dominate the loss. This encourages the agent to neglect dimensions with relatively small ranges, leading to degraded performance. To address this issue, we propose a novel normalization method tailored to online RL, which normalizes low-dimensional observations and balances the resulting losses and gradients. Beyond balancing reconstruction losses, observation normalization enables dynamics prediction to be performed in a normalized observation space, thereby providing a unified treatment of low- and high-dimensional inputs (e.g., physical states and images). Building on this idea, we further introduce Normalized Observation Space Dynamics-Augmented Q-learning (NASDAQ), a framework for observation-predictive RL applicable across diverse domains. NASDAQ learns state-action representations by coupling value learning with two auxiliary tasks: short-term value prediction and next normalized observation prediction. Extensive experiments demonstrate that NASDAQ achieves competitive or superior performance compared with state-of-the-art model-based and self-predictive RL methods, while requiring significantly less training wall-time.

Problem

• reconstruction loss imbalance
• low-dimensional observations
• observation dynamics prediction
• sample efficiency
• reinforcement learning

Innovation

• observation normalization
• reconstruction loss balancing
• dynamics-augmented reinforcement learning
• sample efficiency
• online RL