Geometric Time Score Estimators: Bridging Self-Predictive Representations and Generative Occupancy Models For Reinforcement Learning

Aravind Venugopal; Grace Liu; Jeff Schneider; Benjamin Eysenbach

cs.LGOct 28, 2026

Geometric Time Score Estimators: Bridging Self-Predictive Representations and Generative Occupancy Models For Reinforcement Learning

Authors: Aravind Venugopal, Grace Liu, Jeff Schneider, Benjamin Eysenbach

Organizations: Carnegie Mellon University · Princeton University

Abstract

In reinforcement learning (RL), temporal reasoning is essential for estimating the density of future outcomes under a policy's occupancy measure. While self-predictive representations have been a dominant paradigm for occupancy density estimation, we find that they yield surprisingly poor estimates. On the other hand, generative occupancy models can produce remarkably accurate samples from the occupancy measure. Yet, they cannot easily recover occupancy densities, limiting their use for control despite their expressivity. In this paper, we propose Geometric Time Score Estimators (GTSE), a framework that bridges the two paradigms by learning occupancy density estimators along probability paths. We analyze how GTSE implicitly induces better occupancy estimators through improved representation learning and show how to directly extract self-predictive representations from this framework. Through experiments across several locomotion and manipulation tasks, we show that GTSE improves performance over state-of-the-art methods by $2\times$ to $6\times$ in goal-conditioned RL settings, and that the representations learned using our framework transfer effectively to downstream RL tasks. Website: https://aravindvenu7.github.io/website/gtse/

Figures & tables

Figure 1: The time score function, $\partial_{k}\log p_{k}(g^{+}\mid s,a)$ , captures the instantaneous rate of change of the log-density of the occupancy measure along a probability path. Integrating it induces accurate occupancy density estimates useful for downstream control applications.

Figure 2: Normalized rank of estimated occupancy density vs. the normalized rank of ground-truth $\log d^{\pi_{D}}$ over 50 goals and 300 state-action pairs per goal on a 5-step toy MDP (colored by distance of state from origin). GTSE variants are introduced in Secs. 5 – 6 .

	CRL	FB	TD-JEPA	BYOL- $\gamma$	Dual GR	TMD	TD-flows	GTSE- $b_{0}$	GTSE- $b_{2}$	GTSE- $b_{10}$	GTSE- $b$	GTSE- $r$	GTSE
$\rho$ (near)	$0.56$	$0.11$	$0.25$	$0.33$	$-0.07$	$0.56$	$0.00$	$0.56$	$0.60$	$0.60$	$0.60$	$0.60$	$\mathbf{0.63}$
$\rho$ (far)	$0.14$	$0.00$	$0.08$	$0.16$	$-0.09$	$0.30$	$0.05$	$0.50$	$0.51$	$0.52$	$0.53$	$0.52$	$\mathbf{0.55}$
Success rate	$3\pm 6$	$0\pm 0$	$0\pm 0$	$0\pm 0$	$0\pm 0$	$0\pm 0$	$0\pm 0$	$12\pm 5$	$21\pm 10$	$19\pm 9$	$25\pm 4$	$25\pm 2$	$\mathbf{38\pm 4}$

Table 1: Spearman rank correlation

\rho

with ground-truth

\log d^{\pi_{D}}

for nearby and far-away goals along with overall average success rate (%) across the 5 test-time goals on pointmaze-large-diverse over 4 seeds. We report the 95% bootstrapped C.I. half-width after the

\pm

sign. GTSE variants are introduced in Secs. 5 – 6 .

Dataset

GCBC

CRL

TD-JEPA

BYOL-

\gamma

Dual GR

TMD

SMORE

TD-flows

GTSE-

r

GTSE

scene-play

20\pm 4

25\pm 5

6\pm 2

31\pm 4

17\pm 2

19\pm 15

7\pm 2

13\pm 4

35\pm 4

9\pm 3

60\pm 5

\mathbf{69\pm 6}

scene-noisy

10\pm 1

7\pm 3

2\pm 1

24\pm 3

9\pm 3

12\pm 7

7\pm 1

3\pm 2

19\pm 4

5\pm 2

37\pm 4

\mathbf{55\pm 6}

puzzle-4x4-play

1\pm 1

2\pm 1

22\pm 8

\mathbf{36\pm 3}

2\pm 1

17\pm 6

1\pm 1

0\pm 1

4\pm 2

1\pm 1

26\pm 6

\mathbf{40\pm 4}

puzzle-4x4-noisy

1\pm 0

0\pm 0

6\pm 2

9\pm 3

0\pm 0

2\pm 3

0\pm 0

0\pm 0

1\pm 0

1\pm 0

9\pm 2

\mathbf{22\pm 2}

cube-triple-play

5\pm 2

2\pm 1

1\pm 0

1\pm 1

0\pm 1

13\pm 2

1\pm 1

2\pm 3

16\pm 3

0\pm 0

7\pm 1

\mathbf{18\pm 1}

cube-triple-noisy

2\pm 1

1\pm 1

0\pm 1

0\pm 0

0\pm 1

4\pm 2

0\pm 0

1\pm 1

\mathbf{8\pm 2}

1\pm 1

4\pm 1

\mathbf{8\pm 1}

Table 2: GCRL tasks: Overall average success rate (%) across the 5 test-time goals over 8 seeds per task per algorithm. We report the 95% bootstrapped C.I. half-width after the ± sign.

Task	CRL	FB	TD-JEPA	BYOL- $\gamma$	Dual GR	TMD	Vanilla FQL	GTSE- $r$
task 1	67 $\pm$ 6	58 $\pm$ 33	98 $\pm$ 5	51 $\pm$ 27	37 $\pm$ 7	1 $\pm$ 2	100 $\pm$ 0	100 $\pm$ 1
task 2	14 $\pm$ 12	17 $\pm$ 19	32 $\pm$ 22	8 $\pm$ 8	4 $\pm$ 3	0 $\pm$ 0	83 $\pm$ 23	81 $\pm$ 12
task 3	27 $\pm$ 13	1 $\pm$ 2	85 $\pm$ 11	10 $\pm$ 6	11 $\pm$ 9	0 $\pm$ 0	99 $\pm$ 3	94 $\pm$ 5
task 4	10 $\pm$ 5	0 $\pm$ 0	18 $\pm$ 14	10 $\pm$ 4	7 $\pm$ 4	0 $\pm$ 0	9 $\pm$ 7	36 $\pm$ 18
task 5	2 $\pm$ 2	0 $\pm$ 0	0 $\pm$ 0	0 $\pm$ 0	4 $\pm$ 7	0 $\pm$ 0	0 $\pm$ 0	0 $\pm$ 0
Mean	24.0	15.2	46.6	15.8	12.6	0.2	58.2	62.2

Table 3: Reward-based RL tasks with frozen representation encoders on scene-play . We report success rate (%) with the 95% bootstrapped C.I. half-width after the

\pm

sign, averaged over 4 seeds.

Figure 3: Effect of

K_{max}

. GTSE performance over varying

K_{max}

, the upper limit of the time-score integral, on puzzle-4x4-play (left) and antmaze-large-navigate (right). Lines show the mean over 8 seeds; shaded regions show 95% C.I.

Figure 4: Effect of

N

. GTSE performance over varying

N

, time-score evaluations during inference, on puzzle-4x4-play (left) and antmaze-large-navigate (right). Lines show the mean over 8 seeds; shaded regions show 95% C.I.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Training Iterations
toy MDP	250k
puzzle-4x4-play	1M
scene-play	1M
cube-triple-play	3M
antmaze-large-navigate	3M
pointmaze-large-diverse	3M

Appendix

Table 4: Training iterations per task. All algorithms are trained for the same number of iterations.

Hyperparameter	Value
Optimizer	Adam ( Kingma and Ba, 2014 )
Learning rate	0.0003
Batch size	1024
Architecture	MLP
MLP size	{512, 512, 512, 512}
Representation dimensions (GCRL expts)	512

Appendix

Table 5: Common hyperparameters.

Hyperparameter	Value
Nonlinearity	ELU ( Clevert et al., 2015 ) for GTSE, GELU for GTSE- $r$
Discount factor $\gamma$	0.995 for cube-triple and cube-triple-noisy, 0.99 otherwise
Time embedding size	64 (toy-MDP and pointmaze) & 96 (other experiments)
Goal noise std.	0.0001

Appendix

Table 6: GTSE hyperparameters.

Dataset	$\tau$ (target update rate)	GTSE ( $p^{\mathcal{D}}$ )	N (Eq. 6 )	policy ( $p^{\mathcal{D}}_{\mathrm{cur}},p^{\mathcal{D}}_{\mathrm{traj}},p^{\mathcal{D}}_{\mathrm{rand}}$ )	$K_{max}$
toy MDP	0.005	(0, 1, 0)	30	(0, 1, 0)	0.9999
puzzle-4x4-play	0.005	(0, 1, 0)	30	(0, 1, 0)	0.99
puzzle-4x4-noisy	0.005	(0, 1, 0)	30	(0, 1, 0)	0.95
cube-triple-play	0.003	(0, 1, 0)	30	(0, 0.5, 0.5)	0.925
cube-triple-noisy	0.005	(0, 1, 0)	30	(0, 0.5, 0.5)	0.925
scene-play	0.003	(0, 1, 0)	30	(0, 0.5, 0.5)	0.925

Appendix

Table 7: Task-specific hyperparameters for GTSE training

	CRL	FB	TD-JEPA	BYOL- $\gamma$	Dual GR	TMD	SMORE	TD-flows	NF	GTSE- $b$	GTSE- $r$	GTSE
Training (ms) $\downarrow$	3.12	3.66	5.12	3.77	4.51	44.55	2.53	2.26	4.62	4.49	3.59	2.17
Inference (ms) $\downarrow$	0.24	0.29	0.45	0.29	0.26	0.23	0.19	0.70	0.64	0.55	0.36	0.31

Appendix

Table 8: Wall-clock runtime comparison reported in milliseconds.

Figure 5: Difference in time score per dimension, averaged over $k$ , between the in-trajectory goal and randomly sampled goals on antmaze-large-navigate .

Dataset	GTSE- $b_{0}$	GTSE- $b_{2}$	GTSE- $b_{10}$	GTSE-b
puzzle-4x4-play	$36\pm 3$	$30\pm 2$	$29\pm 7$	$30\pm 3$
antmaze-large-navigate	$63\pm 5$	$77\pm 4$	$77\pm 3$	$77\pm 1$

Appendix

Table 9: Performance comparison of GTSE-

b_{L}

and GTSE-

b

on OGBench GCRL tasks. While performance improves with truncation order on antmaze , it degrades with truncation order on puzzle .

Table 15

Predictive Representation Learning

Predictive representation learning trains models to encode useful structure by predicting unseen or future representations from context.

See all 176 papers

Occupancy Reward Shaping: Improving Credit Assignment for Offline Goal-Conditioned Reinforcement Learning
Apr 22, 2026Aravind Venugopal, Jiayu Chen, Xudong Wu +3Reward Shaping Temporal Credit Assignment
Carnegie Mellon University · The University of Hong Kong · INFIFORCE Intelligent Technology +1
GEM: Gaussian Evolution Model for Occupancy Forecasting and Motion Planning
May 17, 2026Cheng Chen, Hao Huang, Saurabh Bagchi3D Gaussian Splatting World Model Learning
Purdue University · New York University Abu Dhabi
Occupancy Modeling via Drifting: One-Step Generation Enables Direct Policy Extraction
Oct 14, 2026Robin Uhrich, Marco Bagatella, Núria Armengol Urpí +2Reinforcement Learning Policy Learning
University of Tübingen · ETH Zürich · Max Planck Institute for Intelligent Systems

Apr 22, 2026J/K move · Enter open · S save

Occupancy Reward Shaping: Improving Credit Assignment for Offline Goal-Conditioned Reinforcement Learning

Aravind Venugopal, Jiayu Chen, Xudong Wu, Chongyi Zheng+2

Carnegie Mellon University · The University of Hong Kong · INFIFORCE Intelligent Technology+1

Reward Shaping Temporal Credit Assignment Offline RL Sparse-Reward RL Policy Learning Reward Design for RL+1

Open paper AskPDF arXiv

The temporal lag between actions and their long-term consequences makes credit assignment a challenge when learning goal-directed behaviors from data. Generative world models capture the distribution of future states an agent may visit, indicating that they have captured temporal information. How can that temporal information be extracted to perform credit assignment? In this paper, we formalize how the temporal information stored in world models encodes the underlying geometry of the world. Leveraging optimal transport, we extract this geometry from a learned model of the occupancy measure into a reward function that captures goal-reaching information. Our resulting method, Occupancy Reward Shaping, largely mitigates the problem of credit assignment in sparse reward settings. ORS provably does not alter the optimal policy, yet empirically improves performance by 2.2x across 13 diverse long-horizon locomotion and manipulation tasks. Moreover, we demonstrate the effectiveness of ORS in the real world for controlling nuclear fusion on 3 Tokamak control tasks. Code: https://github.com/aravindvenu7/occupancy_reward_shaping; Website: https://aravindvenu7.github.io/website/ors/

Abstract

Predictive Representation Learning

Explore similar work

Occupancy Reward Shaping: Improving Credit Assignment for Offline Goal-Conditioned Reinforcement Learning

GEM: Gaussian Evolution Model for Occupancy Forecasting and Motion Planning

Occupancy Modeling via Drifting: One-Step Generation Enables Direct Policy Extraction