SAVGO: Learning State–Action Value Geometry with Cosine Similarity for Continuous Control
arXiv:2605.00787v1 [cs.LG] 1 May 2026
Stavros Orfanoudakis , Pedro P. Vergara Delft University of Technology, The Netherlands {s.orfanoudakis, p.p.vergarabarrios}@tudelft.nl
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 geometryconsistent 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.
1
Introduction
Deep Reinforcement Learning (RL) has demonstrated remarkable success across a wide range of domains, from Atari games to continuous-control robotics [Mnih et al., 2015; Schulman et al., 2017]. However, in realistic highdimensional environments, training is still frequently limited by sample inefficiency and instability. Off-policy actorcritic methods, such as Soft Actor Critic (SAC) and Twin Delayed Deep Deterministic Policy Gradient (TD3), are strong RL algorithms for continuous control due to their scalability and robustness [Haarnoja et al., 2018; Fujimoto et al., 2018]. More recent algorithms further improve sample efficiency through critic ensembles and distributional objectives [Kuznetsov et al., 2020; Chen et al., 2021]. Even with these
improvements, performance often remains highly sensitive to representation quality and hyperparameter choices, while function approximation can obscure or distort the local decision structure needed for reliable policy improvement. A prominent line of work addresses these limitations by improving representations learned from observations. In pixel-based control, data augmentation and contrastive objectives have substantially boosted data efficiency [Yarats et al., 2021; Laskin et al., 2020; Chen et al., 2020], while selfpredictive objectives leverage temporal structure to regularize latent dynamics [Schwarzer et al., 2021]. Latent-dynamics auxiliary objectives have also been studied through reward and transition prediction in compressed state spaces [Gelada et al., 2019]. More recently, return and value aware auxiliary tasks aim to align features with decision-making signals rather than purely visual invariance [Liu et al., 2021; Yue et al., 2023]. State–action representation learning has also been used to improve actor–critic learning in both online and offline continuous-control settings [Fujimoto et al., 2023], suggesting that explicitly modeling interactions between states and actions can provide useful structure for value estimation and policy learning. Related work has also explored value-consistent and reward-aware representation objectives that more directly couple learned features to downstream control performance [Yue et al., 2023; Tse et al., 2025]. Despite these advances, most methods still treat representation learning as an auxiliary component, where the policy improvement step typically remains a local gradient update that does not explicitly exploit the geometry of the learned latent space. Distance and similarity-based approaches offer a complementary perspective by defining behavioral notions of similarity and learning representations that preserve them. Bisimulation-inspired objectives relate latent distances to control-relevant equivalence classes, yielding more robust representations [Zhang et al., 2021; Kemertas and Aumentado-Armstrong, 2021]. Sampling-based statesimilarity objectives further improve scalability by enabling direct optimization from experience [Castro et al., 2021]. More recent work has extended this direction through actionbased and controllability-oriented representations, predictive or variational bisimulation metrics, and large-scale analyses of behavioral metric learning under distractors [Liang et al., 2026; Freed et al., 2026]. In parallel, episodic memory
1.0 0.8
1
0.6
0
0.4
Policy update steps Sampled actions (final step) Weighted update direction
1 2
tinguish between similar and dissimilar pairs based on their utility difference, rather than using representation learning only as an auxiliary regularizer. • Deriving a geometry-aware policy optimization operator that performs similarity-weighted aggregation over candidate actions. • Empirically evaluating SAVGO on continuous-control benchmarks, including comparisons to established RL baselines and targeted ablations isolating the effects of the value-geometry objective, kernel weighting, and key temperature/curvature hyperparameters.
Q(st , at )
Action Dim. 2
2
2
1
0
Action Dim. 1
1
0.2 2
0.0
Figure 1: Policy improvement with value-aware geometry for a single state. Colored contours depict the critic landscape Q(st , ·). Candidate actions (colored dots) are sampled and reweighted using the learned cosine-similarity geometry, resulting in a kernel-aggregated update direction (green) and stable policy update steps (black) that lead to higher-value regions beyond the pointwise gradient.
and nearest-neighbor methods exploit similarity to accelerate value estimation and improve learning dynamics [Pritzel et al., 2017; Shen and Yang, 2021]. This idea has been extended in continuous control by using neighboring state–action structure, nonparametric value estimates, or action-value discrepancy information to improve policy exploitation [Gao et al., 2026]. Despite these advances, a key limitation remains, as distances are primarily used to regularize representations or shape value targets and only rarely serve as first-class operators for policy improvement in the action space. As a result, policy updates remain largely local, gradient-driven, and weakly informed by the global structure of the learned value function, contributing to instability and inefficient exploration in high-dimensional action spaces. To address this limitation, we propose State–Action Value Geometry Optimization (SAVGO), a geometry-aware RL algorithm that explicitly integrates value-based similarity into the policy update itself. SAVGO learns a joint state–action embedding space in which pairs with similar state-action Qvalues exhibit high cosine similarity, while dissimilar pairs are mapped to different directions. In parallel, this learned geometry can be used to create a similarity kernel over K candidate actions sampled at each update, enabling the policy to be directly pulled toward regions associated with higher values (Figure 1). In this way, SAVGO tightly couples representation learning, value estimation, and policy improvement through a single geometry-consistent objective, while retaining the scalability and practicality of replay-based offpolicy actor–critic training. We evaluate SAVGO on standard continuous-control benchmarks from MuJoCo, comparing against strong RL baselines and demonstrating consistent gains on challenging, high-dimensional tasks, such as Humanoid. In addition, we conduct targeted ablation studies to isolate the contributions of the value-geometry objective, similarity-based kernel weighting, and key hyperparameters. The main contributions of this work can be summarized as: • Proposing SAVGO an RL algorithm that learns a cosinesimilarity geometry over state-action embeddings to dis-
2
Related Work
2.1
Representation Learning for RL
Learning task-relevant representations has long been recognized as a key ingredient for stable and sample-efficient deep RL. For visual control, contrastive objectives combined with augmentation have been particularly effective. SimCLR popularized simple instance discrimination in vision [Chen et al., 2020], and CURL adapted contrastive learning to off-policy RL by training an encoder jointly with the control objective [Laskin et al., 2020]. Related augmentation-driven regularization methods, such as DrQ, further improved data efficiency from pixels [Yarats et al., 2021]. Beyond contrastive learning, predictive self-supervision has also been employed to stabilize representation learning, e.g., by predicting future latent features with target encoders [Schwarzer et al., 2021]. Several works have advocated for decoupling representation learning from policy optimization using strong unsupervised objectives [Stooke et al., 2021]. A central limitation of many such objectives is that they are often generic and do not explicitly encode controlrelevant structure. This has motivated value-aware representation learning methods, including the VIP perspective [Dabney et al., 2021], return-based contrastive learning [Liu et al., 2021], and value-consistent representation learning [Yue et al., 2023]. Representation learning in low-dimensional continuous-control domains has also been used, where observations are already compact, but the interaction between states and actions can still benefit from learned structure. In particular, SALE learns state–action embeddings for deep RL and integrates them into TD3 to form TD7, demonstrating that state–action representation learning can substantially improve both online and offline continuous-control performance [Fujimoto et al., 2023]. This line of work is closely related to SAVGO because it highlights the importance of modeling state–action interactions rather than learning state-only features. SAVGO follows the same overarching goal of making representations control-relevant, but differs in how the learned structure is used. Rather than treating representation learning as an auxiliary regularizer or using it only to improve value targets, a cosine-similarity geometry over state–action embeddings is learned and then reused directly to construct similarity-weighted policy updates.
2.2
Distances, Metrics, and State Abstraction
Distances over states (and state–action pairs) provide a principled route to representation learning and state abstraction,
as behavioral similarity can be formalized by requiring similar rewards and similar future evolution. This idea is captured by bisimulation metrics, in which states are considered close when their immediate rewards and transition dynamics match under an appropriate coupling. DeepMDP connects latent dynamics modeling objectives to bisimulationstyle guarantees by learning representations that support reward and transition prediction in latent space [Gelada et al., 2019]. Deep Bisimulation for Control (DBC) minimizes a bisimulation metric loss to learn invariant representations and improve robustness to distractors [Zhang et al., 2021]. MICo introduces a practical sampling-based behavioral distance that scales to deep RL and uses it to shape representations [Castro et al., 2021]. Furthermore, action-bisimulation encoding can learn controllability-oriented representations through a recursive invariance constraint, allowing representations to capture long-horizon action-relevant structure without relying directly on reward supervision [Rudolph et al., 2024]. Predictive and variational bisimulation objectives have also been explored as mechanisms for learning taskrelevant or distractor-robust latent spaces [Liang et al., 2026; Freed et al., 2026]. In parallel, large-scale empirical studies of behavioral metric learning have examined how well metric-based objectives filter distractors and have highlighted the importance of design choices in bridging theoretical guarantees and practical deep RL performance [Luo et al., 2025]. Despite their appeal, bisimulation-driven objectives can be brittle in practice, exhibiting robustness issues and representation pathologies such as collapse or exploding norms [Kemertas and Aumentado-Armstrong, 2021]. Additional pitfalls have been reported in offline settings, where missing transitions can bias metric estimation and degrade learned distances [Zang et al., 2023]. In contrast to these methods, SAVGO does not aim to learn a full behavioral metric for state abstraction, instead, a value-aware cosine geometry over state– action embeddings is learned.
2.3
Kernel and Nearest-Neighbor Value Search
A separate but closely related field uses similarity directly in value estimation. Kernel-Based Reinforcement Learning approximates Bellman backups using kernel regression in continuous state spaces [Ormoneit and Sen, 2002]. Neural Episodic Control stores a memory of past embeddings with fast-updated value estimates, effectively implementing a differentiable nearest-neighbor value function [Pritzel et al., 2017]. Nearest Neighbor Actor-Critic proposes a theoretically grounded nearest-neighbor module that can replace or augment parametric value networks [Shen and Yang, 2021], whereas Instant Retrospect Action (IRA), uses Qrepresentation discrepancy evolution to learn discriminative representations for neighboring state–action pairs, together with k-nearest-neighbor action-value estimates [Gao et al., 2026]. Collectively, these methods demonstrate that learned similarity can support effective bootstrapping and credit assignment. Building on this insight, SAVGO extends the role of similarity beyond value estimation to policy optimization. Rather than using distances solely to interpolate or retrieve value estimates, SAVGO learns a control-relevant state–action geometry and employs it directly as the driving
signal for policy improvement.
3
State–Action Value Geometry Optimization
Building on the limitations of existing representationlearning, similarity-based, and kernel-driven approaches discussed above, SAVGO is introduced as a geometry-aware offpolicy actor–critic framework in which learned value-based similarity is made directly actionable during policy improvement. SAVGO is composed of two tightly coupled components: (i) a state–action value geometry learned in a normalized embedding space using cosine similarity, and (ii) a geometry-aware policy improvement operator that aggregates a set of candidate actions through similarity-dependent weights. In the remainder of this section, the value-geometry learning objective and the resulting similarity-based policy improvement operator are formally derived.
3.1
Preliminaries: Off-Policy Actor–Critic
A discounted Markov decision process (MDP) M = (S, A, P, r, γ) is considered, where S and A denote the state and action spaces, P (st+1 | st , at ) denotes the transition probabilities, r : S × A → R is the reward function, and γ ∈ [0, 1) is the discount factor. A stochastic policy πθ (a | s), parameterized by θ is used, and transitions are stored in a replay buffer D. The training loop is aligned with SAC [Haarnoja et al., 2018], in which, two critics Qϕ1 , Qϕ2 are trained using slowly soft-updated target networks Qϕ̃1 , Qϕ̃2 . Following common practice, the minimum over the two target critics is used to obtain conservative value estimates [Fujimoto et al., 2018], i.e., Q(st , at ) ≜ minm∈1,2 Qϕ̃m (si , ai ). Accordingly, the critics are updated as in SAC. Given a minibatch {(si , ai , ri , si+1 , di )}B i=1 ∼ D of size B, where di refers to a terminal state flag, a next action ai+1 ∼ πθ (· | si+1 ) is sampled and the temporal difference (TD) target is: yi = ri + γ di Q(si+1 , ai+1 ) − η log πθ (ai+1 | si+1 ) (1) where η ≥ 0 denotes the entropy temperature [Haarnoja et al., 2018]. Each critic is then updated by minimizing B
LQm (ϕm ) =
2 1 X Qϕm (si , ai ) − yi , m ∈ {1, 2}. (2) B i=1
All remaining SAC components (reparameterized sampling, temperature tuning, and target updates) are adopted from [Haarnoja et al., 2018] and are not restated.
3.2
Value-Aware State–Action Geometry Learning
A value-aware state–action geometry is learned to capture control-relevant relationships between pairs of state–action tuples (si , ai ) and (sj , aj ). To this end, a representation encoder function zψ : S × A → Rd , parameterized by ψ, is trained such that similarity in the embedding space reflects similarity in expected utility. For notational convenience, we denote ζi ≜ zψ (si , ai ) and ζj ≜ zψ (sj , aj ). Similarity between embeddings is measured using cosine similarity, cos(ζi , ζj ) =
ζi⊤ ζj ∥ζi ∥2 ∥ζj ∥2
∈ [−1, 1].
(3)
0.0
Anchor action embedding Candidate embeddings
= 1.5 =2 = 2.5
0.0
0.2
0.4
0.6
Normalized Value Gap
i, j
0.8
1.0
Cosine similarity is adopted due to its bounded range, scale invariance, and compatibility with normalized embeddings, which together yield a stable and interpretable similarity measure on the unit hypersphere. To ensure that embedding similarities reflect value proximity between state–action pairs, the representation function zψ is trained using value-based supervision. For a pair of state–action tuples with action-value estimates Q(si , ai ) and Q(sj , aj ), a normalized value gap ∆i,j ∈ [0, 1] is defined as |Q(si , ai ) − Q(sj , aj )| ∆i,j = clip , 0, 1 , (4) β where β > 0 is a robust scale parameter, such as an exponential moving average of observed value differences. This normalization ensures invariance to reward scaling and prevents large value discrepancies from dominating the learned value geometry. With ∆i,j then providing a normalized measure of value dissimilarity on [0, 1], it can be directly mapped to a target cosine similarity so that the supervision signal matches the range and semantics of cosine similarity, i.e., +1 for highly similar pairs and −1 for strongly dissimilar pairs. This is achieved with the bounded curvature transformation ∈ [−1, 1],
(5)
with λ > 0 controlling the curvature of this mapping and thus the resolution assigned to different gap regimes. As illustrated in Figure 2, smaller values of λ allocate more dynamic range to small gaps, sharpening discrimination among nearoptimal actions, while larger values compress small gaps and yield smoother geometries by reducing sensitivity to minor value differences. Finally, the encoder parameters ψ are optimized by matching predicted similarities cos(ζi , ζj ) to their targets Yi,j , using common regression losses ℓ(·), such as the Huber or ℓ2 loss. Lz (ψ) ≜ ℓ(cos(ζi , ζj ) − Yi,j )
∀(i, j) ∈ B.
0.8
(6)
Through this construction, the learned embedding space becomes explicitly control-relevant, because state–action pairs with similar utility are encouraged to occupy nearby directions on the unit hypersphere, while pairs with dissimilar utility are pushed apart. This value-aware geometry forms the foundation for the similarity-based policy improvement operator introduced in the following section.
0.6
0 1
0.4 0.2 1
Dim.
Figure 2: Target cosine similarity Yi,j as a function of the normalized value gap ∆i,j for different curvature parameters λ.
Yi,j = 1 − 2 (∆i,j )λ
1.0
1
0.5 1.0
Weighted update direction Gradient of update
Q(s, a)
= 0.25 = 0.5 =1
0.5
Dim. 3
Target Cos. Sim. Yi, j
1.0
0
2
1
1
0
Dim.
1
0.0
1
Figure 3: Unit-sphere view of the geometry-aware actor update in the embedding space for state st . The anchor embedding z(st , ât ) (black) and candidate embeddings z(st , akt ) (colored by Q(st , akt )) lie on the learned embedding space. Cosine similarities to the anchor define kernel weights wk , yielding a similarity-weighted update direction (green) that steers the actor toward high Q-value actions.
3.3
Policy Improvement in Learned Geometry
Policy improvement is performed by leveraging the valueaware geometry from Section 3.2 to construct similarityweighted updates in the action space. An abstract view is provided in Figure 1, while Figure 3 illustrates the same mechanism on the learned unit-sphere embedding space. In detail, for a given state st , an anchor action ât ∼ πθ (· | st ) is first sampled from the current policy and mapped to its embedding ζ̂t ≜ zψ (st , ât ) (black vector in Figure 3). In addition, a set of K candidate actions {akt }K k=1 is drawn from a proposal distribution (e.g., the policy or a mixture for broader coverage), yielding candidate embeddings ζtk ≜ zψ (st , akt ) (colored vectors). Each candidate pair (st , akt ) is scored by a conservative critic estimate, qtk = minm∈1,2 Qϕ̃m (st , akt ), which is showed by the vector coloring in Figure 3. The influence of each candidate on the update is determined by its geometric proximity to the anchor in the learned cosine space. Cosine similarities cos(ζ̂t , ζtk ) ≜ cos(zψ (st , ât ), zψ (st , akt )) quantify how close a candidate lies to the anchor direction on the unit sphere. As illustrated in Figure 3, these similarities are later converted into kernel weights, so that nearby candidates contribute more strongly to the aggregated update direction than distant ones. For additional stability, similarity computations can be performed with a slowly-updated target encoder zψ̃ . The pairwise similarities are then transformed into a probability simplex through a temperature-controlled softmax, yielding a similarity kernel over candidate actions, exp(cos(ζ̂t , ζtk )/ρ) ε wk = (1 − ε) PK + , j K j=1 exp(cos(ζ̂t , ζt )/ρ)
(7)
where the temperature ρ > 0 controls the sharpness of the kernel and ε ∈ [0, 1] prevents weight collapse. In our implementation, ρ is adapted over training using a cosine schedule from ρmax to ρmin , producing smoother weighting early on while allowing sharper selection later when the representation encoder is more informative.
Algorithm 1 State-Action Value Geometry Optimization 1: Initialize actor πθ , critics Qϕm , encoder zψ 2: Initialize target net. ψ̃ ← ψ and ϕ̃m ← ϕm , m ∈ {1, 2} 3: for t = 1, 2, . . . do 4: Sample at ∼ πθ (·|st ) 5: Step env and store (st , at , rt , st+1 , dt ) ∈ D 6: Sample minibatch {(si , ai , ri , si+1 , di )}B i=1 ∼ D 7: Critics Update 8: Compute targets yi (Eq. 1) 9: Update ϕm by ∇ϕm LQm (Eq. 2) for m ∈ {1, 2} 10: Value-geometry Encoder Update 11: Form pairs (i, j) from the minibatch 12: Compute ∆ij (Eq. 4) and Yij (Eq. 5) 13: Update ψ by ∇ψ Lz (Eq. 6) 14: Actor Update (for each si ) 15: Sample anchor âi ∼ πθ (·|si ) 16: Sample candidates {aki }K k=1 ∼ µ(·|si ) 17: Compute kernel weights wk (Eq. 7) using zψ̃
b i , âi ) (Eq. 8) 18: Compute Q(s 19: Update θ by ∇θ Lπ (Eq. 9) 20: Target Networks Update 21: ϕ̃m ← τ ϕ̃m + (1 − τ ) ϕm 22: ψ̃ ← τ ψ̃ + (1 − τ ) ψ 23: end for
for m ∈ {1, 2}
4
Using the weights {wk }K k=1 , a similarity-weighted value estimate is computed as, Q̂(st , ât ) ≜
K X
wk qtk
(8)
k=1
This can be viewed as a local, similarity-weighted average of critic values around ât in the learned geometry (Figure 1 and 3). This aggregation smooths the update signal by combining multiple candidates rather than relying on a single point estimate, compared to classic algorithms such as SAC and TD3. Finally, the policy parameters θ are updated by minimizing h i Lπ (θ) ≜ Est ,ât ∼πθ (·|st ) η log πθ (ât | st ) − Q̂(st , ât ) , (9) where η controls the strength of entropy regularization, thereby promoting exploration in maximum-entropy stochastic actor–critic methods, such as SAC. From this perspective, policy improvement is no longer driven by a pointwise gradient at a single action, but by a similarity-weighted aggregation over a set of candidate actions. This approach is related to weighted policy improvement and advantage-weighted regression methods shown in [Peters et al., 2010; Abdolmaleki et al., 2018], but differs fundamentally in that the weights are determined by learned geometric similarity in state–action space, rather than solely by action-value magnitude.
3.4
The method follows a standard replay-based off-policy actor–critic loop, while coupling critic learning. At each environment interaction step, an action is sampled from the current stochastic policy and executed (lines 4–5), and the resulting transition is stored in the replay buffer (line 5). A minibatch is then drawn from the buffer (line 6) and used for the parameter updates. The critics are first trained with a SAC-style TD target (lines 7–9), yielding conservative estimates via the minimum of two critics. The valuegeometry encoder is updated using paired samples from the same minibatch (lines 10–13) by regressing cosine similarities toward bounded targets derived from value gaps, so that nearby embedding directions reflect similar expected utility. The geometry-aware policy improvement is performed (lines 14–19) by sampling, for each state si , an anchor action âi ∼ πθ (· | si ) (line 15) and K candidate actions (line 16), converting anchor–candidate cosine similarities into a soft kernel (line 17), forming a kernel-aggregated value estimate b i , âi ) (line 18), and updating the actor to increase this Q(s similarity-weighted value signal (line 19) rather than relying on a single pointwise estimate. Finally, target networks are updated via Polyak averaging (soft-updates) with rate τ (lines 20–22) to stabilize value and similarity estimations.
Overall Algorithm and Training Procedure
Having introduced the value-aware geometry objective (Section 3.2) and the geometry-aware policy improvement operator (Section 3.3), the complete training procedure is summarized in Algorithm 1.
Experimental Evaluation
This section evaluates SAVGO on continuous-control benchmarks, comparing against strong baselines and analyzing the effect of key design choices through targeted ablations.
4.1
Experimental Setup
Experiments are conducted on the standard MuJoCo [Towers et al., 2025] continuous-control suite (v5) under a fixed budget of 1M environment steps per run, spanning tasks of varying action dimensionality and difficulty. SAVGO1 is compared against widely used continuous-control baselines, including Proximal Policy Optimization (PPO) [Schulman et al., 2017], TD3 [Fujimoto et al., 2018], SAC [Haarnoja et al., 2018], and TQC [Kuznetsov et al., 2020], which collectively represent strong on-policy and off-policy training pipelines and are sufficient to characterize performance across the suite. Distance-based representation methods such as MICo [Castro et al., 2021] and DBC [Zhang et al., 2021] are also closely related. However, in our adaptation of the original codebases to the present training stack, MICo and DBC methods did not achieve competitive performance on the more complex MuJoCo tasks and are therefore not included as primary baselines. To ensure a fair comparison, all methods are trained under identical environment settings and step budgets and are evaluated with the same protocol. Implementations are based on Stable Baselines 3 [Raffin et al., 2021]. Results are aggregated over multiple random seeds, and performance is reported as mean and standard deviation. Observation normalization is enabled across tasks to reduce sensitivity to scale and improve training stability. Two-layer actor, critic, and 1 The implementation code is available at https://github.com/ StavrosOrf/DistanceRL.
PPO
Walker2d
1e3
Episode Return
6.0 1e3 4.5 3.0 1.5 0.0 0.00
0.25
0.50
SAC
0.75
Training Steps
7.5 6.0 4.5 3.0 1.5 0.0 1.00 0.00 1e6
SAVGO (ours)
Humanoid
0.25
0.50
TD3
1e3
0.75
Training Steps
1.00 1e6
5 4 3 2 1 0 0.00
TQC
Ant
0.25
0.50
HalfCheetah
1e3
0.75
Training Steps
12.5 10.0 7.5 5.0 2.5 0.0 1.00 0.00 1e6
0.25
0.50
0.75
Training Steps
1.00 1e6
Figure 4: Training curves over 1M environment steps on representative MuJoCo (v5) tasks. Solid lines show the mean evaluation return over seeds, and shaded regions indicate the standard deviation. The baseline implementations from Stable Baselines 3 are used. Table 1: Best evaluation performance (maximum over training up to 1,000,000 steps) across MuJoCo v5 environments. Entries report mean ± standard deviation over 10 seeds. Bold indicates the best-performing algorithm per environment. The Total Rewards row sums environmentwise means, its ± value is the sum of standard deviations across environments (for a rough aggregate variability indicator). Environment (v5)
TD3
SAC
TQC
Ant HalfCheetah Hopper Humanoid Inverted2DPendulum InvertedPendulum Reacher Swimmer Walker2d
12 ± 72 2363 ± 1022 3527 ± 284 554 ± 62 9350 ± 16 1000 ± 0 −3 ± 0 114 ± 25 4321 ± 635
PPO
4366 ± 1296 10 259 ± 838 3584 ± 77 799 ± 266 9358 ± 2 1000 ± 0 −3 ± 0 102 ± 46 4503 ± 752
4086 ± 789 10 388 ± 1069 3371 ± 84 5351 ± 75 9360 ± 0 1000 ± 0 −2 ± 0 77 ± 10 4693 ± 572
3041 ± 668 12 114 ± 1006 3571 ± 65 6330 ± 227 9360 ± 0 1000 ± 0 −2 ± 0 106 ± 25 4952 ± 374
4653 ± 10 545 ± 3573 ± 7687 ± 9360 ± 1000 ± −3 ± 88 ± 5626 ±
Total Rewards
21 238 ± 2116
33 968 ± 3277
38 324 ± 2599
40 472 ± 2365
42 529 ± 2059
encoder networks with 256 nodes per layer are used across algorithms. For SAVGO, the in-state candidate set size is set to K ∈ {64, . . . , 256} depending on action dimension, with kernel smoothing ε = 0.05. SAVGO’s encoder network consists of two layers, each with 256 nodes. The weight smoothing temperature ρ is cosine annealed over the first 200,000 training steps with ρmax = 0.75 and ρmin = 0.05. Detailed hyperparameter settings for each environment and baseline can be found in the supplementary material. All the experiments were run on the “Anonymous” supercomputer, consisting of Nvidia V100 GPUs and Intel Xeon 48-core CPU nodes.
4.2
Main Results on MuJoCo
Figure 4 reports evaluation episode return over 1M environment training steps averaged over 10 seeds. Across environments, faster early learning and smoother convergence were observed for SAVGO relative to PPO, TD3, SAC, and TQC, indicating improved sample-efficiency and a more stable policy-improvement signal. The clearest gains appeared in the higher-dimensional tasks. On Walker2d, SAVGO learned rapidly and stabilized at a higher return band than the baselines, while on Humanoid, it was the only method to consistently achieve strong improvement over training, whereas the others progressed more slowly and/or plateaued earlier. Competitive behavior was also observed on Ant and HalfCheetah, where SAVGO tracked the top-performing baselines while exhibiting reduced variance across training. Overall runtimes stayed practical even on the hardest MuJoCo tasks (1M training steps): PPO < 2h, TD3 ≈ 2h, SAC ≈ 2.5h, TQC ≈ 4.5h,
SAVGO 235 616 114 319 0 0 0 28 747
and SAVGO ≈ 8h. SAVGO was the slowest due to extra critic and encoder evaluations for similarity weighting, while PPO was the fastest, followed by TD3 and SAC. Across MuJoCo benchmarks, the best evaluation returns within the 1M-step budget (Table 1) show that SAVGO is consistently among the strongest methods and achieves the best overall performance, with the largest gains concentrated on challenging high-dimensional tasks such as Humanoid and Walker2d, where policy improvement is particularly sensitive to critic noise. On mid-complexity environments (e.g., Ant), improvements remain clear but smaller, while on HalfCheetah performance is competitive rather than best. As expected, on simpler tasks where near-ceiling behavior is reached, differences largely vanish. Overall, these trends suggest that similarity-weighted updates are most beneficial when action ranking is difficult and local gradients are less reliable.
4.3
Ablation Studies
Ablation studies can quantify the contribution and sensitivity of key SAVGO design choices, with a focus on the representation curvature parameter λ (Eq. 5) and the number of candidate actions K used in the kernelized update (Eq. 7). Evaluations were performed on HalfCheetah and Humanoid, which were selected to contrast a lower-dimensional, fast-learning locomotion task, and a high-dimensional setting with higher critic noise and action-ranking difficulty. Each configuration was trained for 1M environment steps over 5 random seeds, and performance was summarized by the maximum episode return achieved within the budget.
HalfCheetah
9 1e3
Humanoid
8 7 6 1
1.5
2
2.5
5
.25 .5
1
1.5
2
2.5
Figure 5: Sensitivity to the representation curvature parameter λ. Points denote per-seed maxima and error bars show mean±std of the maximum evaluation return within 1M steps over 5 seeds.
Effect of representation curvature λ. A focused ablation isolates the effect of the representation curvature parameter λ in Eq. 5, while keeping the architecture, optimization settings, replay configuration, and all other SAVGO hyperparameters fixed. Figure 5 shows a strong dependence on curvature, with values close to 1 yielding the most reliable performance. Small curvature values (λ ≤ 0.5) increase seed variance, consistent with an overly sharp target mapping that over-weights small value gaps and destabilizes geometry learning (See Figure 2). Large curvature values (λ ≥ 2) reduce peak value gaps, suggesting that the mapping becomes too smooth to preserve informative action-ranking structure. Overall, λ ≈ 1–1.5 offers the best trade-off between stability and final performance across both tasks. Effect of sample candidates K. The ablation here studies the effect of the candidate set size K in the kernelized policy improvement operator (Eq. 7) by sweeping K ∈ {16, 64, 128, 256, 512} while holding all other hyperparameters fixed. Figure 6 indicates that strong performance is typically obtained with a moderate number of candidates. On HalfCheetah, returns remain largely unchanged for K ≥ 64, since adding more samples doesn’t benefit anymore once the local neighborhood is sufficiently covered. On Humanoid, smaller candidate sets (K = 16) reduce peak performance and increase variability, whereas intermediate values (K ≈ 128) yield the most consistent results. Increasing to K = 512 does not produce systematic gains, suggesting that additional candidates become redundant and mainly increase computation. Overall, K is best chosen to balance neighborhood coverage and efficiency, with larger action spaces typically requiring larger candidate sets. Design choice ablation. To quantify the contribution of each design choice in SAVGO and verify that the intended policy-improvement mechanism is realized in practice, a targeted ablation study is performed. Kernel sharpness is ablated by disabling the adaptive temperature schedule in Eq. 7 and keeping ρ fixed. Value-aware geometry learning is ablated by freezing the encoder zψ , i.e., stopping gradients from Eq. 6, which tests whether similarity must be learned to track value proximity rather than used as a fixed heuristic. Robust normalization is ablated by removing the adaptive scaling in Eq. 4 and using a fixed β, which tests sensitivity to critic’s value magnitude. Finally, a stress-test replaces the similarity kernel with uniform weights wk = 1/K to verify
Max Episode Return
Max Episode Return
12 1e3 11 10 9 8 7 6 .25 .5
12 1e3
HalfCheetah
9 1e3
11
8
10
7
9 8
Humanoid
16 64 128
256 K
512
6
16 64 128
256 K
512
Figure 6: Sensitivity to candidate set size K in the similarityweighted policy improvement operator. Points denote per-seed maxima and error bars show mean±std of the maximum evaluation return within 1M environment steps over 5 seeds. Table 2: Design choice ablation study on two MuJoCo environments. Each row removes or modifies one architectural component of the proposed method. Results are reported as mean evaluation return ± standard deviation over 5 seeds after 1M steps. Variant
HalfCheetah
Humanoid
SAVGO (full) w/o adaptive ρ w/o representation loss w/o adaptive β scaling uniform kernel weighting
10 546 ± 616 9184 ± 941 8466 ± 1247 4974 ± 830 0± 0
7687 ± 319 7559 ± 667 6241 ± 706 5765 ± 259 218 ± 21
that gains come from geometry-aware weighting rather than unweighted averaging. As anticipated, Table 2 shows that each component was necessary for strong performance on HalfCheetah and Humanoid after 1M steps (5 seeds). Fixing ρ or removing adaptive β scaling reduced returns, and freezing the encoder (removing Eq. 6) caused a larger drop, indicating that the value-aware geometry must be learned rather than treated as a fixed heuristic. The stress test collapsed performance (uniform weights), confirming that policy improvement strictly depends on similarity-based weighting.
5
Discussion
Computation and Complexity. Relative to standard offpolicy actor–critic methods such as SAC, SAVGO introduces extra per-update cost by sampling and scoring K candidate actions per state in order to form the similarity kernel (Eqs. 7– 8). Let B denote the minibatch size, CQ the cost of one critic forward pass, and Cz the cost of one encoder forward pass. Ignoring constants and minor bookkeeping, a SACstyle update requires O(B CQ ) critic evaluation for targets and losses (plus O(B) policy log-prob terms). In SAVGO, the critic update remains unchanged up to constants, but the policy improvement step additionally evaluates the critic and encoder on K candidates per state, yielding an overhead of O BK(CQ +Cz ) on top of the baseline critic update. Computing cosine similarities and the softmax weights adds only a few more arithmetic computations, which are typically dominated by network forward passes. Thus, SAVGO’s runtime scales approximately linearly in K, and the relative overhead is most pronounced when CQ and Cz are large (e.g., highdimensional tasks). Memory overhead remains modest, as
candidates and similarities are computed on-the-fly and the replay buffer continues to dominate storage. Assumptions, Practical Considerations, and Limitations. SAVGO assumes that the critic provides a reasonably consistent ranking over sampled candidate actions. Under severe critic noise, the resulting similarity weights can become unreliable even with temperature smoothing. The geometry objective likewise assumes that normalized value gaps yield stable supervision, motivating robust scaling (β) and curvature control (λ), since overly sharp mappings can amplify noise while overly smooth mappings can obscure action-ranking structure. Performance further depends on candidate coverage. In detail, small K may under-sample the local neighborhood, whereas large K increases computation with diminishing returns, suggesting that K should be tuned relative to action-space dimensionality. Empirical gains are smallest on easy tasks where all methods reach near-ceiling performance, and scaling to very high-dimensional actions may require a larger K with proportional cost. Finally, extending SAVGO to discrete-action domains is non-trivial, as learning a useful state–action geometry and designing effective candidate proposals over discontinuous actions likely requires different encoders and sampling strategies.
6
Conclusion
In this paper, SAVGO was introduced as a geometry-aware off-policy actor–critic method in which state–action similarity is made directly actionable during policy improvement. On the MuJoCo continuous-control suite under a fixed 1Mstep budget, consistent gains in sample efficiency and stability were observed, with the greatest improvements on challenging high-dimensional tasks such as Humanoid and Walker2d. Ablation results further verified that value-aware geometry learning, conservative candidate evaluation, adaptive kernel sharpness, and robust value-gap normalization each contribute materially to performance, and removing any component leads to clear degradations or failure cases. Future work will focus on more structured candidate proposal mechanisms for improved coverage in very high-dimensional action spaces, increased robustness under stronger stochasticity or partial observability, and extensions to discrete-action domains (e.g., Atari) through geometry-aware encoders and candidate selection tailored to discontinuous action sets.
Acknowledgment The study was funded by the DriVe2X research and innovation project from the European Commission with grant number 101056934. The authors acknowledge the use of computational resources of the DelftBlue supercomputer, provided by Delft High Performance Computing Centre (https://www.tudelft.nl/dhpc). This work also used the Dutch national e-infrastructure with the support of the SURF Cooperative, using grant no. EINF-5716.
References [Abdolmaleki et al., 2018] Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess,
and Martin Riedmiller. Maximum a posteriori policy optimisation. In International Conference on Learning Representations, 2018. [Castro et al., 2021] Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland. MICo: Improved representations via sampling-based state similarity for markov decision processes. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. [Chen et al., 2020] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, ICML’20. JMLR.org, 2020. [Chen et al., 2021] Xinyue Chen, Che Wang, Zijian Zhou, and Keith W. Ross. Randomized ensembled double qlearning: Learning fast without a model. In International Conference on Learning Representations, 2021. [Dabney et al., 2021] Will Dabney, André Barreto, Mark Rowland, Robert Dadashi, John Quan, Marc G. Bellemare, and David Silver. The value-improvement path: Towards better representations for reinforcement learning. volume 35, pages 7160–7168, May 2021. [Freed et al., 2026] Benjamin Freed, Roberto Calandra, Jeff Schneider, and Howie Choset. Distractor-robust reinforcement learning via variational bisimulation, 2026. [Fujimoto et al., 2018] Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 1582–1591. PMLR, 2018. [Fujimoto et al., 2023] Scott Fujimoto, Wei-Di Chang, Edward J. Smith, Shixiang Shane Gu, Doina Precup, and David Meger. For sale: state-action representation learning for deep reinforcement learning. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS ’23, Red Hook, NY, USA, 2023. Curran Associates Inc. [Gao et al., 2026] Gong Gao, Weidong Zhao, Xianhui Liu, and Ning Jia. Improving policy exploitation in online reinforcement learning with instant retrospect action. Neural Networks, 199:108667, 2026. [Gelada et al., 2019] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2170– 2179. PMLR, 09–15 Jun 2019. [Haarnoja et al., 2018] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic:
Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1861–1870. PMLR, 10–15 Jul 2018. [Kemertas and Aumentado-Armstrong, 2021] Mete Kemertas and Tristan Aumentado-Armstrong. Towards robust bisimulation metric learning. In Proceedings of the 35th International Conference on Neural Information Processing Systems, NIPS ’21, Red Hook, NY, USA, 2021. Curran Associates Inc. [Kuznetsov et al., 2020] Arsenii Kuznetsov, Pavel Shvechikov, Alexander Grishin, and Dmitry Vetrov. Controlling overestimation bias with truncated mixture of continuous distributional quantile critics. In Proceedings of the 37th International Conference on Machine Learning, ICML’20. JMLR.org, 2020. [Laskin et al., 2020] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representations for reinforcement learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5639– 5650. PMLR, 13–18 Jul 2020. [Liang et al., 2026] Dayang Liang, Ruihan Liu, Lipeng Wan, Yunlong Liu, and Bo An. Task-aware exploration via a predictive bisimulation metric, 2026. [Liu et al., 2021] Guoqing Liu, Chuheng Zhang, Li Zhao, Tao Qin, Jinhua Zhu, Li Jian, Nenghai Yu, and Tie-Yan Liu. Return-based contrastive representation learning for reinforcement learning. In Ninth International Conference on Learning Representations(ICLR), January 2021. [Luo et al., 2025] Ziyan Luo, Tianwei Ni, Pierre-Luc Bacon, Doina Precup, and Xujie Si. Understanding behavioral metric learning: A large-scale study on distracting reinforcement learning environments. CoRR, abs/2506.00563, 2025. [Mnih et al., 2015] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nat., 518(7540):529–533, 2015. [Ormoneit and Sen, 2002] Dirk Ormoneit and Śaunak Sen. Kernel-based reinforcement learning. Machine Learning, 49:161–178, 2002. [Peters et al., 2010] Jan Peters, Katharina Mülling, and Yasemin Altün. Relative entropy policy search. In Proceedings of the Twenty-Fourth AAAI Conference on Artificial Intelligence, AAAI’10, page 1607–1612. AAAI Press, 2010. [Pritzel et al., 2017] Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adrià Puigdomènech Badia, Oriol
Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 2827– 2836. PMLR, 06–11 Aug 2017. [Raffin et al., 2021] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1–8, 2021. [Rudolph et al., 2024] Max Rudolph, Caleb Chuck, Kevin Black, Misha Lvovsky, Scott Niekum, and Amy Zhang. Learning action-based representations using invariance. RLJ, 1:342–365, 2024. [Schulman et al., 2017] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. [Schwarzer et al., 2021] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with selfpredictive representations. 2021. [Shen and Yang, 2021] Junhong Shen and Lin F. Yang. Theoretically principled deep rl acceleration via nearest neighbor function approximation. volume 35, pages 9558–9566, May 2021. [Stooke et al., 2021] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9870– 9879. PMLR, 18–24 Jul 2021. [Towers et al., 2025] Mark Towers, Ariel Kwiatkowski, John U. Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulão, Kallinteris Andreas, Markus Krimmel, Arjun KG, Rodrigo De Lazcano Perez-Vicente, J K Terry, Andrea Pierré, Sander V Schulhoff, Jun Jet Tai, Hannah Tan, and Omar G. Younis. Gymnasium: A standard interface for reinforcement learning environments. In The Thirty-ninth Annual Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2025. [Tse et al., 2025] Hon Tik Tse, Siddarth Chandrasekar, and Marlos C. Machado. Reward-aware proto-representations in reinforcement learning. CoRR, abs/2505.16217, 2025. [Yarats et al., 2021] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021. [Yue et al., 2023] Yang Yue, Bingyi Kang, Zhongwen Xu, Gao Huang, and Shuicheng Yan. Value-consistent representation learning for data-efficient reinforcement learning. In Proceedings of the Thirty-Seventh AAAI Conference on Artificial Intelligence and Thirty-Fifth Conference
on Innovative Applications of Artificial Intelligence and Thirteenth Symposium on Educational Advances in Artificial Intelligence, AAAI’23/IAAI’23/EAAI’23. AAAI Press, 2023. [Zang et al., 2023] Hongyu Zang, Xin Li, Leiji Zhang, Yang Liu, Baigui Sun, Riashat Islam, Rémi Tachet des Combes, and Romain Laroche. Understanding and addressing the pitfalls of bisimulation-based representations in offline reinforcement learning. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS ’23, Red Hook, NY, USA, 2023. Curran Associates Inc. [Zhang et al., 2021] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.