Multi-step Proximal Policy Improvement in Offline Reinforcement Learning

Multi-step Proximal Policy Improvement in Offline Reinforcement Learning

Soohyun Choi*1

arXiv:2609.03842v1 [cs.LG] 3 Sep 2026

1

Seonvin Cho*1

Department of Electronic Engineering, Hanyang University, Seoul, Republic of Korea, {petersun0221,seonbin0319,snhong}@hanyang.ac.kr, * Equal contribution., † Corresponding author.

Abstract Offline reinforcement learning (RL) must reconcile two competing requirements: policy updates should stay near dataset-supported actions to keep value estimates reliable, yet meaningful gains often require moving beyond the behavior distribution. We develop a geometric view of offline actor updates by modeling policies as a probability manifold endowed with a chosen metric geometry. Under this lens, a broad class of offline actor objectives can be interpreted as a single proximal policy improvement step (SPI), i.e., an implicit discretization of a manifold gradient flow induced by a critic-defined energy. Building on this insight, we propose multi-step proximal policy improvement (MPI), a plug-in refinement mechanism that composes sequential re-centered proximal steps. MPI enables controlled policy improvement beyond dataset support while retaining proximal control at each refinement. The framework accommodates multiple policy geometries and admits practical instantiations for deterministic and diagonal-Gaussian policies. Experiments on D4RL benchmarks show that small numbers of MPI refinements improve strong offline baselines, including TD3+BC, ReBRAC, and IQL, on many tasks. Focused diagnostics further distinguish re-centered refinement from fixed-objective update scheduling and characterize limitations under critic error.

1

Songnam Hong†1

INTRODUCTION

Offline reinforcement learning (RL) aims to learn effective policies from a fixed dataset without further interaction with the environment. This setting is particularly appealing in domains where online exploration is costly, unsafe, or imPreprint. This manuscript has not been peer reviewed.

practical, such as healthcare, robotics, and autonomous driving. A fundamental challenge, however, is distribution shift: when a learned policy places probability mass on actions (or state–action pairs) weakly supported by the dataset, value estimation must extrapolate beyond observed data, often leading to overestimation and unstable policy improvement [Fujimoto et al., 2019b]. Consequently, a dominant theme in offline RL is to constrain policy updates to remain close to dataset support. A large body of practical algorithms addresses this issue by tempering policy improvement, spanning a spectrum from actor-side regularization to critic-side conservatism. On the actor side, advantage-weighted methods such as IQL [Kostrikov et al., 2021] and AWAC [Nair et al., 2020] moderate updates by reweighting dataset actions according to estimated value or advantage, thereby implicitly anchoring policy optimization to in-distribution behavior. Behavioranchored approaches, exemplified by TD3+BC [Fujimoto and Gu, 2021], impose an explicit behavior cloning term alongside value maximization to directly restrict policy deviation from the dataset. ReBRAC [Tarasov et al., 2023] extends TD3+BC by adding a critic penalty, coupling behavior anchoring with critic-side conservatism. On the critic side, conservative or uncertainty-aware methods such as CQL [Kumar et al., 2020] and EDAC [An et al., 2021] mitigate extrapolation by reshaping the value landscape itself to penalize or downweight unsupported actions. Despite their methodological differences, these methods share a common structural principle: policy improvement is moderated—either through actor-level anchoring (implicit or explicit) or through critic-level conservatism—to prevent overly aggressive extrapolation beyond data support. This stabilization mechanism, however, introduces a deeper challenge on the actor side. In offline RL, policy improvement must balance stability and exploration beyond the dataset: policy updates must remain sufficiently close to ensure reliable value estimation, yet must also move beyond the behavior distribution to achieve meaningful gains. Ex-

cessive anchoring can lead to under-improvement, whereas insufficient control results in instability under distribution shift. This tension raises a natural question: can we design a principled policy update mechanism that preserves stability while enabling controlled improvement beyond the dataset support? Our contributions. To address this tension, we develop a unifying geometric perspective on offline actor updates. We model the policy class as a statistical manifold M endowed with a chosen geometry, and show that a broad family of offline actor objectives can be interpreted as a single proximal policy improvement (SPI) step on M. In this view, SPI minimizes an energy functional (e.g., negative expected Q or advantage) regularized by a geometry-induced distance dM , corresponding to a single implicit gradient-flow step. This interpretation clarifies that many existing actor updates achieve stability via strong anchoring, yet can suffer from under-improvement. Building on this insight, we propose multi-step proximal policy improvement (MPI), which composes multiple sequential, re-centered proximal steps under the same geometry within each iteration. From the gradient-flow viewpoint, MPI amounts to a finer implicit discretization of the underlying manifold flow, enabling controlled advancement beyond dataset support while retaining local displacement control. MPI is implemented as a modular plug-in refinement stage that preserves the baseline actor-critic update, making it compatible with a wide range of existing offline RL algorithms. Empirically, small numbers of MPI refinements improve behavior-anchored baselines (e.g., TD3+BC [Fujimoto and Gu, 2021], ReBRAC [Tarasov et al., 2023], and IQL [Kostrikov et al., 2021]) on many standard D4RL tasks. We also characterize settings in which critic error limits these gains.

2

RELATED WORK

Statistical manifold in RL. A classical geometric viewpoint on policy optimization comes from information geometry: the Fisher information metric (Fisher–Rao) underlies the natural policy gradient [Kakade, 2001] and provides the local metric associated with KL-based trust-region formulations such as TRPO [Schulman et al., 2017]. Recent extensions further connect Fisher–Rao geometry to modern policy optimization and entropy-regularized MDPs [Lascu et al., 2025, Kerimkulov et al., 2025]. Another line of work leverages optimal transport (OT) geometry. Online policy updates have been formulated as gradient flows in Wasserstein space [Zhang et al., 2018], and Wasserstein Actor-Critic [Likmeta et al., 2023] leverages Wasserstein distances in the critic to propagate Q-posterior uncertainty and guide exploration through optimistic Q-value bounds. Trust-region methods have also been revisited under OT geometry, including OT-TRPO [Terpin et al., 2022] and metric-aware TRPO [Song et al., 2023], while Moskovitz et al. [2020] develop Wasserstein natural gradient methods for improved efficiency. In the offline setting, W-BRAC [Wu et al., 2019] employs the Wasserstein-1 distance as a behavior regularizer, and Asadulaev et al. [2024] reinterpret it as a partial OT objective that selectively matches high-value dataset regions. These works demonstrate that manifold-based metrics can enrich policy optimization by imposing intrinsic geometric structure. However, most prior developments focus on online RL, where geometry primarily shapes exploration dynamics or trust regions under direct environment interaction. In offline RL, geometric constructions are typically introduced as regularization terms that constrain policy deviation from the dataset, rather than as mechanisms defining the improvement dynamic themselves. In contrast, we interpret offline actor updates as implicit proximal steps induced by gradient flows on statistical manifolds, providing a unified view across metric choices.

Our main contributions are summarized as follows: • We provide a geometric interpretation of behavioranchored offline actor updates as an SPI on a statistical manifold, unifying diverse objectives under a common proximal framework. • We propose MPI, a modular plug-in refinement mechanism that advances the same manifold gradient flow through re-centered implicit steps, enabling controlled improvement beyond dataset support. • We demonstrate that MPI integrates seamlessly with existing offline RL algorithms and improves behavioranchored baselines on many D4RL tasks, while identifying limitations under critic error.

Offline RL algorithms. A central challenge in offline RL is distribution shift, which arises when learned policies deviate from dataset support. Existing methods address this issue either by anchoring policy updates to the dataset or by mitigating critic extrapolation error. Behavior-anchored approaches include TD3+BC [Fujimoto and Gu, 2021] and ReBRAC [Tarasov et al., 2023], which regularize actor updates toward dataset actions. Related behavior-constrained methods enforce dataset support in different ways: BCQ [Fujimoto et al., 2019a] restricts policy improvement to actions sampled from a learned behavior model, BEAR [Kumar et al., 2019] constrains the policy via distribution matching, and BRAC [Wu et al., 2019] penalizes divergence from the behavior policy. Advantageweighted approaches such as IQL [Kostrikov et al., 2021]

and AWAC [Nair et al., 2020] implicitly anchor updates through value-based weighting. Complementary approaches focus on critic conservatism. CQL [Kumar et al., 2020] penalizes high values on outof-distribution actions, and ensemble-based methods such as SAC-N and EDAC [An et al., 2021] reduce overestimation via multiple critics. More expressive policy parameterizations have also been explored, including flowbased methods (FQL [Park et al., 2025], FlowQ [Alles et al., 2025]), diffusion-based updates (Diffusion Q-learning [Wang et al., 2023], QSM [Psenka et al., 2025], BDPO [Gao et al., 2025]), and TD-based generative optimization (QAM [Li and Levine, 2026]). Sequence-model approaches further extend offline RL beyond standard actor–critic forms, including QT [Hu et al., 2024] and Q-Transformer [Chebotar et al., 2023]. In contrast to approaches that enforce strong actor anchoring or increase critic-side conservatism and complexity, we resolve the stability–improvement tension through a geometric multi-step refinement that advances a manifold gradient flow on the actor side, without modifying the critic.

policy π = π(· | s). When unambiguous, we omit the conditioning on s for simplicity. Let M denote a policy manifold equipped with a metric inducing a geodesic distance dM and the corresponding Riemannian gradient operator gradπ . Given an energy functional Es [π] defined on M, the gradient flow starting from π0 is the continuous-time dynamics: dπt = − gradπ Es [πt ]. dt

Closed-form solutions to (1) are generally unavailable, so practical algorithms rely on time discretization. An explicit Euler step on M takes the form πk+1 = Rπk (−η gradπ Es [πk ]) ,

BACKGROUND

3.1

OFFLINE REINFORCEMENT LEARNING

A Markov decision process (MDP) is defined as ⟨S, A, T, r, γ⟩, where S and A are Polish state and action spaces, T : S × A → P(S) is the transition kernel with P(S) the set of probability measures on S, r : S × A → R is the reward function, and γ ∈ [0, 1) is the discount factor. A (stochastic) policy is a Markov kernel π : S → P(A) assigning to each state s ∈ S a probability measure π(·|s) over actions. The RL objective is to find a policy π ⋆ thatPmaximizes the expected discounted ∞ return J(π) = Eπ [ t=0 γ t r(st , at )], where the expectation is taken over trajectories generated from an initial state distribution under policy π. In offline RL, the agent has access only to a fixed dataset D = {(si , ai , ri , s′i )}N i=1 collected by one or more behavior policies [Sutton and Barto, 2018]. All value estimates and policy updates must be derived exclusively from D. This induces distribution shift between the dataset distribution and the occupancy measure of the learned policy [Fujimoto et al., 2019b], making it challenging to maximize J(π) without extrapolation error. 3.2

GRADIENT FLOWS ON POLICY MANIFOLDS

We interpret policy optimization as a gradient flow, where a policy evolves as a probability distribution following the steepest descent of an energy functional under a chosen geometry. Fix a state s ∈ S and consider the state-conditioned

(2)

where η > 0 is a step size and Rπk is a retraction at πk [Hu et al., 2019]. This scheme typically requires sufficient smoothness of Es and may be sensitive to the choice of η. An alternative is the implicit Euler (minimizing movement) discretization, known as the Jordan–Kinderlehrer– Otto (JKO) scheme in Wasserstein spaces. Given a current policy πk , the next iterate πk+1 is defined as πk+1 ∈ arg min proxτM (Es | πk ) (π),

3

(1)

(3)

π∈M

where the proximal energy functional is defined as 1 2 d (π, πk ). (4) 2τ M Here, τ > 0 is the step size and dM is the geodesic distance induced by the chosen policy geometry. proxτM (Es | πk ) (π) := Es [π] +

In offline RL, E is instantiated using learned function approximators (e.g., a neural network critic), yielding an estimated functional Êϕ . We use the shorthand Es when the dependence on the critic parameters ϕ is clear, and all proximal/gradient-flow interpretations above should be read with respect to this estimated energy. This recursion produces a discrete sequence {πk }k≥0 . To relate it to the continuous-time flow, we define a piecewiseconstant interpolation π̂t = πk for t ∈ ((k − 1)τ, kτ ], with π̂0 = π0 . Under standard assumptions, a minimizer exists for each τ > 0, the (estimated) energy decreases monotonically along the iterates, and differentiability of Es is not required [Ambrosio et al., 2005]. Moreover, as τ → 0, the interpolated trajectory converges to a solution of the continuous gradient flow in (1). In the sequel, we leverage this discretized gradient-flow perspective to interpret offline policy updates as proximal steps on the manifold M, where the choice of distance dM determines the underlying policy geometry. 3.3

POLICY MANIFOLDS

We endow the policy manifold M with a geometry induced by a metric on probability measures. Different geometries

define distinct notions of proximity and steepest descent on distributions, thereby leading to different proximal update rules. In this work, we focus on two representative geometries: Fisher–Rao geometry from information geometry and the Wasserstein geometry from OT. Fisher–Rao geometry. Let Ω be a Polish space and let ∞ P+ (Ω) denote the set of smooth strictly positive probability densities on Ω. Equipped with the Fisher–Rao metric, this space forms an infinite-dimensional statistical manifold [Amari and Nagaoka, 2000, Pistone and Sempi, 1995, Ay et al., 2007]. The induced Fisher–Rao distance admits the closed-form expression Z dFR (ρ0 , ρ1 ) = 2 arccos

p

 ρ0 (x)ρ1 (x) dx ,

(5)

Ω

where BC(ρ0 , ρ1 ) :=

Z p

ρ0 (x)ρ1 (x) dx

Ω

4

MULTI-STEP PROXIMAL POLICY IMPROVEMENT

Building on the implicit Euler discretization in (3), we formulate offline policy improvement as proximal updates on the policy manifold. Under this view, many behavioranchored actor objectives can be interpreted as an SPI. We then extend this perspective to MPI, which composes several re-centered proximal steps within each training iteration. 4.1

OFFLINE RL AS SINGLE-STEP PROXIMAL IMPROVEMENT

We interpret a broad class of behavior-anchored actor updates as instances of SPI on the policy manifold. Under the discretized gradient-flow perspective of Section 3.2, SPI corresponds to an implicit Euler step on (M, dM ), updating the policy by minimizing an energy functional regularized by a geometry-induced proximity term. We formalize this update as the SPI operator in Algorithm 1.

denotes the Bhattacharyya coefficient. Moreover, dFR is locally equivalent to the Kullback–Leibler divergence: KL(ρ0 ∥ρ1 ) provides a second-order approximation to the squared Fisher–Rao distance. The detailed derivation is shown in Appendix A.2. Wasserstein geometry. Let P2 (Ω) denote the space of probability measures on Ω with finite second moments. The quadratic Wasserstein distance is defined as W22 (ρ0 , ρ1 ) =

Z

∥x − y∥22 dγ(x, y),

inf γ∈Γ(ρ0 ,ρ1 )

(6)

Ω×Ω

where Γ(ρ0 , ρ1 ) is the set of couplings (i.e., joint distributions with prescribed marginals) of ρ0 and ρ1 . Under standard regularity conditions, optimal couplings exist [Villani, 2008, Thm. 4.1]. Through Otto calculus [Otto, 2001, Chen and Li, 2020], the space P2 (Ω) admits a formal Riemannian structure. W2 is often approximated via entropic regularization, leading to a Sinkhorn distance

Algorithm 1 SPI 1: function SPI(Es , dM , πξ , τ ) 2: πξ+1 ∈ arg minπ∈M proxτM (Es | πξ ) (π) 3: return πξ+1 4: end function

We now instantiate SPI under representative choices of energy and geometry, showing that common offline actor updates arise as SPI in appropriate limits. Lemma 1 (Deterministic operator identification). Let πβ (· | s) denote the behavior policy that induces the offline dataset D. Let πTD3+BC denote the actor obtained by the anchored policy update of TD3+BC (and similarly ReBRAC), with anchor weight λ > 0. Assume deterministic policies πTD3+BC (a | s) = δ(a − µTD3+BC (s)) and πβ (a | s) = δ(a − µβ (s)) over a ∈ Rd . Under Wasserstein–2 geometry on MW and the energy functional Es [π] := Ea∼π(·|s) [−Q(s, a)], the anchored actor update matches a single SPI step: πTD3+BC ∈ SPI(Es , W2 , πβ ,

W2 2ε (ρ0 , ρ1 ) =

Z

∥x − y∥22 dγ(x, y)

min γ∈Γ(ρ0 ,ρ1 )

1 ), 2λ

up to multiplying the objective by a positive constant.

Ω×Ω

+ε KL(γ ∥ ρ0 ⊗ ρ1 ) , (7) which can be solved efficiently by Sinkhorn iterations for ε > 0 [Cuturi, 2013]. In what follows, we endow the policy family M with either Fisher–Rao or Wasserstein geometry, and use the associated geodesic distance as the proximal regularizer in the implicit update (3).

Proof. For deterministic policies π(a | s) = δ(a − µ(s)), the energy reduces to Es [π] = −Q(s, µ(s)) and the Wasserstein–2 proximity becomes W22 (π, πβ ) = ∥µ(s) − µβ (s)∥22

(Appendix A.3).

Hence the SPI subproblem is equivalent to
min −Q(s, µ(s)) + λ∥µ(s) − µβ (s)∥22 , µ

which matches the TD3+BC/ReBRAC anchored actor objective up to a positive rescaling of the objective. Lemma 2 (Local stochastic operator identification). Let ∞ π, πβ ∈ P+ (Ω) satisfy dFR (π, πβ ) < ϵ. Under Fisher– Rao geometry, consider the energy functional Es [π] := Ea∼π(·|s) [−A(s, a)]. Then, under the local approximation, the corresponding proximal step yields an advantageweighted policy
π(a | s) ∝ πβ (a | s) exp τ A(s, a) ,

re-centered proximal steps. In addition to distribution-shift considerations, MPI faces two intrinsic sources of error: (i) approximation error from optimizing a critic-induced estimated energy, and (ii) discretization error from using finite proximal step sizes. Figure 1 illustrates these effects and how the discretization error increases with larger τ .

up to higher-order terms in ϵ, recovering IQL/AWAC-style extraction (with temperature absorbed by rescaling A or τ ). Proof. For π in a small Fisher–Rao neighborhood of πβ , d2FR (π, πβ ) = 2 KL(π∥πβ ) + O(ϵ3 )

(Appendix A.2).

(a) Approximation (critic) error: gradient flow of the estimated energy differs from the true flow.

Therefore, up to higher-order terms, the Fisher–Rao proximal subproblem admits the KL-proxy form min π(·|s)

− Ea∼π(·|s) [A(s, a)] +

1 KL π(· | s) ∥ πβ (· | s) . τ

The optimizer of this KL-regularized expected-advantage objective is given by π(a | s) ∝ πβ (a | s) exp(τ A(s, a)), with temperature absorbed by rescaling A or τ , which recovers the advantage-weighted policy extraction used in IQL/AWAC-style methods. In summary, Lemmas 1 and 2 identify representative offline actor updates as SPI operators: TD3+BC/ReBRAC correspond to Es [π] = −Ea∼π(·|s) [Q(s, a)] with Wasserstein-2 proximity in the Dirac limit, while IQL/AWAC correspond locally to Es [π] = −Ea∼π(·|s) [A(s, a)] under Fisher–Rao geometry. 4.2

MULTI-STEP PROXIMAL IMPROVEMENT

(b) Approximation + discretization error in MPI: small τ (left) vs. large τ (right).

Figure 1: Visualization of MPI error sources. (a) Continuoustime gradient flow under the true energy (black solid) versus the critic-induced estimated energy (green dashed), illustrating approximation error. (b) Discretized proximal trajectories for the estimated energy (red dash-dotted) with different step sizes τ . Definition 3 (MPI). Given a base policy πbase ∈ M, an energy functional Es [·], a policy geometry dM , and step sizes τ = (τ1 , . . . , τK ) with τi > 0, multi-step proximal policy improvement (MPI) generates a sequence {πi }K i=0 with π0 = πbase and, for i = 1, . . . , K, i πi ∈ arg min proxτM (Es | πi−1 ) (π).

(8)

π∈M

The SPI interpretation reveals that many offline actor updates correspond to advancing a gradient flow by a single implicit step on (M, dM ). In offline RL, this step is typically chosen conservatively to ensure stability under distribution shift, which may lead to under-improvement under a fixed critic. Simply enlarging the proximal region is not a principled solution, since it can expose the actor to outof-distribution actions where value (energy) estimates are unreliable. Motivated by the implicit Euler viewpoint, we instead advance the same manifold dynamics over a longer effective horizon within each training iteration by composing multiple

Equivalently, MPI repeatedly applies SPI while updating the proximal center to the latest iterate: πi = SPI(Es [·], dM , πi−1 , τi ) ,

i = 1, . . . , K. (9)

In our offline RL setting, πbase denotes the actor produced by the base algorithm’s anchored update (i.e., a single SPItype update from the behavior policy). Algorithm 2 summarizes MPI. Connection to implicit Euler discretization. Each update in (8) is an implicit Euler step of the gradient flow

Environment πbase

IQL π1

π2

HalfCheetah-M HalfCheetah-MR HalfCheetah-ME

48.42 ± 0.19 43.84 ± 0.88 90.44 ± 4.06

53.78 ± 0.24 49.80 ± 0.50 91.02 ± 3.72

55.94 ± 0.67 52.30 ± 0.87 88.58 ± 5.86

Hopper-M Hopper-MR Hopper-ME

60.44 ± 6.17 90.54 ± 9.40 101.40 ± 3.51 98.35 ± 5.09 101.05 ± 1.49 101.22 ± 2.29 59.18 ± 3.56 91.44 ± 7.43 100.60 ± 1.54 86.72 ± 7.58 100.50 ± 3.36 103.04 ± 0.42 90.28 ± 20.12 100.70 ± 0.29 99.53 ± 1.54 71.02 ± 23.88 99.32 ± 1.49 99.12 ± 1.64 102.20 ± 21.41 106.14 ± 10.82 107.38 ± 10.84 111.65 ± 0.45 109.65 ± 2.73 107.22 ± 7.96 94.62 ± 9.48 96.82 ± 12.32 105.48 ± 8.26

Walker2d-M Walker2d-MR Walker2d-ME

80.54 ± 3.57 78.06 ± 8.59 111.86 ± 1.06

74.94 ± 15.15 81.82 ± 3.39 112.62 ± 0.77

Locomotion avg

78.06

84.57

πbase

ReBRAC π1

64.90 ± 0.82 65.92 ± 1.51 50.62 ± 2.26 50.73 ± 0.95 103.00 ± 4.76 97.12 ± 4.34

π2

πbase

TD3+BC π1

π2

66.00 ± 1.58 49.98 ± 0.49 97.08 ± 2.04

48.40 ± 0.38 44.70 ± 0.68 90.62 ± 6.51

54.46 ± 0.36 49.76 ± 0.83 89.88 ± 3.72

57.66 ± 0.51 51.98 ± 1.06 74.42 ± 8.52

85.26 ± 3.34 85.75 ± 0.64 85.55 ± 1.05 86.97 ± 0.97 83.78 ± 2.77 84.30 ± 6.49 89.44 ± 2.10 84.46 ± 5.58 79.55 ± 5.40 86.33 ± 5.04 88.35 ± 3.73 79.30 ± 6.07 85.66 ± 2.55 89.42 ± 1.89 113.78 ± 1.22 111.83 ± 0.28 112.08 ± 0.79 113.38 ± 0.24 110.24 ± 0.18 112.20 ± 0.64 113.62 ± 1.50 88.02

88.44

89.90

89.97

75.76

84.87

86.86

Table 1: πbase refers to the base policy trained with the original algorithm, and π1 /π2 are obtained by policy extension with step size τ . D4RL score returns are mean±std over 5 seeds. Algorithm 2 MPI 1: function MPI(Es [·], dM , πbase , τ , K) 2: π0 ← πbase 3: for k = 1, . . . , K do 4: πk ← SPI(Es [·], dM , πk−1 , τk ) 5: end for 6: return {πk }K k=1 7: end function

associated with Es [·] under the geometry induced by dM . Thus, MPI corresponds to a K-step implicit Euler discretization with step sizes {τi }K i=1 , and the total advancement in PK “flow time” is i=1 τi . For sufficiently small step sizes, the piecewise-constant interpolation of {πi } approximates the continuous-time gradient flow in Section 3.2, while retaining the stability properties of proximal updates. MPI as longer flow-time. Within each training iteration, MPI performs multiple implicit proximal steps with respect to the same estimated energy functional Eˆsϕ . Geometrically, MPI composes local trust-region moves: each step remains close to the previous iterate, while re-centering enables a larger net displacement along the discretized trajectory. This also clarifies the role of critic reliability: increasing K can reduce under-improvement when Eˆsϕ is locally accurate, but may over-optimize critic artifacts when the estimate is unreliable. Re-centering versus a single enlarged step. Because MPI updates the proximal center after every step, composing re-centered maps is generally different from weakening the regularizer in one fixed-anchor problem. For example, for the Euclidean energy E(µ) = 21 (µ21 +2µ22 ) and µ0 = (1, 1), two unit proximal steps reach (1/4, 1/9), whereas no single proximal step reaches the same point. Thus, MPI does not enlarge the policy class; it changes the endpoints reachable by sequential re-centered maps. Appendix B.1 formalizes this observation.

Descent of the estimated energy. MPI inherits the standard descent property of exact proximal minimization. Since πi−1 is feasible in (8), the optimality of πi yields Es [πi ] +

1 2 d (πi , πi−1 ) ≤ Es [πi−1 ]. 2τi M

(10)

Hence, the estimated energy decreases along exact MPI iterates. When Es [·] is instantiated as the negative of a criticdriven objective, this corresponds to successive actor-side improvement under a fixed critic, while each step remains controlled by a geometry-induced proximity term. To state the effect of critic and optimization errors explicitly, let F̂s = −Ês denote the estimated critic objective, let Fs be a reference local objective, and suppose that the ith proximal subproblem is solved up to error ξi ≥ 0. For the visited neighborhood BK , define ϵF (BK ) = supπ∈BK |F̂s (π) − Fs (π)|. Then Fs (πK ) − Fs (π0 ) ≥

K X 1

d2M (πi , πi−1 ) 2τ i i=1

− 2ϵF (BK ) −

K X

(11)

ξi .

i=1

This gives a conditional local explanation: refinement is beneficial when the accumulated proximal gain exceeds local critic uncertainty and optimization error. Appendix B.2 provides the corresponding estimated-objective bound and derivation. 4.3

PRACTICAL INSTANTIATION

MPI is implemented as a plug-in refinement stage that augments an existing offline RL algorithm without modifying its critic update. At each training iteration, we first run the base algorithm’s standard critic and actor updates to obtain a base policy, denoted by π0 . Using the current estimated energy Ê[·] induced by the critic and the chosen geometry dM , we then apply MPI by repeatedly solving re-centered proximal subproblems, producing refined policies {πi }K i=1 from

π0 . Algorithm 3 summarizes this per-iteration refinement procedure. Algorithm 3 MPI-based Policy Extension Require: Base offline RL algorithm Algo, refinement steps K, step sizes τ , geometry dM 1: for each gradient step do 2: Algo.UpdateCritic() 3: Algo.UpdateActor() 4: π0 ← Algo.Actor   ˆ 5: {πi }K i=1 ← MPI Es [·], dM , π0 , τ , K 6: end for 7: Outputs {πi }K i=0

medium-expert tasks, where the base policy is already close to optimal, even small approximation errors in the critic can cause the refinement step to increase the true energy, leading to non-monotone behavior. This behavior is less pronounced for ReBRAC. Because ReBRAC explicitly regularizes the critic with respect to the behavior distribution, the resulting value function is intended to be most reliable near dataset-supported actions and less prone to optimistic extrapolation. Accordingly, πbase is typically already optimized within the region where the critic is well-calibrated, so additional refinement steps offer limited room for improvement. 5.2

When the policy class is diagonal Gaussian or deterministic, we compute the geometry-induced distance dM in closed form following Appendix A.3. Otherwise, when M is endowed with Wasserstein geometry, we instantiate dM using Sinkhorn distance (7). For advantage-weighted methods such as IQL, our SPI interpretation is most naturally expressed under Fisher–Rao geometry via its local connection to KL divergence (Appendix A.2). Nevertheless, for practical refinement we use Wasserstein-2 geometry, motivated by the analysis in Appendix A.1 which explains why Fisher–Rao refinement can be unfavorable in near-deterministic continuous-control regimes; Fisher–Rao refinement results for IQL are reported in Appendix D.

5

EXPERIMENTS

We evaluate MPI-based policy extension on standard offline RL benchmarks from D4RL [Fu et al., 2020], including locomotion tasks and AntMaze tasks. We consider three representative offline RL methods: IQL, TD3+BC, and ReBRAC. We report D4RL-normalized scores of the base policies πbase and their MPI extensions π1 and π2 averaged over five random seeds. All base algorithms and training hyperparameters follow the standard CORL offline RL benchmark configurations [Tarasov et al., 2022]. 5.1

MPI-BASED POLICY EXTENSION

MPI introduces a single additional hyperparameter τ . For each environment, we sweep τ over a fixed grid [0.001, 0.005, 0.01, 0.05] and report the best-performing setting selected by the score of π2 , as summarized in Table 1. Comparing the score of π1 to that of π2 reveals that a second re-centered refinement step often provides additional gains beyond a single step (e.g., HalfCheetah-medium, Hopper-medium, Walker2d-medium-replay), but improvements are not strictly monotone across all tasks. In

RE-CENTERING DIAGNOSTICS

MPI changes the proximal subproblem rather than only increasing the number of optimizer steps on a fixed actor objective. As focused diagnostics on Hopper-medium-replay-v2, re-centered refinement improves the policy under the original TD3+BC step size, whereas compute-matched optimization of the fixed-anchor objective does not reproduce the MPI result. Since these diagnostics cover one environment, we use them only to distinguish re-centered MPI from fixed-objective update scheduling. Full protocols and results are provided in Appendix B.3. 5.3

COMPARISON WITH AN EXPLICIT SCHEME

To isolate the effect of implicit proximal refinement, we compare against an explicit action-gradient update that performs a first-order Euler step on the critic in (2). When dM = W2 and Es [π] = −Ea∼π(·|s) [Q(s, a)], the explicit update is given as πk+1 (s) = πk (s) + τ ∇a Q(s, a)|a=πk (s) .

(12)

This update can be interpreted as a direct gradient ascent step in action space, and is a natural baseline when one views refinement as moving actions along ∇a Q. These experiments are conducted to study the numerical stability of the implicit update. As mentioned in Section 3.2, the explicit update is more sensitive to hyperparameter tuning and local curvature of the energy function. In particular, when the critic is imperfect, the estimated gradient field may deviate from the true flow, and large step sizes can amplify such approximation errors and lead to overshooting or unstable trajectories. Figure 2 visualizes these effects: as τ increases, the explicit trajectory can deviate substantially, whereas the implicit proximal refinement remains comparatively stable. For a fair comparison, within each environment we use the same step size τ for both the implicit (MPI) and explicit updates. Table 2 reports the percentage changes relative to πbase ,

TD3+BC

Environment

ReBRAC

Implicit

Explicit

Implicit

Explicit

π1

π1

π1

π1

π2

π2

π2

π2

HalfCheetah-M +12.6 +19.1 +0.9 +2.6 +1.6 +1.7 +0.4 +0.8 HalfCheetah-MR +11.4 +16.1 +3.0 +3.0 +1.8 +0.6 −21.8 −25.9 HalfCheetah-ME −2.9 −17.3 +0.9 −2.2 −5.7 −5.7 +1.9 +2.8 Hopper-M Hopper-MR Hopper-ME

+53.2 +72.8 +26.5 +50.6 −11.1 −2.1 +0.6 −5.0 +23.3 +36.1 +15.3 +22.2 +1.6 +1.9 +8.2 +14.9 +6.6 +6.8 +19.8 +23.2 −2.2 −5.3 −4.7 −8.5

Walker2d-M Walker2d-MR Walker2d-ME

+3.6 +9.7 +1.7

+6.7 +5.8 +7.2 +9.2 +11.4 +4.4 +3.1 +1.4 +2.4

0.0 +1.3 +0.3 −2.1 +5.7 +0.2 +9.8 +11.3 +0.2 +1.2 +0.6 +0.9

Environment

Implicit

Explicit

π1

π1

π2

π2

Umaze 0.0 +2.0 −0.4 −0.8 Umaze-Diverse +15.0 +15.0 +0.9 +1.2 +10.0 +15.0 −3.2 +0.5 Medium-Play Medium-Diverse 0.0 0.0 +0.7 −4.3 Large-Play +12.5 +50.0 −6.8 −1.6 Large-Diverse +14.3 +5.7 −2.4 −1.2 (b) ReBRAC AntMaze

(a) Locomotion (TD3+BC, ReBRAC)

Table 2: Comparison of implicit policy extension (MPI) and explicit policy extension. (a) TD3+BC & ReBRAC for Locomotion, (b) ReBRAC for AntMaze. Values indicate the percentage change relative to πbase . pirically less sensitive to critic-induced errors than a naive explicit gradient step.

6

(a) Implicit (left) and explicit (right) gradient-flow trajectories with small τ .

(b) Implicit (left) and explicit (right) gradient-flow trajectories with large τ .

Figure 2: Effect of the step size τ on the explicit scheme. Figure 2a shows that with small τ , implicit (red dash-dotted) and explicit (orange dash-dotted) updates yield similarly stable trajectories. In contrast, Figure 2b shows that for large τ , the implicit scheme remains stable while the explicit scheme becomes unstable and diverges.

showing that implicit MPI generally yields larger and more reliable improvements than the explicit scheme. These results support the central motivation for MPI in offline RL: within the evaluated settings, an implicit proximal step behaves like a trust-region improvement operator and is em-

CONCLUSION & FUTURE WORKS

We presented a geometric perspective on offline actor updates, showing that a broad class of behavior-anchored objectives can be interpreted as SPI on a statistical manifold. This viewpoint clarifies how strong anchoring stabilizes training but may induce under-improvement. Motivated by this insight, we proposed MPI, which composes sequential, re-centered proximal updates under a fixed geometry. From the gradient-flow perspective, MPI corresponds to a finer implicit discretization of the same manifold flow, enabling controlled policy advancement while retaining local displacement control. MPI is implemented as a modular plugin refinement stage that leaves the base actor–critic training unchanged. Empirically, small numbers of refinements improve behavior-anchored baselines on many standard D4RL tasks without introducing significant additional complexity. Our critic-error analysis also shows that refinement can saturate or degrade when the learned objective is insufficiently aligned with true return. Overall, our results suggest that carefully controlling the discretization of geometry-induced policy flows provides a useful path toward reducing underimprovement in offline reinforcement learning, provided that refinement remains local to a reliable critic region. A promising direction is to extend MPI to generative policy classes (e.g., flow- or diffusion-based actors), where policy improvement is still often implemented as a conservative single-step update despite the expressive sampler (e.g. Diffusion Q-learning [Wang et al., 2023]). Incorporating MPI could enable multi-step, re-centered refinement directly in distribution space, but will require scalable prox-

ies for geometry-induced distances and careful control of critic-induced artifacts under highly expressive policies.

via autoregressive q-functions, 2023. URL https:// arxiv.org/abs/2309.10150.

Another direction is to incorporate predictor–corrector refinement, where an explicit predictor proposes a tentative improvement step and an implicit proximal corrector re-centers and stabilizes it [Diao et al., 2023]. Such schemes could reduce the cost of repeatedly solving proximal subproblems while retaining MPI’s robustness to critic approximation error. The main challenge is to design a critic-induced predictor that remains reliable off-support and a corrector that meaningfully contracts harmful extrapolation.

Yifan Chen and Wuchen Li. Optimal transport natural gradient for statistical manifolds with continuous sample space, 2020. URL https://arxiv.org/abs/ 1805.08380.

Bibliography Marvin Alles, Nutan Chen, Patrick van der Smagt, and Botond Cseke. Flowq: Energy-guided flow policies for offline reinforcement learning, 2025. URL https: //arxiv.org/abs/2505.14139. S. Amari and H. Nagaoka. Methods of Information Geometry. Translations of mathematical monographs. American Mathematical Society, 2000. ISBN 9780821843024. URL https://books.google. co.kr/books?id=vc2FWSo7wLUC. L. Ambrosio, N. Gigli, and G. Savare. Gradient Flows: In Metric Spaces and in the Space of Probability Measures. Lectures in Mathematics. ETH Zürich. Birkhäuser Basel, 2005. ISBN 9783764324285. URL https://books. google.co.kr/books?id=HZqhWIq1-jgC. Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble, 2021. URL https: //arxiv.org/abs/2110.01548. Arip Asadulaev, Rostislav Korst, Alexander Korotin, Vage Egiazarian, Andrey Filchenkov, and Evgeny Burnaev. Rethinking optimal transport in offline reinforcement learning, 2024. URL https://arxiv.org/abs/2410. 14069. Nihat Ay, Jürgen Jost, Hông Vân Lê, and Lorenz Schwachhöfer. Information Geometry. Springer Cham, 2007. ISBN 978-3-319-564784. URL https://link.springer.com/ book/10.1007/978-3-319-56478-4# bibliographic-information. Yevgen Chebotar, Quan Vuong, Alex Irpan, Karol Hausman, Fei Xia, Yao Lu, Aviral Kumar, Tianhe Yu, Alexander Herzog, Karl Pertsch, Keerthana Gopalakrishnan, Julian Ibarz, Ofir Nachum, Sumedh Sontakke, Grecia Salazar, Huong T Tran, Jodilyn Peralta, Clayton Tan, Deeksha Manjunath, Jaspiar Singht, Brianna Zitkovich, Tomas Jackson, Kanishka Rao, Chelsea Finn, and Sergey Levine. Q-transformer: Scalable offline reinforcement learning

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transportation distances, 2013. URL https: //arxiv.org/abs/1306.0895. Michael Diao, Krishnakumar Balasubramanian, Sinho Chewi, and Adil Salim. Forward-backward gaussian variational inference via jko in the bures-wasserstein space, 2023. URL https://arxiv.org/abs/ 2304.05398. Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020. Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132–20145, 2021. Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration, 2019a. URL https://arxiv.org/abs/1812.02900. Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration, 2019b. URL https://arxiv.org/abs/1812.02900. Chen-Xiao Gao, Chenyang Wu, Mingjun Cao, Chenjun Xiao, Yang Yu, and Zongzhang Zhang. Behaviorregularized diffusion policy optimization for offline reinforcement learning, 2025. URL https://arxiv. org/abs/2502.04778. Jiang Hu, Xin Liu, Zaiwen Wen, and Yaxiang Yuan. A brief introduction to manifold optimization, 2019. URL https://arxiv.org/abs/1906.05450. Shengchao Hu, Ziqing Fan, Chaoqin Huang, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Q-value regularized transformer for offline reinforcement learning, 2024. URL https://arxiv.org/abs/2405. 17098. Sham M. Kakade. A natural policy gradient. In NIPS, 2001. Bekzhan Kerimkulov, James-Michael Leahy, David Siska, Lukasz Szpruch, and Yufei Zhang. A fisher–rao gradient flow for entropy-regularised markov decision processes in polish spaces: B. kerimkulov et al. Foundations of Computational Mathematics, pages 1–75, 2025.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169, 2021. Aviral Kumar, Justin Fu, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction, 2019. URL https://arxiv.org/abs/ 1906.00949. Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. CoRR, abs/2006.04779, 2020. URL https: //arxiv.org/abs/2006.04779. Razvan-Andrei Lascu, David Šiška, and Łukasz Szpruch. Ppo in the fisher-rao geometry, 2025. URL https:// arxiv.org/abs/2506.03757. Qiyang Li and Sergey Levine. Q-learning with adjoint matching, 2026. URL https://arxiv.org/abs/ 2601.14234. Amarildo Likmeta, Matteo Sacco, Alberto Maria Metelli, and Marcello Restelli. Wasserstein actor-critic: Directed exploration via optimism for continuous-actions control, 2023. URL https://arxiv.org/abs/2303. 02378. Ted Moskovitz, Michael Arbel, Ferenc Huszar, and Arthur Gretton. Efficient wasserstein natural gradients for reinforcement learning. arXiv preprint arXiv:2010.05380, 2020. Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. Accelerating online reinforcement learning with offline datasets. CoRR, abs/2006.09359, 2020. URL https://arxiv.org/abs/2006.09359. Felix Otto. The geometry of dissipative evolution equations: The porous medium equation. Communications in Partial Differential Equations, 26(1-2):101–174, 2001. doi: 10. 1081/PDE-100002243. URL https://doi.org/10. 1081/PDE-100002243. Seohong Park, Qiyang Li, and Sergey Levine. Flow qlearning, 2025. URL https://arxiv.org/abs/ 2502.02538. Giovanni Pistone and Carlo Sempi. An infinite-dimensional geometric structure on the space of all the probability measures equivalent to a given one. The annals of statistics, pages 1543–1561, 1995. Michael Psenka, Alejandro Escontrela, Pieter Abbeel, and Yi Ma. Learning a diffusion model policy from rewards via q-score matching, 2025. URL https://arxiv. org/abs/2312.11752.

John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Pieter Abbeel. Trust region policy optimization, 2017. URL https://arxiv.org/abs/1502. 05477. Jun Song, Niao He, Lijun Ding, and Chaoyue Zhao. Provably convergent policy optimization via metric-aware trust region methods, 2023. URL https://arxiv. org/abs/2306.14133. Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018. URL http://incompleteideas. net/book/the-book-2nd.html. Denis Tarasov, Alexander Nikulin, Dmitry Akimov, Vladislav Kurenkov, and Sergey Kolesnikov. CORL: Research-oriented deep offline reinforcement learning library. In 3rd Offline RL Workshop: Offline RL as a ”Launchpad”, 2022. URL https://openreview. net/forum?id=SyAS49bBcv. Denis Tarasov, Vladislav Kurenkov, Alexander Nikulin, and Sergey Kolesnikov. Revisiting the minimalist approach to offline reinforcement learning, 2023. URL https: //arxiv.org/abs/2305.09836. Antonio Terpin, Nicolas Lanzetti, Batuhan Yardim, Florian Dörfler, and Giorgia Ramponi. Trust region policy optimization with optimal transport discrepancies: Duality and algorithm for continuous actions, 2022. URL https://arxiv.org/abs/2210.11137. C. Villani. Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008. ISBN 9783540710493. URL https://books.google.co.kr/books?id= NZXiNAEACAAJ. Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning, 2023. URL https://arxiv. org/abs/2208.06193. Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning, 2019. URL https://arxiv.org/abs/1911.11361. Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization as wasserstein gradient flows. CoRR, abs/1808.03030, 2018. URL http://arxiv. org/abs/1808.03030.

Multi-step Proximal Policy Improvement in Offline Reinforcement Learning (Supplementary Material)

Soohyun Choi*1 1

Seonvin Cho*1

Songnam Hong†1

Department of Electronic Engineering, Hanyang University, Seoul, Republic of Korea, {petersun0221,seonbin0319,snhong}@hanyang.ac.kr, *

Equal contribution., † Corresponding author.

A

PROOFS

A.1

FIRST-ORDER ANALYSIS OF GAUSSIAN SPI

We derive the first-order approximation of one-step SPI updates for diagonal Gaussian policies under the Wasserstein and Fisher–Rao (FR) geometries. Fix a state s and expand Q(s, a) at a = µbase (s) up to second order. For a local one-step analysis, we keep only terms that contribute at O(τ ) when initialized at πbase (the quadratic term in µ − µbase affects the update only at O(τ 2 )). The resulting local approximation of the energy functional becomes 1 Es [π] = −GT (µ − µbase ) − tr(HΣ) + const., 2 where G = ∇a Q(s, ·)|a=µbase and H = ∇2a Q(s, ·)|a=µbase . Since Σ is diagonal, only the diagonal of H contributes through tr(HΣ); define e := diag(diag(H)). H Local proximal objectives. Let πbase = N (µbase , Σbase ) and π = N (µ, Σ) with Σ = diag(σ 2 ). (We use the closed-form W2 distance for diagonal Gaussians; see Appendix A.3.) The local proximal objectives can be written as Wasserstein:

1 proxτW2 (Es | πbase ) (π) = −GT (µ − µbase ) −

Fisher–Rao:

proxτdFR (Es | πbase ) (π) = −GT (µ − µbase ) −

d X

2 i=1

e ii σ 2 + 1 W 2 (π, πbase ) + const. H i 2τ 2

d

1Xe 2 1 2 Hii σi + d (π, πbase ) + const. 2 i=1 2τ FR

First-order one-step updates. Differentiating the proximal objectives and evaluating at π = πbase yields the following first-order expansions for the one-step SPI update π1 = N (µ1 , Σ1 ) as τ → 0: Wasserstein:

µ1 = µbase + τ G + O(τ 2 ), e Σbase + O(τ 2 ). Σ1 = Σbase + 2τ H

Fisher–Rao:

µ1 = µbase + τ Σbase G + O(τ 2 ), 2 e Σbase Σ1 = Σbase + τ H + O(τ 2 ).

These relations describe the first-order evolution of the mean and covariance (and thus the standard deviation) under each geometry for small step size τ .

Implication for Fisher–Rao refinement in near-deterministic regimes. The Fisher–Rao updates exhibit a structural limitation when the base actor is near-deterministic. In particular, the mean update is preconditioned by the current covariance, µ1 − µbase = O(τ Σbase G), so as Σbase shrinks, the effective step size for the mean collapses and multi-step refinement becomes increasingly conservative. The covariance update further clarifies how variance can contract. From the Fisher–Rao expansion, (Σ1 )ii = (Σbase )ii + τ Hii (Σbase )2ii + O(τ 2 ). Hence, for coordinates where Q is locally concave around the base action (i.e., Hii < 0), we have (Σ1 )ii < (Σbase )ii

for sufficiently small τ,

so refinement can further concentrate the policy once it enters a small-variance regime (while the opposite trend occurs when Hii > 0). In implementations that realize Fisher–Rao refinement via local KL-based proxies (Appendix A.2), the proximity term introduces inverse-variance weighting and variance-ratio/log-determinant factors; as σbase,i → 0, the resulting objective can become ill-conditioned, leading to numerical instability and occasional optimization divergence. We empirically observe this variance-collapse/instability pattern in our Fisher–Rao refinement ablations on continuous-control benchmarks, motivating our use of Wasserstein-2 geometry for practical MPI refinement in near-deterministic regimes, where W2 remains well-behaved in the vanishing-variance (Dirac) limit. A.2

KL DIVERGENCE APPROXIMATION OF THE FISHER–RAO DISTANCE

∞ Let ξ ∈ Tπ⋆ P+ (Ω) and define

π = Expπ⋆ (ξ),

δ := ∥ξ∥FR,π⋆ .

(Equivalently, δ = dFR (π, π ⋆ ).) Set F (π) := 2 KL(π∥π ⋆ ). Then F (π ⋆ ) = 0,

∇FR F (π ⋆ ) = 0,

HessFR F (π ⋆ ) = 2 gFR,π⋆ .

By Taylor’s theorem in Fisher–Rao normal coordinates (for π in a sufficiently small Fisher–Rao neighborhood of π ⋆ ), F (π) = F (π ⋆ ) + 21 HessFR F (π ⋆ )[ξ, ξ] + R3 = δ 2 + R3 ,

|R3 | ≤ C δ 3 .

Hence,
2 KL(π∥π ⋆ ) = d2FR (π, π ⋆ ) + O d3FR (π, π ⋆ ) . Equivalently, writing ε := d2FR (π, π ⋆ ) gives 2 KL(π∥π ⋆ ) = ε + O(ε3/2 ). A.3

CLOSED-FORM DISTANCES FOR DIAGONAL GAUSSIAN POLICIES

We collect closed-form expressions for the Wasserstein-2 and Fisher–Rao distances used in our proximal objectives, specialized to diagonal Gaussian policies on Rd . Throughout, we write ρi = N (µi , Σi ) with Σi = diag(σi2 ) and σi ∈ Rd+ . Wasserstein-2 distance. Let ρ0 = N (µ0 , Σ0 ) and ρ1 = N (µ1 , Σ1 ) be Gaussians with positive definite covariances. A standard formula gives   1/2 1/2
1/2 W22 (ρ0 , ρ1 ) = ∥µ0 − µ1 ∥22 + tr Σ0 + Σ1 − 2 Σ1 Σ0 Σ1 . (13) 1/2

1/2

(Equivalently, one may use (Σ0 Σ1 Σ0 )1/2 inside the trace; both yield the same value.) For diagonal covariances Σ0 = diag(σ02 ) and Σ1 = diag(σ12 ), this simplifies to
W22 N (µ0 , diag(σ02 )), N (µ1 , diag(σ12 )) = ∥µ0 − µ1 ∥22 + ∥σ0 − σ1 ∥22 .

(14)

In particular, deterministic actors correspond to the Dirac limit σ0 , σ1 → 0, yielding W22 (δ(· − µ0 ), δ(· − µ1 )) = ∥µ0 − µ1 ∥22 .

(15)

Fisher–Rao distance.

Let Σ̄ := 21 (Σ0 + Σ1 ). The Bhattacharyya coefficient between Gaussians admits BC(ρ0 , ρ1 ) =

|Σ0 |1/4 |Σ1 |1/4 exp − 18 (µ0 − µ1 )⊤ Σ̄−1 (µ0 − µ1 ) , 1/2 |Σ̄|

(16)

dFR (ρ0 , ρ1 ) = 2 arccos BC(ρ0 , ρ1 ) .

(17)

and the Fisher–Rao distance is For diagonal covariances, Σ̄ = diag((σ02 + σ12 )/2) and the coefficient factorizes as BC(ρ0 , ρ1 ) =

d Y j=1

s

(µ0,j − µ1,j )2 2σ0,j σ1,j exp − 2 + σ2 2 + σ2 ) σ0,j 4(σ0,j 1,j 1,j

! ,

(18)

which can be substituted into (17) to obtain a closed-form dFR for diagonal Gaussian policies. Remark (domain). Unlike W2 , which is defined on P2 (Rd ) and naturally accommodates Dirac measures, dFR is typically ∞ defined on strictly positive densities (e.g., P+ (Ω)). As a result, deterministic (Dirac) policies lie outside the Fisher–Rao manifold, even though one may consider limiting behavior as σ → 0.

B

RE-CENTERING AND CRITIC-ERROR ANALYSES

B.1

REACHABLE SETS UNDER RE-CENTERED PROXIMAL MAPS

Let Tτ denote the exact proximal map associated with a fixed energy and geometry. For a base policy π0 , define the K-step reachable set as RK (π0 ) = {TτK ◦ · · · ◦ Tτ1 (π0 ) : τi > 0} . (19) Proposition 4 (Reachable-set nesting). Assume Tϵ (π) → π as ϵ ↓ 0 for every policy in the reachable region. Then R0 (π0 ) ⊆ R1 (π0 ) ⊆ R2 (π0 ) ⊆ · · · . Moreover, the inclusion can be strict even for a quadratic energy in Euclidean space. Proof. Take any πK ∈ RK (π0 ). Appending the proximal map Tϵ gives a point in RK+1 (π0 ), and Tϵ (πK ) → πK as ϵ ↓ 0. Hence every point in RK (π0 ) belongs to RK+1 (π0 ), which proves the nesting after taking closures. For strictness, consider E(µ) =

1 2 (µ + 2µ22 ), 2 1

µ0 = (1, 1).





The Euclidean proximal map satisfies Tτ (µ0 ) =

1 1 , 1 + τ 1 + 2τ

.

Two unit steps yield T1 ◦ T1 (µ0 ) = (1/4, 1/9). A single step would require τ = 3 from the first coordinate and τ = 4 from the second coordinate, which is impossible. The point is also not obtained as a boundary limit of the one-step curve. Therefore, sequential re-centering can reach an endpoint that no single enlarged step can reach. This argument concerns the set of endpoints induced by the proximal maps. It does not enlarge the policy parameterization or establish that every additional endpoint has higher true return. B.2

LOCAL CRITIC-OBJECTIVE SENSITIVITY

Suppose the approximately solved ith subproblem satisfies Ês (πi ) +

1 2 d (πi , πi−1 ) ≤ Ês (πi−1 ) + ξi . 2τi M

With F̂s = −Ês , rearranging and summing over i = 1, . . . , K gives F̂s (πK ) − F̂s (π0 ) ≥

K X 1 i=1

2τi

d2M (πi , πi−1 ) −

K X

ξi .

(20)

i=1

If |F̂s (π) − Fs (π)| ≤ ϵF (BK ) throughout the visited neighborhood, the two endpoint discrepancies contribute at most 2ϵF (BK ), yielding (11). The bound separates three quantities: the accumulated proximal displacement, local critic-objective error, and approximatesolve error. It is deliberately local and conditional; converting it to a true-return guarantee would require additional assumptions linking Fs to the environment return and controlling distribution shift along the full trajectory. B.3

PROTOCOLS FOR THE ADDITIONAL DIAGNOSTICS

Same-τ diagnostic. For deterministic TD3+BC under Wasserstein geometry, the actor regularization coefficient λ corresponds to τ = 1/(2λ). We keep the original TD3+BC actor structure and the same λ on Hopper-medium-replay-v2, so no MPI-specific step-size sweep is used. This diagnostic isolates re-centering under a fixed nominal step size; it is not an exhaustive per-task comparison against an oracle one-step coefficient. Table 3 reports the final D4RL means.

Same τ

πbase

π1

π2

π3

π4

74.725

94.450

100.775

101.225

100.750

Table 3: Same-τ diagnostic on Hopper-medium-replay-v2. Compute-matched actor updates. On Hopper-medium-replay-v2, we compare the base TD3+BC update, five actor updates per delayed actor-update slot under the original fixed-anchor objective, and TD3+BC-MPI with K = 4 (one base actor update followed by four re-centered refinements). The latter two use the same number of actor optimization stages, the same TD3+BC loss template, and the same hyperparameters. Table 4 reports the final mean D4RL scores over four seeds. The purpose of this experiment is to distinguish re-centered MPI from additional optimization of one fixed-anchor objective, not to claim universal dominance over update-scheduling methods.

Score

Base

Fixed objective

MPI (K = 4)

62.97

53.71

100.75

Table 4: Compute-matched actor-update diagnostic on Hopper-medium-replay-v2. Larger refinement depth. We extend the locomotion medium-replay evaluation from K = 2 to K = 4 while retaining fixed step sizes within each run. The averages in Table 5 show improvement, saturation, and degradation across different base algorithms. Since each proximal subproblem is solved approximately, larger K can accumulate optimization error and imperfect centers. These results motivate the small-K interpretation used throughout the main text. Algorithm

πbase

π1

π2

π3

π4

IQL ReBRAC TD3+BC

71.9 71.2 58.3

78.8 75.5 79.7

81.2 73.1 80.1

80.9 73.1 81.0

83.1 61.9 80.3

Table 5: Locomotion medium-replay averages for larger refinement depth.

C

ADDITIONAL EXPERIMENTS

πbase

π1

π2

τ

HalfCheetah-M HalfCheetah-MR HalfCheetah-ME

48.54 ± 0.09 42.94 ± 1.44 94.44 ± 0.32

48.78 ± 0.23 44.86 ± 0.18 93.28 ± 2.56

48.74 ± 0.18 45.34 ± 0.70 90.30 ± 3.62

0.05 0.05 0.001

Hopper-M Hopper-MR Hopper-ME

58.70 ± 5.50 88.60 ± 10.68 111.30 ± 1.47

95.64 ± 7.31 102.04 ± 2.40 110.24 ± 2.42

93.86 ± 10.37 103.36 ± 1.02 108.68 ± 4.95

0.05 0.05 0.001

Walker2d-M Walker2d-MR Walker2d-ME

81.98 ± 4.25 79.02 ± 10.24 108.05 ± 0.62

81.90 ± 5.59 85.84 ± 3.67 −0.36 ± 0.03

85.04 ± 4.40 84.12 ± 7.08 −0.35 ± 0.02

0.05 0.05 0.01

Environment

Table 6: IQL (Fisher–Rao) locomotion with MPI. D4RL-normalized scores are mean±std over 5 seeds.

Environment

HalfCheetah-M

HalfCheetah-MR

HalfCheetah-ME

Hopper-M

Hopper-MR

Hopper-ME

Walker2d-M

Walker2d-MR

Walker2d-ME

πbase

π1

π2

τ

48.42 ± 0.19

48.54 ± 0.22 49.40 ± 0.29 50.22 ± 0.36 53.78 ± 0.24

48.32 ± 0.38 49.86 ± 0.35 51.14 ± 0.32 55.94 ± 0.67

0.001 0.005 0.010 0.050

43.84 ± 0.88

44.42 ± 0.36 44.96 ± 0.22 45.78 ± 0.26 49.80 ± 0.50

44.56 ± 0.54 45.26 ± 0.29 46.76 ± 0.43 52.30 ± 0.87

0.001 0.005 0.010 0.050

90.44 ± 4.06

91.02 ± 3.72 83.26 ± 11.32 63.26 ± 6.75 15.62 ± 6.05

88.58 ± 5.86 62.98 ± 10.53 35.20 ± 10.24 6.86 ± 3.46

0.001 0.005 0.010 0.050

60.44 ± 6.17

59.28 ± 6.00 63.14 ± 9.41 60.92 ± 12.91 90.54 ± 9.40

58.60 ± 6.05 61.66 ± 10.25 70.60 ± 10.62 101.40 ± 3.51

0.001 0.005 0.010 0.050

86.72 ± 7.58

91.64 ± 6.42 92.94 ± 11.94 92.32 ± 12.95 100.50 ± 3.36

98.24 ± 6.43 91.52 ± 9.84 100.64 ± 2.71 103.04 ± 0.42

0.001 0.005 0.010 0.050

102.20 ± 21.41

106.14 ± 10.82 95.36 ± 24.06 87.78 ± 42.29 28.58 ± 23.79

107.38 ± 10.84 93.82 ± 16.32 69.36 ± 56.39 14.56 ± 28.54

0.001 0.005 0.010 0.050

80.54 ± 3.57

82.38 ± 3.56 74.94 ± 15.15 86.68 ± 0.79 87.50 ± 0.66

82.84 ± 3.03 85.26 ± 3.34 76.22 ± 22.94 84.72 ± 5.48

0.001 0.005 0.010 0.050

78.06 ± 8.59

79.98 ± 4.26 81.82 ± 3.39 81.80 ± 6.16 88.52 ± 3.98

83.72 ± 3.50 84.46 ± 5.58 84.40 ± 6.96 83.36 ± 17.98

0.001 0.005 0.010 0.050

111.86 ± 1.06

111.72 ± 1.20 112.58 ± 0.73 112.62 ± 0.77 115.36 ± 0.65

112.10 ± 0.67 112.56 ± 1.04 113.78 ± 1.22 111.82 ± 8.93

0.001 0.005 0.010 0.050

Table 7: IQL (Wasserstein) locomotion sweep with MPI D4RL-normalized scores are mean±std over 5 seeds.

Environment

HalfCheetah-M

HalfCheetah-MR

HalfCheetah-ME

Hopper-M

Hopper-MR

Hopper-ME

Walker2d-M

Walker2d-MR

Walker2d-ME

πbase

π1

π2

τ

48.28 ± 0.22

48.06 ± 0.42 48.95 ± 0.43 49.80 ± 0.55 54.36 ± 0.29

48.36 ± 0.20 49.83 ± 0.34 51.24 ± 0.39 57.50 ± 0.61

0.001 0.005 0.010 0.050

44.77 ± 0.83

44.42 ± 0.55 45.28 ± 0.36 45.65 ± 0.58 49.85 ± 0.88

44.13 ± 0.72 45.51 ± 0.56 46.77 ± 0.52 51.99 ± 1.11

0.001 0.005 0.010 0.050

87.93 ± 7.96

85.34 ± 7.57 77.45 ± 4.90 69.03 ± 11.84 48.83 ± 10.28

72.72 ± 6.57 62.00 ± 11.18 47.58 ± 3.58 41.38 ± 9.74

0.001 0.005 0.010 0.050

58.63 ± 2.38

60.93 ± 4.38 63.26 ± 3.74 66.37 ± 4.49 89.81 ± 9.11

60.56 ± 3.00 62.49 ± 3.18 70.40 ± 6.96 101.30 ± 0.59

0.001 0.005 0.010 0.050

72.85 ± 23.01

76.40 ± 14.36 84.33 ± 14.03 89.84 ± 12.87 99.24 ± 1.53

76.40 ± 24.28 92.42 ± 4.50 99.12 ± 2.63 99.07 ± 1.62

0.001 0.005 0.010 0.050

94.63 ± 6.74

90.17 ± 3.89 100.85 ± 4.71 97.06 ± 9.45 60.11 ± 32.37

94.03 ± 9.65 101.03 ± 6.14 82.41 ± 12.77 16.69 ± 12.07

0.001 0.005 0.010 0.050

83.70 ± 0.99

83.01 ± 2.20 84.49 ± 0.71 84.95 ± 1.09 86.70 ± 1.58

84.61 ± 0.66 85.50 ± 1.22 85.50 ± 0.81 89.32 ± 2.19

0.001 0.005 0.010 0.050

81.73 ± 4.86

80.50 ± 5.66 82.71 ± 4.48 85.17 ± 1.99 89.63 ± 1.11

82.64 ± 7.92 84.98 ± 4.34 88.29 ± 1.93 89.24 ± 2.91

0.001 0.005 0.010 0.050

110.25 ± 0.14

110.21 ± 0.62 110.46 ± 0.44 110.92 ± 0.41 112.14 ± 0.65

110.98 ± 0.65 111.34 ± 1.24 111.80 ± 1.04 113.65 ± 1.40

0.001 0.005 0.010 0.050

Table 8: TD3+BC (Wasserstein) locomotion sweep with MPI D4RL-normalized scores are mean±std over 5 seeds.

Environment

HalfCheetah-M

HalfCheetah-MR

HalfCheetah-ME

Hopper-M

Hopper-MR

Hopper-ME

Walker2d-M

Walker2d-MR

Walker2d-ME

πbase

π1

π2

τ

48.19±0.17

48.60±0.23 49.59±0.34 50.47±0.31 48.60±0.90

48.84±0.24 50.27±0.42 51.70±0.34 49.44±0.55

0.001 0.005 0.010 0.050

44.39±1.06

45.46±0.74 46.08±0.70 46.81±0.90 45.71±1.26

45.50±0.91 47.06±0.89 48.02±0.58 45.73±1.09

0.001 0.005 0.010 0.050

89.23±5.94

90.03±3.69 82.22±9.88 74.87±9.71 28.43±4.93

87.29±9.89 69.27±13.67 64.72±6.78 26.27±7.26

0.001 0.005 0.010 0.050

57.94±1.67

59.81±3.85 58.54±3.40 66.32±3.97 73.30±7.81

59.19±6.40 62.78±6.34 69.51±6.12 87.24±8.30

0.001 0.005 0.010 0.050

73.47±23.62

76.23±21.45 78.59±21.51 84.72±22.27 97.34±3.91

76.13±20.87 85.58±22.26 89.83±18.60 99.20±0.77

0.001 0.005 0.010 0.050

89.71±14.34

100.82±15.26 107.48±7.20 107.30±6.14 74.46±12.29

109.89±3.34 110.53±2.02 108.20±4.10 36.41±14.85

0.001 0.005 0.010 0.050

83.07±2.95

83.65±2.94 85.00±2.46 84.56±3.13 87.91±1.05

85.61±1.48 85.15±3.31 86.99±1.06 89.08±1.15

0.001 0.005 0.010 0.050

80.56±10.59

85.54±2.78 85.14±3.60 85.68±4.30 89.73±1.98

81.21±6.58 87.41±2.17 88.89±1.84 84.10±8.36

0.001 0.005 0.010 0.050

110.24±0.16

110.26±0.17 110.50±0.32 110.72±0.43 111.80±0.44

110.34±0.26 110.71±0.42 111.04±0.47 112.89±0.46

0.001 0.005 0.010 0.050

Table 9: TD3+BC (Wasserstein) locomotion explicit-update sweep D4RL-normalized scores are mean±std over 5 seeds.

Environment

HalfCheetah-M

HalfCheetah-MR

HalfCheetah-ME

Hopper-M

Hopper-MR

Hopper-ME

Walker2d-M

Walker2d-MR

Walker2d-ME

πbase

π1

π2

τ

64.58 ± 1.01

65.80 ± 1.33 55.76 ± 23.91 66.54 ± 1.04 65.90 ± 0.92

65.52 ± 1.74 51.98 ± 29.51 65.14 ± 1.38 65.02 ± 1.14

0.001 0.005 0.010 0.050

49.56 ± 1.39

50.26 ± 1.09 50.18 ± 0.98 49.82 ± 0.79 50.44 ± 1.04

49.10 ± 1.76 49.72 ± 0.83 49.70 ± 0.65 49.84 ± 0.52

0.001 0.005 0.010 0.050

103.44 ± 4.24

97.34 ± 3.79 100.86 ± 1.28 98.06 ± 5.28 98.86 ± 4.98

95.38 ± 4.18 93.98 ± 4.55 90.74 ± 5.31 91.40 ± 7.23

0.001 0.005 0.010 0.050

101.14 ± 3.22

102.44 ± 0.29 100.24 ± 2.81 99.72 ± 3.24 89.96 ± 17.04

80.52 ± 42.82 93.36 ± 12.18 98.94 ± 5.48 99.06 ± 4.62

0.001 0.005 0.010 0.050

97.48 ± 6.93

100.64 ± 1.01 100.50 ± 0.51 99.04 ± 2.76 97.26 ± 7.70

98.36 ± 1.79 90.40 ± 12.74 99.36 ± 0.77 99.54 ± 1.33

0.001 0.005 0.010 0.050

110.62 ± 1.60

108.66 ± 5.13 108.18 ± 4.05 109.68 ± 3.57 109.78 ± 2.17

102.38 ± 8.64 104.76 ± 8.83 103.86 ± 9.50 99.54 ± 12.97

0.001 0.005 0.010 0.050

85.68 ± 0.58

85.66 ± 0.94 82.54 ± 2.88 84.20 ± 2.73 81.82 ± 4.33

86.76 ± 0.97 85.20 ± 1.58 77.90 ± 11.52 84.12 ± 3.51

0.001 0.005 0.010 0.050

80.12 ± 5.75

87.14 ± 4.63 78.48 ± 10.58 84.68 ± 5.71 78.08 ± 10.91

66.94 ± 21.35 64.52 ± 26.11 80.26 ± 18.38 70.16 ± 23.52

0.001 0.005 0.010 0.050

111.80 ± 0.21

112.06 ± 0.64 112.00 ± 0.71 112.06 ± 0.71 112.06 ± 0.69

112.92 ± 0.80 113.02 ± 0.47 113.14 ± 0.32 113.16 ± 0.52

0.001 0.005 0.010 0.050

Table 10: ReBRAC (Wasserstein) locomotion sweep with MPI D4RL-normalized scores are mean±std over 5 seeds.

πbase

π1

π2

τ

Umaze

100.00 ± 0.00 94.00 ± 5.48 98.00 ± 4.47 98.00 ± 4.47

94.00 ± 8.94 98.00 ± 4.47 98.00 ± 4.47 98.00 ± 4.47

96.00 ± 8.94 98.00 ± 4.47 98.00 ± 4.47 100.00 ± 0.00

0.001 0.005 0.010 0.050

Umaze-Diverse

80.00 ± 12.25 80.00 ± 23.45 92.00 ± 4.47 90.00 ± 7.07

88.00 ± 8.37 92.00 ± 8.37 96.00 ± 8.94 86.00 ± 16.73

80.00 ± 12.25 88.00 ± 16.43 92.00 ± 4.47 74.00 ± 24.08

0.001 0.005 0.010 0.050

Medium-Play

80.00 ± 18.71 80.00 ± 18.71 90.00 ± 12.25 80.00 ± 18.71

90.00 ± 10.00 86.00 ± 5.48 90.00 ± 12.25 88.00 ± 8.37

84.00 ± 13.42 84.00 ± 11.40 86.00 ± 11.40 92.00 ± 8.37

0.001 0.005 0.010 0.050

Medium-Diverse

84.00 ± 30.50 90.00 ± 7.07 76.00 ± 23.02 66.00 ± 27.93

76.00 ± 33.62 88.00 ± 13.04 84.00 ± 30.50 84.00 ± 15.17

74.00 ± 36.47 76.00 ± 11.40 84.00 ± 30.50 78.00 ± 32.71

0.001 0.005 0.010 0.050

Large-Play

48.00 ± 23.87 62.00 ± 30.33 72.00 ± 23.87 52.00 ± 29.50

54.00 ± 26.08 52.00 ± 30.33 82.00 ± 14.83 62.00 ± 31.14

72.00 ± 8.37 68.00 ± 37.68 68.00 ± 34.93 64.00 ± 25.10

0.001 0.005 0.010 0.050

Large-Diverse

70.00 ± 10.00 74.00 ± 18.17 50.00 ± 14.14 80.00 ± 14.14

64.00 ± 11.40 54.00 ± 16.73 64.00 ± 16.73 64.00 ± 15.17

62.00 ± 19.24 58.00 ± 21.68 56.00 ± 20.74 74.00 ± 15.17

0.001 0.005 0.010 0.050

Environment

Table 11: ReBRAC (Wasserstein) AntMaze sweep with MPI D4RL-normalized scores are mean±std over 5 seeds.

Environment

HalfCheetah-M

HalfCheetah-ME

HalfCheetah-MR

Hopper-M

Hopper-ME

Hopper-MR

Walker2d-M

Walker2d-ME

Walker2d-MR

πbase

π1

π2

τ /w

64.43±1.35

64.45±1.71 64.62±1.03 64.71±1.12 64.88±2.82

64.77±1.68 64.91±0.97 64.95±1.14 65.05±2.85

10 50 100 500

104.05±6.46

43.07±1.17 102.71±3.56 104.26±4.32 105.26±2.37

37.85±5.12 101.27±1.71 105.11±3.96 104.92±4.19

10 50 100 500

50.06±1.39

50.22±0.83 50.22±0.93 50.97±0.89 39.14±3.42

50.54±0.41 50.68±0.51 51.29±0.97 37.09±5.06

0.001 0.005 0.010 0.050

99.29±5.01

101.74±1.30 102.08±1.13 102.17±1.08 99.85±4.91

101.90±1.51 102.36±0.66 101.81±1.91 94.35±12.12

0.001 0.005 0.010 0.050

111.28±0.85

110.49±1.47 106.04±7.16 90.50±12.40 60.50±17.50

108.19±3.55 101.83±15.08 88.73±13.92 56.14±22.58

0.001 0.005 0.010 0.050

87.82±18.54

96.27±8.73 93.38±8.62 95.06±12.06 100.85±0.49

96.26±8.71 99.38±2.74 100.89±0.46 101.25±0.48

0.001 0.005 0.010 0.050

85.54±0.44

85.78±0.59 85.68±0.43 85.17±2.20 87.19±0.29

83.77±1.58 86.21±0.21 86.45±0.38 85.18±6.23

0.001 0.005 0.010 0.050

111.80±0.26

111.84±0.20 111.79±0.32 111.90±0.26 112.45±0.40

111.84±0.22 111.83±0.30 112.02±0.29 112.81±0.73

0.001 0.005 0.010 0.050

79.63±6.40

78.19±2.92 82.48±6.66 87.43±2.05 89.50±1.79

81.62±8.46 84.52±12.35 88.66±2.64 87.53±7.23

0.001 0.005 0.010 0.050

Table 12: ReBRAC (Wasserstein) locomotion explicit-update sweep. D4RL-normalized scores are mean±std over 5 seeds.

πbase

π1

π2

τ

Umaze

95.80±3.35 97.20±1.64 96.50±1.00 97.40±0.55

96.80±1.79 96.40±1.95 97.50±1.29 97.00±2.45

97.20±0.84 98.00±1.22 96.25±2.06 96.60±1.52

0.001 0.005 0.010 0.050

Umaze-Diverse

82.60±7.02 87.00±7.97 84.50±8.39 82.80±10.55

83.60±7.37 84.60±7.33 85.25±9.00 86.80±8.01

84.60±6.66 84.80±8.64 85.50±7.19 82.60±8.88

0.001 0.005 0.010 0.050

Medium-Play

83.20±6.61 86.80±2.86 84.25±6.18 87.20±8.11

86.60±4.56 87.80±4.97 89.25±1.71 84.40±7.86

84.00±4.53 89.00±3.54 86.25±4.57 87.60±3.44

0.001 0.005 0.010 0.050

Medium-Diverse

81.20±21.38 88.20±4.76 76.25±18.96 77.20±20.38

80.20±19.80 85.00±6.63 76.75±17.50 77.20±15.45

75.60±23.75 86.00±6.78 73.00±21.15 80.60±13.79

0.001 0.005 0.010 0.050

Large-Play

62.20±19.51 59.60±21.29 66.75±2.99 56.80±19.52

58.00±20.06 59.80±20.68 70.50±2.08 55.80±18.47

61.20±18.85 61.20±22.45 69.75±3.30 58.60±20.28

0.001 0.005 0.010 0.050

Large-Diverse

68.80±10.57 64.60±10.36 61.00±4.55 67.00±9.82

66.80±11.73 67.20±10.06 64.50±8.02 65.40±8.20

68.80±8.90 70.40±9.56 64.75±6.95 66.20±6.42

0.001 0.005 0.010 0.050

Environment

Table 13: ReBRAC (Wasserstein) AntMaze explicit-update sweep. D4RL-normalized scores are mean±std over 5 seeds.

D

TRAINING CURVES

We report D4RL-normalized scores versus time steps for the baseline algorithm, where (+) and (++) denote π1 , and π2 , respectively. Shaded regions indicate the standard deviation over seeds. Figure 3: IQL (Wasserstein) training curves on D4RL locomotion.

(a) HalfCheetah-M

(b) HalfCheetah-MR

(c) HalfCheetah-ME

(d) Hopper-M

(e) Hopper-MR

(f) Hopper-ME

(g) Walker2d-M

(h) Walker2d-MR

(i) Walker2d-ME

Figure 4: IQL (Fisher-Rao) training curves on D4RL locomotion.

(a) HalfCheetah-M

(b) HalfCheetah-MR

(c) HalfCheetah-ME

(d) Hopper-M

(e) Hopper-MR

(f) Hopper-ME

(g) Walker2d-M

(h) Walker2d-MR

(i) Walker2d-ME

Figure 5: TD3+BC (Wasserstein) training curves on D4RL locomotion.

(a) HalfCheetah-M

(b) HalfCheetah-MR

(c) HalfCheetah-ME

(d) Hopper-M

(e) Hopper-MR

(f) Hopper-ME

(g) Walker2d-M

(h) Walker2d-MR

(i) Walker2d-ME

Figure 6: ReBRAC (Wasserstein) training curves on D4RL locomotion.

(a) HalfCheetah-M

(b) HalfCheetah-MR

(c) HalfCheetah-ME

(d) Hopper-M

(e) Hopper-MR

(f) Hopper-ME

(g) Walker2d-M

(h) Walker2d-MR

(i) Walker2d-ME

Figure 7: ReBRAC (Wasserstein) training curves on D4RL AntMaze.

(a) Umaze

(b) Umaze-Diverse

(c) Medium-Play

(d) Medium-Diverse

(e) Large-Play

(f) Large-Diverse

Related documents

• On- and Off-Policy Monotonic Policy Improvement#478153
• Lifted Bellman Linear Programming for Offline Reinforcement Learning#1050662
• Lifted Bellman Linear Programming for Offline Reinforcement Learning#1028719
• Second-Order Policy Effects as State Transitions: A Source-Linked Benchmark for Policy Simulation#612528
• Second-Order Policy Effects as State Transitions: A Source-Linked Benchmark for Policy Simulation#611051