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

Abstract: While representation and similarity learning have improved the sample efficiency of Reinforcement Learning (RL), they are rarely used to shape policy updates directly in the action space. To bridge this gap, a geometry-aware RL algorithm that explicitly incorporates value-based similarity into the policy update, State-Action Value Geometry Optimization (SAVGO), is proposed. In detail, SAVGO learns a joint state-action embedding space in which pairs with similar action-value estimates exhibit high cosine similarity, while dissimilar pairs are mapped to distinct directions. This learned geometry enables the generation of a similarity kernel over candidate actions sampled at each update, allowing policy improvement to be guided directly toward higher-value regions beyond local gradient-based updates. As a result, representation learning, value estimation, and policy optimization are unified within a single geometry-consistent objective, while preserving the scalability of off-policy actor-critic training. The proposed method is evaluated on standard MuJoCo continuous-control benchmarks, demonstrating improvements over strong baselines on challenging high-dimensional tasks. Ablation studies are done to analyze the contributions of value-geometry learning and similarity-based policy updates.

• The paper introduces a geometry-aware RL algorithm that leverages cosine similarity to map state–action pairs onto a unit-norm hypersphere for enhanced policy updates.
• It unifies representation learning, value estimation, and policy improvement by employing a similarity-weighted kernel to achieve robust and sample-efficient updates.
• Empirical evaluations on MuJoCo tasks demonstrate SAVGO’s superior performance, especially in high-dimensional environments, with effective hyperparameter tuning.

SAVGO: State–Action Value Geometry Optimization with Cosine Similarity for Continuous Control

Introduction and Motivation

State–action representation learning has traditionally been used in deep reinforcement learning (RL) to improve value estimation and sample efficiency, often as an auxiliary regularizer rather than a direct decision variable in the policy improvement process. The paper "SAVGO: Learning State-Action Value Geometry with Cosine Similarity for Continuous Control" (2605.00787) proposes a geometry-aware RL algorithm that explicitly utilizes learned state–action value similarity to guide policy updates. SAVGO introduces a framework where representation learning, value estimation, and policy optimization are unified through a geometry-consistent objective, leveraging a novel use of cosine similarity in a joint embedding space.

Methodology: Value-Aware State–Action Geometry and Policy Improvement

SAVGO learns an embedding function mapping state–action pairs into a unit-norm hypersphere, where the cosine similarities between embeddings precisely reflect action-value proximity. The core innovation is the construction of a geometry-aware kernel defining action similarity within this space, which then acts as the basis for policy improvement.

Formally, for state–action pairs $(s_i, a_i)$ and $(s_j, a_j)$, their embeddings $\zeta_i = z_\psi(s_i, a_i)$ and $\zeta_j = z_\psi(s_j, a_j)$ are trained with a regression loss so that their cosine similarity $\cos(\zeta_i, \zeta_j)$ matches a target $Y_{i,j}$ determined by their normalized $Q$-value gap. The normalized value gap is defined by:

$$\Delta_{i,j} = \mathrm{clip}\left(\frac{|Q(s_i, a_i) - Q(s_j, a_j)|}{\beta}, 0, 1\right),$$

and the target similarity is mapped through a curvature-parameterized transformation:

$Y_{i,j} = 1 - 2(\Delta_{i,j})^\lambda,$

where $\lambda$ controls the discrimination sharpness in the learned metric. This ensures that the geometry of the embedding space aligns with differences in action values and is robust to scale.

Policy improvement exploits this geometry by sampling a set of candidate actions for each state, computing embeddings, and aggregating $(s_j, a_j)$0-estimates through a similarity-weighted kernel. The final policy gradient is driven by this aggregated, geometry-aware signal, effectively smoothing and stabilizing updates beyond conventional local gradients.

This is illustrated in the following depiction:

Figure 1: Policy improvement with value-aware geometry—candidate actions are sampled, reweighted by learned cosine similarity, and aggregated for stable updates beyond single-point gradients.

The mapping from $(s_j, a_j)$1-value gap to target cosine similarity is parameterized and critical for geometry learning:

Figure 2: Target cosine similarity $(s_j, a_j)$2 as a function of normalized value gap $(s_j, a_j)$3, showing the effect of curvature parameter $(s_j, a_j)$4 on the discrimination regime.

Furthermore, the learned updates manifest as smooth, geometry-driven transitions in the embedding space:

Figure 3: Unit-sphere view—anchor embeddings and candidates yield a direction (green) directly oriented toward high-Q actions via similarity weighting.

Empirical Evaluation: Performance and Ablations

SAVGO is benchmarked on multiple MuJoCo continuous control tasks against PPO, TD3, SAC, and TQC, using identical training protocols and hyperparameter search spaces. The primary empirical findings are:

• SAVGO achieves the highest total score across the MuJoCo benchmark suite, with the most prominent gains on high-dimensional tasks such as Humanoid and Walker2d.
• In sample efficiency and learning stability, SAVGO demonstrates accelerated learning and lower return variance, notably outperforming baselines on challenging high-dimensional problems.
• When assessing maximum evaluation returns within 1M steps, SAVGO either matches or exceeds strong baseline methods in almost all environments, with the disparity increasing with environment complexity.

The learning curves directly illustrate the comparative sample efficiency and convergence dynamics:

Figure 4: Learning curves over 1M steps on MuJoCo (v5); SAVGO consistently matches or outperforms baselines, especially in the Humanoid and Walker2d environments.

Crucial ablation studies analyze the sensitivity of results to hyperparameters and architectural choices. The curvature parameter $(s_j, a_j)$5 in the similarity mapping is essential: excessively sharp or smooth mappings hurt performance and stability, with optimal performance near $(s_j, a_j)$6–$(s_j, a_j)$7.

Figure 5: Ablation of the representation curvature parameter $(s_j, a_j)$8—strongly influences peak returns and learning variance.

The candidate set size $(s_j, a_j)$9 balances kernel coverage versus computational cost: larger $\zeta_i = z_\psi(s_i, a_i)$0 improves stability, especially in higher-dimensional tasks, but with diminishing returns.

Figure 6: Sensitivity to candidate set size $\zeta_i = z_\psi(s_i, a_i)$1 in the similarity-weighted update operator; sufficient $\zeta_i = z_\psi(s_i, a_i)$2 is necessary for high-dimensional action spaces.

Design ablations further confirm that all proposed components (adaptive temperature, representation learning, normalization, similarity kernel) are required for full performance. Removing geometry learning or using uniform weighting negates observed improvements.

Discussion and Implications

SAVGO introduces a distinct mechanism whereby state–action similarity, learned via value-aware cosine geometry, directly shapes policy improvement in continuous control. Unlike prior approaches using value-consistent or reward-aware auxiliary objectives (e.g., [fujimoto2023sale]; [yue2023vcr]), SAVGO operationalizes similarity as an actionable policy improvement operator, not merely as a regularizer or for value target shaping.

Theoretical and empirical implications include:

• Representation learning, value estimation, and policy optimization can be unified within a geometry-consistent, actor–critic pipeline, enabling policy improvement to exploit global structure rather than relying solely on unreliable local gradients.
• The similarity-weighted kernel approach enables more robust and sample-efficient updates, particularly valuable under noisy critics or in high-dimensional action spaces, as demonstrated in the empirical results.
• Adapting hyperparameters central to geometry design (e.g., the curvature $\zeta_i = z_\psi(s_i, a_i)$3, temperature schedule, robust scaling) is critical. Both underfitting and overfitting the value-aware similarity mapping degrade results.
• The computational cost scales linearly with the number of sampled candidates $\zeta_i = z_\psi(s_i, a_i)$4, necessitating careful balancing in large action spaces.

Practically, this framework leverages controlled geometric aggregation for improved policy exploitation, aligning with advances in kernel- and similarity-based RL for efficient value propagation (see [gao2026ira], [shen2021nnac], [castro2021mico]). However, limitations remain for very high-dimensional or discrete action spaces due to the cost of candidate sampling and the need for effective coverage.

Future Prospects

Potential advances include:

• Extending geometry-aware policy optimization to discrete or hybrid action spaces, leveraging alternative encoders or kernel constructions.
• Improved candidate sampling heuristics to maintain efficiency in very high-dimensional settings.
• Integration with temporally extended abstraction methods to accelerate long-horizon credit assignment.
• Application in partially observable or highly stochastic domains, with further robustness studies on critic and geometry noise.

Conclusion

SAVGO presents a coherent algorithmic design for integrating value-aware state–action geometry directly into the core of policy improvement in continuous control. Empirical results support the efficacy of the approach, especially in complex domains, and ablation studies validate the necessity and effectiveness of each design choice. The proposed framework opens new avenues for geometry-aware and similarity-driven RL agents, with practical implications for both research and large-scale RL application deployment.