RepRL: Theory-Informed Reinforcement Learning

• RepRL is a reinforcement learning framework that embeds policies into a compact latent space, reformulating exploration as a representation-exploitation problem.
• It employs a variational encoder-decoder architecture to predict policy values linearly in the latent space, using UCB or Thompson sampling for efficient bandit-based policy selection.
• Empirical results on benchmarks like MuJoCo and MinAtar demonstrate that RepRL improves exploration efficiency and stability over traditional methods through precise policy representation.

Searching arXiv for the primary paper and closely related theory-informed RL frameworks to ground the article with current citations. Representation-Driven Reinforcement Learning (RepRL) is a theory-informed reinforcement-learning framework that represents policies as estimates of their expected values and leverages techniques from contextual bandits to guide exploration and exploitation. Its central claim is that embedding a policy network into a linear feature space allows the exploration-exploitation problem to be reframed as a representation-exploitation problem, so that good policy representations enable optimal exploration. In this formulation, the critical object is not a state-action feature map, but a latent representation of entire policies, on top of which optimism or posterior sampling can be applied for policy selection (Nabati et al., 2023).

1. Formal setting and value-based policy representation

RepRL is formulated on an infinite-horizon discounted Markov Decision Process

$M = (S, A, r, T, B, \gamma),$

where $S$ is a possibly large state space, $A$ is the action space, $T(s'|s,a)$ is the transition kernel, $r(s,a)\in[0,1]$ is the reward function, $B$ is the initial state distribution, and $\gamma\in[0,1)$ is the discount factor (Nabati et al., 2023). A stationary stochastic policy $\pi_\theta$ maps states to distributions over actions, $\pi_\theta(a|s)$, and the optimization objective is $J(\pi)=V(\pi)$ in episodic or continuing discounted tasks.

The discounted return of a trajectory generated by $S$0 is

$S$1

with $S$2, $S$3, and $S$4. The policy value is

$S$5

and the state-value function is

$S$6

RepRL also uses the discounted state-action occupancy measure

$S$7

The foundational representational statement is the linear embedding assumption:

$S$8

This exactly holds with $S$9 by taking $A$0 and $A$1, yielding

$A$2

but RepRL targets $A$3 (Nabati et al., 2023). The framework therefore treats policy value prediction as a problem of learning a compact latent representation in which value is approximately linear.

2. Learned latent space and variational objective

RepRL parameterizes a representation encoder

$A$4

that maps policy parameters $A$5 to a latent vector $A$6. Values are then predicted linearly in $A$7 through a linear-Gaussian decoder,

$A$8

where $A$9 are linear decoder parameters (Nabati et al., 2023). This induces a per-policy value representation shared across policies through the same linear predictor.

The encoder can be stochastic or deterministic. With deterministic $T(s'|s,a)$0, the representation reduces to $T(s'|s,a)$1. With stochastic $T(s'|s,a)$2, RepRL trains the representation by the evidence lower bound

$T(s'|s,a)$3

with a Gaussian prior $T(s'|s,a)$4 and linear-Gaussian likelihood $T(s'|s,a)$5 (Nabati et al., 2023). In the implementation-oriented form used for learning,

$T(s'|s,a)$6

Within this formulation, representation quality is defined by the realizability error

$T(s'|s,a)$7

High-quality representations yield tight confidence sets and well-calibrated optimism. This makes representation learning the central lever of the framework rather than an auxiliary component (Nabati et al., 2023).

3. Contextual-bandit grounding and theoretical guarantees

RepRL grounds policy selection in linear contextual bandits. After collecting latent-return pairs $T(s'|s,a)$8, it maintains the regularized least-squares statistics

$T(s'|s,a)$9

Uncertainty is quantified by

$r(s,a)\in[0,1]$0

and upper-confidence action selection is

$r(s,a)\in[0,1]$1

The Thompson-sampling alternative samples

$r(s,a)\in[0,1]$2

and selects the embedding maximizing $r(s,a)\in[0,1]$3 (Nabati et al., 2023).

The crucial embedding step is that, at episode $r(s,a)\in[0,1]$4, RepRL constructs a decision set $r(s,a)\in[0,1]$5 of candidate policies or candidate embeddings, evaluates each candidate by $r(s,a)\in[0,1]$6 or Thompson sampling, and selects the optimistic policy for rollout. The latent context is therefore the policy representation $r(s,a)\in[0,1]$7 rather than a state-action feature. Bandit updates occur once per episode after observing a scalar return signal $r(s,a)\in[0,1]$8 (Nabati et al., 2023).

Under standard linear bandit assumptions—linear realizability, bounded features $r(s,a)\in[0,1]$9, $B$0-regularization, and sub-Gaussian reward noise—the cumulative regret after $B$1 selections satisfies

$B$2

This guarantee applies to the bandit stage of selecting policies through the latent embedding. RepRL does not establish regret bounds for the full RL process with function approximation and MDP-induced dependencies; its guarantees are confined to the policy-selection bandit stage (Nabati et al., 2023). This suggests that the framework is theoretically strongest as a representation-mediated selection mechanism, rather than as a complete end-to-end RL theory.

4. Algorithmic instantiations and inner trajectory sampling

RepRL is instantiated in two main algorithmic forms. Representation-Driven Evolution Strategies (RepES) samples $B$3 Gaussian perturbations around the current parameters $B$4,

$B$5

maps each candidate to $B$6, scores it by UCB or Thompson sampling, and updates parameters using a soft weighted finite-difference gradient,

$B$7

$B$8

Mixing with standard ES gradients can improve stability; the reported example is 20% ES and 80% RepRL bandit gradient (Nabati et al., 2023).

Representation-Driven Policy Gradient (RepPG) starts from a baseline policy-gradient loss $B$9 from an on-policy method such as PPO, and augments it with a regularizer toward an optimistic parameter $\gamma\in[0,1)$0 chosen from a decision set via the bandit score:

$\gamma\in[0,1)$1

$\gamma\in[0,1)$2

An alternative conceptual integration uses a latent-uncertainty exploration bonus, but the regularizer is the primary mechanism in the framework (Nabati et al., 2023).

A distinctive technical device is inner trajectory sampling. Rather than relying only on initial-state returns, RepRL defines

$\gamma\in[0,1)$3

and proves the equivalence

$\gamma\in[0,1)$4

Sampling along trajectories from $\gamma\in[0,1)$5 therefore yields an unbiased target for the linear value model and can be used for representation learning and bandit updates (Nabati et al., 2023). In practice, states are sampled uniformly along stored trajectories, and for truncated rollouts the critic estimates the remainder.

The overall update schedule is alternating. Each episode maps $\gamma\in[0,1)$6 to $\gamma\in[0,1)$7, builds a decision set, selects $\gamma\in[0,1)$8 via the bandit algorithm, rolls out to collect returns, and updates both the encoder-decoder pair $\gamma\in[0,1)$9 and the bandit statistics

$\pi_\theta$0

The inverse can be maintained by Sherman–Morrison updates if desired (Nabati et al., 2023).

5. Empirical behavior and benchmark results

The framework is evaluated on GridWorld, Sparse HalfCheetah, MuJoCo tasks, and MinAtar. In a nonuniform, noisy GridWorld reward landscape with a sparse high-reward goal, RepRL explored more efficiently than ES and found the optimal path. The reported t-SNE visualizations showed that the learned latent representation organizes policies by value better than raw parameter space (Nabati et al., 2023).

On Sparse HalfCheetah, policy-space, latent-space with gradient-based inversion, and history-based decision sets performed similarly, while policy-space sampling offered stable, simple implementations. On the MuJoCo suite, RepES matched or significantly outperformed ES in standard tasks and achieved distant goals more reliably than ES and SAC in sparse variants, despite occasional successes via random search in ES. The sparse settings awarded rewards at distance intervals with a control cost, creating strong exploration challenges (Nabati et al., 2023).

On MinAtar, RepPG regularized PPO and outperformed PPO across tasks such as SpaceInvaders and Breakout, particularly under sparse rewards. Inner trajectory sampling facilitated seamless integration with chunked rollouts (Nabati et al., 2023). Across these experiments, the empirical pattern is that policy-representation quality rather than hand-crafted exploration noise determines the usefulness of optimism or posterior sampling in the search process.

Practical details reported for these experiments include latent dimensions such as $\pi_\theta$1 in MuJoCo, encoder-decoder learning rates such as Adam with $\pi_\theta$2, policy-space Gaussian decision sets with variance $\pi_\theta$3, and exact inverse maintenance by Sherman–Morrison updates when feasible (Nabati et al., 2023).

6. Positioning within theory-informed RL, practical guidance, and limitations

RepRL is positioned against several neighboring lines of work. Relative to neural contextual bandits, prior work maps states to features for action-level exploration, whereas RepRL maps entire policies to a linear space to drive exploration in policy space. Relative to successor features and linear MDPs, it likewise exploits linear structure for value prediction, but does so through a learned policy embedding explicitly optimized for linear value prediction and bandit-derived optimism (Nabati et al., 2023). In that sense, its theory-informed element is the importation of OFUL- or Thompson-style policy selection into RL through learned policy representations.

Other theory-informed RL frameworks organize theory differently. "A Game Theoretic Framework for Model Based Reinforcement Learning" casts model-based RL as a Stackelberg game between a policy player and a model player, with approximate equilibrium guarantees and leader-follower algorithm families such as Policy-as-Leader and Model-as-Leader (Rajeswaran et al., 2020). "Human-Level Reinforcement Learning through Theory-Based Modeling, Exploration, and Planning" formulates Theory-Based Reinforcement Learning through structured, causal, object-centric world models expressed in VGDL and used for exploration and planning (Tsividis et al., 2021). "Learning by Doing: An Online Causal Reinforcement Learning Framework with Causal-Aware Policy" instead places a causal graphical model over factorized state variables and uses interventions plus causal masks to restrict the action search space (Cai et al., 2024). "LESSON: Learning to Integrate Exploration Strategies for Reinforcement Learning via an Option Framework" uses an option-critic architecture to learn a policy over predefined exploration strategies, again separating the theory-guided exploration mechanism from the task policy (Kim et al., 2023). These comparisons indicate that “theory-informed” in contemporary RL can refer to bandit realizability assumptions, game-theoretic bi-level optimization, structured causal simulators, causal graphical constraints, or formal option decompositions, depending on which aspect of exploration or control is being regularized.

RepRL’s main failure modes are explicit. If $\pi_\theta$4 is poorly approximated by a linear function of $\pi_\theta$5, optimism may be miscalibrated and both exploration and stability can degrade. Latent decision sets can be problematic because inverting $\pi_\theta$6 is challenging, and off-manifold parameters can lead to overestimation. Learned representations may drift under non-stationarity or partial observability, and the sub-Gaussian and independence assumptions inherited from linear bandits can be violated by temporally correlated RL data. Extremely sparse environments still remain difficult, and there is sensitivity to optimism parameters $\pi_\theta$7 or $\pi_\theta$8, decision-set variance $\pi_\theta$9, latent dimension $\pi_\theta(a|s)$0, and the regularization weight $\pi_\theta(a|s)$1 in RepPG (Nabati et al., 2023).

The practical guidance follows directly from these limitations. The framework recommends starting with a deterministic encoder $\pi_\theta(a|s)$2 producing bounded-norm $\pi_\theta(a|s)$3, training $\pi_\theta(a|s)$4 via the ELBO, validating representation quality by regressing observed returns on $\pi_\theta(a|s)$5, and monitoring $\pi_\theta(a|s)$6 together with calibration of the confidence widths $\pi_\theta(a|s)$7 (Nabati et al., 2023). It also recommends tracking predictive residuals $\pi_\theta(a|s)$8, the diversity of selected parameters across episodes, and ablations over decision-set size, variance, latent dimension, and optimism scale. This suggests that RepRL is best understood not as a replacement for policy optimization, but as a structured representational layer that determines whether bandit-style exploration in policy space is statistically meaningful.