SAVGO: Learning State-Action Value Geometry with Cosine Similarity for Continuous Control
| Authors | Stavros Orfanoudakis & Pedro P. Vergara |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2605.00787 |
| Download | |
| Categories | cs.LG |
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.
Engineering Breakdown
Plain English
This paper proposes SAVGO (State-Action Value Geometry Optimization), a reinforcement learning algorithm that learns a shared embedding space where state-action pairs with similar value estimates are pulled close together while dissimilar pairs are pushed apart. The core innovation is using this learned geometry to directly shape policy updates—rather than just improving sample efficiency, the algorithm guides policy improvement toward higher-value regions by generating similarity kernels over candidate actions. This bridges a gap in prior work where representation learning and similarity metrics were used for feature learning but not directly incorporated into the policy update mechanism itself. The approach treats policy optimization as a geometry problem, enabling the algorithm to make better action choices by understanding the value landscape in embedding space.
Core Technical Contribution
The novel contribution is embedding state-action pairs into a joint vector space where the cosine similarity between embeddings directly reflects similarity in their Q-values (action values), then using this geometry to modulate policy updates. Unlike prior representation learning approaches in RL that optimize embeddings independently from policy learning, SAVGO tightly couples embedding geometry and policy gradients—the learned similarity kernel over candidate actions becomes a direct input to policy improvement. This is a paradigm shift: instead of learning representations for better feature extraction, SAVGO learns representations where the geometric relationships between actions inform which actions to prioritize. The algorithm effectively creates a value-aware action space where similar-quality actions cluster together, allowing the policy to leverage neighborhood structure rather than treating each action independently.
How It Works
SAVGO maintains a neural network that maps (state, action) pairs into a continuous embedding space, trained so that pairs with similar Q-values are nearby in this space (measured by cosine similarity) and pairs with different Q-values are distant. At each policy update, the algorithm samples a set of candidate actions and computes their embeddings; the cosine similarities between these embeddings form a similarity kernel. This kernel is then used to weight or smooth the policy gradients—actions that are embedded near high-value actions get boosted, while those near low-value actions are suppressed, creating a 'pull' effect toward better regions of the action space. The training objective jointly optimizes the embedding function (to match cosine similarity to actual Q-value differences) and the policy (to exploit the learned geometry), creating a feedback loop where better geometry enables better policy improvement, which in turn produces better Q-estimates for geometry learning. The approach uses standard RL components (value estimates, sampled actions) but connects them through explicit geometric structure rather than implicit gradient flow.
Production Impact
For production RL systems, SAVGO offers a way to improve sample efficiency and convergence speed without requiring larger models or more data—it extracts more signal from existing experience by leveraging action space geometry. In practice, this means faster policy training, more stable updates (because the similarity kernel smooths out noisy individual action estimates), and better generalization to unseen states through the learned embedding structure. The trade-off is moderate: you need to maintain an additional embedding network and compute cosine similarities over candidate actions at each update, adding ~10-20% computational overhead in the forward pass, but this is often worth the speedup in wall-clock training time. Integration is straightforward for existing actor-critic RL codebases—you add an embedding head to your network, define a similarity loss alongside your policy loss, and replace hard action selection with similarity-weighted selection. For large-scale systems (robotics, game playing, autonomous agents), this approach could reduce training time from weeks to days while improving asymptotic performance, making it particularly valuable for expensive-to-simulate environments.
Limitations and When Not to Use This
The paper assumes that action similarity in embedding space should correlate with Q-value similarity, which may not hold in environments with highly non-convex or discrete action spaces, or where locally similar actions lead to very different outcomes. The approach requires differentiable action sampling and embedding computation, limiting applicability to discrete action spaces with large cardinality (e.g., thousands of possible moves) where sampling-based approximation becomes inefficient. SAVGO also assumes sufficient exploration to populate the embedding space meaningfully—in sparse reward or high-dimensional action settings, the learned geometry may be unreliable early in training. The paper doesn't clearly address computational cost under large batch sizes or action dimensionality, nor does it discuss convergence guarantees or theoretical properties of the learned geometry. Additionally, there's limited discussion of how to set hyperparameters for the similarity kernel weight (how much to trust geometry vs. Q-estimates) or how the method scales to very long-horizon tasks where value estimates themselves become unreliable.
Research Context
This work builds on a decade of representation learning in RL (e.g., contrastive methods, inverse models, curiosity-driven learning) and recent interest in geometry-aware optimization (e.g., hyperbolic embeddings, geodesic losses). It directly extends prior work on similarity-based policy improvement but is distinct in that it learns and exploits geometry end-to-end rather than pre-computing similarity from fixed features. The approach is positioned as a middle ground between model-based RL (which explicitly learns transition geometry) and model-free RL (which ignores structure)—you get structure without the cost of full environment modeling. This opens research directions in understanding how different action space geometries affect convergence, whether richer geometric structures (e.g., Riemannian manifolds) could help, and how to extend SAVGO to multi-agent or hierarchical RL where state-action relationships are more complex.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::