SAVGO: Learning State-Action Value Geometry with Cosine Similarity for Continuous Control

1. Problem Statement

Deep Reinforcement Learning (RL), particularly off-policy actor-critic methods like SAC and TD3, has achieved success in continuous control but remains limited by sample inefficiency and instability in high-dimensional environments.

• Current Limitations: While representation learning (e.g., contrastive learning) has improved data efficiency, it is typically treated as an auxiliary task. The policy update step usually relies on local gradient descent, which ignores the global geometric structure of the learned latent space.
• The Gap: Existing similarity-based methods often use distances to regularize representations or shape value targets, but they rarely use these geometric relationships as first-class operators to directly guide policy updates in the action space. Consequently, policy improvements remain local and weakly informed by the global value landscape.

2. Methodology: SAVGO

The authors propose State–Action Value Geometry Optimization (SAVGO), a geometry-aware off-policy actor-critic algorithm that explicitly integrates value-based similarity into the policy update mechanism.

A. Value-Aware State-Action Geometry Learning

SAVGO learns a joint embedding space for state-action pairs $(s, a)$ where the geometric distance (specifically cosine similarity) reflects the similarity in expected utility (Q-values).

• Encoder: A neural network $z_\psi(s, a)$ maps state-action pairs to a $d$-dimensional embedding space.
• Target Construction: For any pair of state-action tuples $(s_i, a_i)$ and $(s_j, a_j)$, a normalized value gap $\Delta_{i,j}$ is computed based on the difference in their Q-values.
• Curvature Mapping: This value gap is mapped to a target cosine similarity $Y_{i,j}$ using a bounded curvature transformation:
  $$Y_{i,j} = 1 - 2(\Delta_{i,j})^\lambda$$
  where $\lambda$ controls the sensitivity. Small gaps (similar values) map to high similarity ($\approx 1$), while large gaps map to low/negative similarity.
• Loss Function: The encoder is trained to minimize the regression loss between the predicted cosine similarity of embeddings and the target $Y_{i,j}$. This forces the embedding space to organize such that actions with similar values cluster together on the unit hypersphere.

B. Geometry-Aware Policy Improvement

Instead of a standard pointwise gradient update, SAVGO performs a similarity-weighted aggregation over candidate actions.

1. Sampling: For a given state $s_t$, an "anchor" action $\hat{a}_t$ is sampled from the current policy. Additionally, $K$ candidate actions $\{a^k_t\}$ are sampled from a proposal distribution.
2. Embedding & Scoring: Both the anchor and candidates are embedded via $z_\psi$. Candidates are scored using a conservative critic estimate ($\min Q$).
3. Kernel Weighting: Cosine similarities between the anchor embedding and candidate embeddings are converted into a probability distribution (similarity kernel) using a temperature-controlled softmax:
   $$w_k = (1-\epsilon)\frac{\exp(\cos(\hat{\zeta}_t, \zeta^k_t)/\rho)}{\sum \exp(\dots)} + \frac{\epsilon}{K}$$
   This creates a weight $w_k$ for each candidate based on its geometric proximity to the anchor in the value-aware space.
4. Aggregated Update: The policy is updated to maximize a similarity-weighted value estimate:
   $$\hat{Q}(s_t, \hat{a}_t) = \sum_{k=1}^K w_k Q(s_t, a^k_t)$$
   This allows the policy to be pulled toward regions of high value that may not be immediately adjacent in the raw action space but are "close" in the learned value geometry.

3. Key Contributions

1. Unified Objective: Proposes SAVGO, which unifies representation learning, value estimation, and policy optimization into a single geometry-consistent objective, moving beyond treating representation learning as a mere regularizer.
2. Novel Policy Operator: Derives a geometry-aware policy improvement operator that uses learned cosine similarities to weight candidate actions, enabling policy updates guided by global value structure rather than just local gradients.
3. Empirical Validation: Demonstrates consistent improvements over strong baselines (SAC, TD3, TQC, PPO) on MuJoCo benchmarks, particularly in high-dimensional tasks.

4. Experimental Results

The method was evaluated on the MuJoCo v5 continuous control suite with a fixed budget of 1 million environment steps.

• Performance: SAVGO achieved the highest total rewards across the suite, with the most significant gains on challenging, high-dimensional tasks:
  • Humanoid: SAVGO (7687) significantly outperformed TQC (6330) and SAC (5351).
  • Walker2d: SAVGO (5626) outperformed TQC (4952) and SAC (4693).
  • Ant & HalfCheetah: Competitive or superior performance with reduced variance.
• Sample Efficiency: SAVGO demonstrated faster early learning and smoother convergence compared to baselines.
• Ablation Studies:
  • Curvature ($\lambda$): Values around 1.0–1.5 provided the best trade-off between stability and performance.
  • Candidate Count ($K$): Moderate values (e.g., $K=128$ for Humanoid) yielded optimal results; too few candidates increased variance, while too many offered diminishing returns.
  • Component Necessity: Removing the adaptive temperature, the representation loss (freezing the encoder), or the adaptive value scaling all caused significant performance drops. Replacing the similarity kernel with uniform weights caused the algorithm to collapse, proving the necessity of geometry-aware weighting.

5. Significance and Limitations

Significance:

• Bridging the Gap: SAVGO successfully bridges the gap between representation learning and policy optimization, showing that the geometry of the latent space can be directly exploited to improve decision-making.
• Robustness: By aggregating over multiple candidates weighted by value-similarity, the method is more robust to critic noise and local minima than standard gradient-based updates.
• Scalability: It retains the scalability of off-policy actor-critic methods while adding a manageable computational overhead (linear in $K$).

Limitations:

• Computational Cost: The method requires evaluating $K$ candidates per state, making it slower than standard SAC/TD3 (approx. 8h vs. 2h for 1M steps in the experiments).
• Critic Dependency: The method assumes the critic provides a reasonably consistent ranking; severe critic noise can destabilize the similarity weights.
• Discrete Actions: Extending this to discrete action spaces (e.g., Atari) is non-trivial due to the difficulty of defining continuous candidate proposals and state-action geometries for discontinuous actions.

In conclusion, SAVGO represents a significant step forward in continuous control by making the "shape" of the value function a direct driver of policy improvement, leading to superior performance in complex, high-dimensional environments.