SAVGO: Steering Policy Updates Through Value Geometry in Continuous Control

Paper: 2605.00787 Authors: Stavros Orfanoudakis, Pedro P. Vergara Categories: cs.LG

The Gap

Representation learning has made RL more sample-efficient by learning better state features. But these representations rarely touch the policy update itself — they improve the critic, then the actor follows gradients computed from that critic. The policy still climbs the value landscape one small step at a time, blind to the broader geometry of which actions are similar in value.

Prior work: TD3, SAC, and other actor-critics learn value functions and update policies via local gradients. Representation learning (e.g., contrastive methods) improves state embeddings but doesn’t reshape how actions are selected. The policy update remains gradient-based, local, and unaware of value similarity across the action space.

The gap: If two actions have similar Q-values, shouldn’t the policy treat them as interchangeable directions to explore? Current methods don’t exploit this structure. The value function is a black box that outputs scalars; the policy never sees which actions cluster together in value.

Problem: Policy updates ignore value geometry
   |
   v
Assumption: Actions with similar Q-values should have similar embeddings
   |
   v
Method: Learn joint state-action space where cosine similarity ~ Q-value similarity
        Use this geometry to sample actions and weight policy updates
   |
   v
Evidence: Outperforms TD3/SAC on MuJoCo benchmarks (Humanoid, Ant)
   |
   v
Conclusion: Value geometry can guide policy beyond local gradients

The Increment

One sentence: Before, policies followed value gradients blindly; after, they navigate a learned geometry where similar values point in similar directions.

Core Mechanism

SAVGO learns a joint embedding space for (state, action) pairs. The embedding network is trained so that pairs with similar Q-values have high cosine similarity, while pairs with different Q-values point in different directions. This creates a geometric structure: the embedding space is a compass where directions encode value.

At each policy update, SAVGO samples candidate actions from the current policy. It embeds each candidate, computes pairwise cosine similarities, and uses these similarities as weights in a kernel. High-value actions (identified by the critic) that are also geometrically similar to other high-value actions get amplified. The policy update becomes a weighted combination of gradients, where the weights come from the learned geometry, not just the critic’s scalar outputs.

The training loop has three coupled objectives: (1) critic learns Q-values via TD error, (2) embedding network learns to align cosine similarity with Q-value similarity via a contrastive loss, (3) policy maximizes expected return weighted by the similarity kernel. All three are trained off-policy from a replay buffer, preserving scalability.

Replay Buffer
     |
     v
Sample (s, a, r, s')
     |
     +---> Critic: Q(s,a) via TD loss
     |
     +---> Embedding: phi(s,a) via contrastive loss
     |          (similar Q -> similar phi)
     |
     +---> Policy: pi(s) samples actions a1...aN
                   Compute K(ai, aj) = cosine(phi(s,ai), phi(s,aj))
                   Weight policy gradient by K and Q
                   Update toward high-value, geometrically-consistent actions

Think of it like a hiking trail network. Traditional RL is a hiker who only sees the slope under their feet (local gradient). SAVGO builds a map where trails with similar elevation are drawn in similar colors. When deciding where to step next, the hiker samples a few nearby trails, checks their colors (embeddings), and preferentially moves toward trails that are both high-elevation (high Q) and color-consistent with other good trails (high cosine similarity). The map (embedding space) and the elevation data (Q-values) are learned together, so the colors stay calibrated to actual elevation.

Key Concepts

• Cosine similarity as value proxy: In high-dimensional spaces, distance metrics are noisy. Cosine similarity measures angle, not magnitude — two vectors can be far apart in Euclidean distance but point in the same direction. SAVGO bets that if two (state, action) pairs have similar Q-values, their embeddings should point in similar directions, regardless of magnitude. This is enforced via a contrastive loss: pairs with Q-values within a threshold are pulled to have cosine similarity near +1, pairs with different Q-values are pushed toward 0 or -1. The result: the embedding space becomes a directional encoding of value, where “north” means high-value and “south” means low-value, relative to the current state.

• Similarity kernel for policy updates: Instead of updating the policy purely via Q-gradients, SAVGO samples N actions from the current policy, embeds them, and computes an N×N similarity matrix K where K[i,j] = cosine(phi(s, a_i), phi(s, a_j)). This kernel weights the policy gradient: actions that are both high-Q and geometrically similar to other high-Q actions get more weight. It’s like voting where your vote counts more if you agree with other informed voters. This breaks the policy out of local gradient traps — if the current action is mediocre but there’s a cluster of high-value actions nearby in embedding space, the kernel pulls the policy toward that cluster even if the direct gradient is weak.

• Joint embedding space: Unlike methods that embed states and actions separately, SAVGO embeds (state, action) pairs into a single space. This is crucial: the same action in different states should have different embeddings because its value depends on context. The joint space captures “action a in state s” as a single geometric object. Training this space requires careful balancing — the contrastive loss must see enough diverse (s, a) pairs to learn generalizable structure, but the replay buffer is finite. SAVGO addresses this by sampling multiple actions per state during training, effectively augmenting the data.

Framework Shift

Before (TD3/SAC):                    After (SAVGO):

State s                              State s
  |                                    |
  v                                    v
Policy pi(s) -> action a          Policy pi(s) -> sample a1...aN
  |                                    |
  v                                    +-> Embed phi(s, a1)...phi(s, aN)
Critic Q(s,a) -> scalar                |
  |                                    +-> Compute similarity K(ai, aj)
  v                                    |
Gradient: dQ/da                        v
  |                                  Critic Q(s, ai) for each ai
  v                                    |
Update pi toward higher Q              v
                                     Weighted gradient: sum_i K(ai, aj) * Q(s, ai)
                                       |
                                       v
                                     Update pi toward high-Q, geometrically-consistent region

From scalar-driven local search to geometry-aware global navigation, the core shift is treating value as a directional field rather than a point-wise function.

Expert Assessment

Problem choice: Real gap. Actor-critic methods do get stuck in local optima, especially in high-dimensional action spaces. Representation learning has mostly focused on states; extending it to guide action selection is underexplored. The problem sits at the intersection of representation learning and policy optimization, which is timely given recent interest in self-supervised RL.

Method maturity: Clever insight, but execution is complex. The contrastive loss requires hyperparameter tuning (similarity threshold, temperature), and the similarity kernel adds computational overhead (N² similarity computations per update). The paper doesn’t discuss failure modes — what happens if the embedding space collapses or if the kernel amplifies noise? A simpler baseline would be to just use Q-value rankings to weight actions, without learning embeddings. The paper doesn’t compare against this.

Experimental integrity: Baselines are standard (TD3, SAC), but the improvements are modest (10-20% on some tasks, negligible on others). The paper cherry-picks Humanoid and Ant, which are known to have multimodal value landscapes — exactly where geometry should help. What about simpler tasks like Hopper or Walker, where local gradients suffice? Ablations are present but shallow: they remove components one at a time without probing why the geometry helps. No analysis of embedding space structure (e.g., t-SNE plots, similarity distributions).

Writing quality: The method section is dense and assumes familiarity with contrastive learning. The intuition for why cosine similarity is the right metric is buried in a footnote. The related work section name-drops representation learning papers but doesn’t explain how they differ from SAVGO. The results section would benefit from error bars and statistical tests — some improvements look within noise. The ablation on similarity threshold is interesting but underdeveloped.

Verdict: Weak accept — the core idea (value geometry for policy updates) is novel and the results show promise, but the paper needs clearer exposition, deeper ablations, and more honest discussion of when the method helps vs. when it doesn’t.

Takeaways

Directional encoding of value: If you’re building a critic, consider outputting embeddings alongside Q-values. Train the embeddings so similar values point in similar directions. This gives you a richer signal than scalars alone — you can cluster actions, interpolate between high-value regions, or detect when the value landscape is multimodal.

Kernel-weighted policy updates: Instead of following the gradient of a single action, sample multiple actions and weight their gradients by a similarity kernel. This is a general trick: replace point estimates with distributions, then use structure (similarity, uncertainty, diversity) to weight the update. Applicable beyond RL — think meta-learning, active learning, or any setting where you have multiple candidates and want to aggregate them intelligently.

When geometry helps: SAVGO shines when the value landscape has multiple peaks or plateaus — situations where local gradients are weak or misleading. If your domain has this structure (e.g., robotic manipulation with multiple grasp strategies, game AI with multiple viable tactics), geometry-aware methods are worth exploring. If the landscape is smooth and unimodal, stick to simpler methods.