1. Reinforcement Learning as an MDP

An MDP is the minimal mathematical story of an agent acting over time. It separates the world into states, possible actions, how the world changes, what the agent is rewarded for, and how much the future matters. This compact tuple is useful because every later model in this note modifies one of these ingredients.

The standard mathematical formulation is a Markov Decision Process:

Read this equation as a contract between the agent and the environment. The agent can only choose from \(\mathcal{A}\), the environment responds according to \(P\), and success is measured through \(R\). The Markov assumption means that \(s_t\) contains all information needed for predicting the next step; the model does not need the full history once the current state is known.

where:

• \(\mathcal{S}\) is the state space.
• \(\mathcal{A}\) is the action space.
• \(P(s' \mid s,a)\) is the transition probability.
• \(R(s,a,s')\) is the reward function.
• \(\gamma \in [0,1)\) is the discount factor.

At time \(t\), the agent observes \(s_t\), chooses an action from a policy, receives reward, and transitions to the next state:

The policy \(\pi\) is the agent's behavior rule. The notation \(a_t \sim \pi(a \mid s_t)\) says that the action may be sampled from a probability distribution, so the same state can sometimes lead to different actions. This is important during exploration and in environments where deterministic behavior is too brittle.

The transition model \(P\) is the environment's response. The dot in \(P(\cdot \mid s_t,a_t)\) means the next state is drawn from the entire distribution over possible next states, not chosen by the agent directly.

The reward is observed after the transition because the consequence of an action often depends on where the action actually leads. For example, the same steering action can be good or bad depending on the next position of the vehicle.

The discounted return is:

The return \(G_t\) folds the whole future into one number. The discount \(\gamma\) controls how quickly future rewards fade in importance: when \(\gamma\) is close to 0, the agent behaves myopically; when it is close to 1, the agent is encouraged to plan farther ahead.

The goal is to find a policy \(\pi\) maximizing expected return:

The expectation is needed because trajectories may be random. Randomness can come from the policy, the environment, partial observability, or all three. Optimizing \(J(\pi)\) therefore means choosing behavior that performs well on average over many possible rollouts.

An optimal policy is:

The symbol \(\arg\max\) returns the policy that achieves the largest objective value. There can be more than one optimal policy, so \(\pi^*\) is written as a member of the set of maximizers rather than as a unique object.

2. Bellman Equations

Bellman equations are the bridge between the global objective \(J(\pi)\) and step-by-step decision making. Instead of evaluating an entire infinite trajectory at once, they say: estimate what happens now, then add the value of the state you land in. This recursive structure is the backbone of dynamic programming, temporal-difference learning, Q-learning, actor-critic methods, and many planning algorithms.

State-value function

The state-value function answers: if the agent starts in state \(s\) and then follows policy \(\pi\), how much discounted reward should it expect? It evaluates states under a fixed behavior rule.

Bellman expectation equation for \(V^\pi\):

This equation expands \(V^\pi(s)\) by averaging over two sources of uncertainty. First, \(\sum_a \pi(a \mid s)\) averages over the policy's possible actions. Second, \(\sum_{s'} P(s' \mid s,a)\) averages over the environment's possible next states. Inside the brackets, \(R(s,a,s')\) is the immediate payoff and \(\gamma V^\pi(s')\) is the discounted value of continuing from the next state.

Action-value function

The action-value function asks a slightly sharper question: if the agent starts in state \(s\), takes action \(a\) first, and only then follows \(\pi\), what return should it expect? This is often more directly useful for control because it lets us compare actions from the same state.

Bellman expectation equation for \(Q^\pi\):

The first action is fixed here, so there is no outer sum over actions at state \(s\). After the environment moves to \(s'\), the policy resumes, and \(\sum_{a'} \pi(a' \mid s')Q^\pi(s',a')\) averages over the next action choices.

Optimal Bellman equations

The starred functions describe best-achievable value rather than the value of a particular policy. \(V^*(s)\) is the best long-term return available from state \(s\), while \(Q^*(s,a)\) is the best long-term return after committing to action \(a\) first.

The optimal state-value equation replaces policy averaging with a maximization over actions. It says that an optimal agent chooses the action whose expected immediate reward plus future value is largest.

The optimal action-value equation has the same idea one step later. After taking \(a\) and arriving in \(s'\), the agent assumes it will choose the best next action \(a'\), hence the \(\max_{a'} Q^*(s',a')\) term.

The optimal greedy policy can be written as:

This final line turns value estimation into action selection. Once \(Q^*\) is known or approximated, acting optimally means choosing an action with maximal \(Q^*(s,a)\) at the current state.

3. World Models

The motivation for a world model is simple: real interaction can be slow, expensive, unsafe, or unavailable. A learned model gives the agent a simulator it can query internally. In standard or generative world models, this simulator often tries to predict future observations, reconstruct sensory data, or model the next state in the same space where the agent observes the world.

The hat notation below is important. \(P\) and \(R\) are the true environment dynamics and reward. \(\hat{P}_{\theta}\) and \(\hat{R}_{\theta}\) are learned approximations with parameters \(\theta\). The agent does not own reality; it owns a model of reality, and that distinction becomes central when we discuss reward hacking.

A simple model-based RL formulation replaces the true MDP components with learned approximations:

This tuple has the same shape as the original MDP, but two pieces have been learned from data. The state space, action space, and discount are kept fixed, while the transition and reward components are replaced by neural or statistical estimators.

where:

The first approximation says the learned model should put probability mass on the same next states that the real environment would produce. The second says the learned reward predictor should score transitions similarly to the true reward function. If either approximation is poor in regions the policy visits, planning inside the learned model can become misleading.

For stochastic dynamics:

Stochastic dynamics preserve uncertainty. Instead of predicting one future, the model represents a distribution over possible futures. This is useful when the environment has noise, hidden variables, or inherently ambiguous outcomes.

For deterministic dynamics:

Deterministic dynamics compress the transition model into a function \(f_{\theta}\). Given the current state and action, the model outputs one predicted next state. This is easier to train and use, but it can hide uncertainty that matters for robust planning.

A common supervised objective for deterministic transition learning is:

The dynamics loss compares the observed next state \(s_{t+1}\) with the model's predicted next state \(f_{\theta}(s_t,a_t)\). Squared error encourages predictions to be close in the chosen state representation. This objective is straightforward when the state is already a meaningful vector, but it can be awkward for raw pixels, graphs, or curved geometric spaces.

With reward prediction:

The combined loss trains two abilities at once: predict what happens next and predict how good that transition is. The coefficient \(\lambda\) sets the tradeoff. A larger \(\lambda\) gives reward accuracy more influence; a smaller \(\lambda\) focuses training more heavily on state prediction.

Planning or policy learning inside the learned model becomes:

This is the model-based control loop. The policy is optimized using imagined trajectories sampled from \(\hat{P}_{\theta}\) and scored by \(\hat{R}_{\theta}\). The benefit is data efficiency: the agent can learn from simulated experience. The risk is model bias: the policy may discover behaviors that look excellent inside \(\hat{\mathcal{M}}_{\theta}\) but fail in the real environment.

4. Abstract World Models

The article's central point is that abstraction is not just compression for smaller storage. It is compression that tries to keep the variables needed for reasoning. In a bouncing-ball video, for example, millions of pixels may be reduced to position, velocity, gravity, contact events, and constraints. The model deliberately ignores details that do not affect the task, such as exact texture, while preserving structure that determines future behavior.

This is why abstract world models are often described as latent manifold models. The latent space \(\mathcal{Z}\) is not merely a bag of hidden features. It can be shaped as a vector space, graph, manifold, occupancy grid, simplicial complex, or another structured object, depending on what kinds of relationships the model must preserve.

The essential shift is from modeling raw state transitions directly:

This is the ordinary environment-level transition. It asks: given observed state \(s\) and action \(a\), what observed state \(s'\) comes next? For high-dimensional observations, this can force the model to track many details that are visually present but irrelevant to decision making.

to modeling transitions in latent space:

Now the model first encodes the observation into \(z\), then predicts the next latent state \(z'\). The transition \(\tau_{\theta}\) works in the abstract space, so prediction can focus on task-relevant variables rather than raw sensory reconstruction. The encoder \(\varphi_{\theta}\) is therefore not a passive compression function; it decides what the model will be able to reason about.

For an MDP defined by state space \(\mathcal{S}\), action space \(\mathcal{A}\), reward function \(R\), and transition function \(T\):

Here \(R\) maps a state-action pair to a scalar score, while \(T\) maps a state-action pair to the next state. This deterministic notation is a clean way to introduce the abstraction; a stochastic version would replace \(T\) with a transition distribution.

a self-supervised abstract world model can be defined in a latent space \(\mathcal{Z}\) with learnable components:

The encoder \(\varphi\) projects observed states into latent states. Its job is to create a representation where important similarities are close, important differences remain distinguishable, and irrelevant variation is suppressed.

The transition function \(\tau\) is the learned rule of motion in latent space. If \(z\) contains position and velocity, \(\tau\) should update those variables in a way that respects action effects and physical constraints.

The reward model \(r\) scores latent states and actions. It lets the agent plan without decoding every imagined latent state back into an observed state.

Here:

• \(\varphi\) is the encoder mapping observed states into latent states.
• \(\tau\) is the latent transition function.
• \(r\) is the latent reward or utility model.

A compact formulation is:

Read \(\mathcal{W}_{\theta}\) as the learned abstract world. It contains the space where reasoning happens, the encoder into that space, the transition model inside that space, the reward model used for evaluation, and the structural assumptions that keep the latent world coherent.

where \(\mathcal{G}\) denotes structural priors such as geometry, topology, group symmetries, invariances, equivariances, physical constraints, or causal structure. These priors are the difference between arbitrary compression and useful abstraction: they tell the model which transformations should preserve meaning and which changes should alter the predicted future.

Latent transition and reward prediction

The operational flow has three steps: encode the current observation, advance the latent state using the action, and estimate the reward or utility of that latent transition.

This line converts the observed state at time \(t\) into the model's internal state. If the observed state is an image, \(z_t\) might encode object positions, motion, scene layout, or affordances rather than individual pixels.

This line performs imagination in latent space. The model asks: if the current abstract state is \(z_t\) and the agent takes action \(a_t\), what abstract state should come next?

This line estimates the immediate value of the action from the latent state. In planning, this predicted reward is what lets the agent compare imagined futures without waiting for real-world feedback.

If the latent space \(\mathcal{Z}\) has a non-Euclidean geometry, the prediction loss should use an appropriate distance \(d_{\mathcal{Z}}\), such as a geodesic distance on a manifold:

The distance function \(d_{\mathcal{Z}}\) is doing important conceptual work. In a flat vector space, squared Euclidean distance may be fine. On a sphere, graph, mesh, or curved manifold, the meaningful distance may follow the structure of the space. A geodesic distance measures how far two latent states are along the manifold, not through empty ambient space.

where:

The target \(z_{t+1}\) is produced by encoding the actual next observation. Training therefore compares the predicted latent future \(\tau_{\theta}(z_t,a_t)\) with the latent encoding of what really happened.

The reward-prediction term can be written as:

This loss says that the latent reward model should match observed reward labels. When the reward is sparse or imperfect, this term can be fragile, because the model may learn the proxy signal rather than the intended task objective.

A general training objective can combine prediction, reward, geometric regularization, and anti-collapse terms:

This combined objective makes the design tradeoffs explicit. \(\mathcal{L}_{pred}\) teaches dynamics, \(\mathcal{L}_{reward}\) teaches scoring, \(\mathcal{L}_{geom}\) enforces structure such as smoothness or symmetry, and \(\mathcal{L}_{collapse}\) prevents the encoder from mapping many distinct observations to the same uninformative latent point.

Geometric priors: invariance and equivariance

Geometric priors express what should remain stable when the world is transformed. They help the model generalize from fewer examples because it does not have to relearn the same situation under every rotation, translation, viewpoint, or scale.

Invariance says the representation or learned function should not change under a transformation \(T\):

For example, if a classifier only needs to know that an object is a ball, rotating the image should not change the answer. The function's output is invariant to the transformation.

A regularized objective is:

The regularization term penalizes the model when transformed and untransformed versions produce different outputs. It is a soft constraint: the data loss still matters, but the model is nudged toward respecting the desired invariance.

Equivariance says a transformation in latent space should correspond to a predictable transformation in output space:

Equivariance is stronger and often more useful for world modeling. If an object moves two meters to the right in the observed world, its latent representation should move in a corresponding way rather than remain unchanged. The model preserves the structure of the transformation instead of ignoring it.

A regularized objective is:

This penalty encourages the learned function to commute with the transformation: transform first and then predict, or predict first and then transform, and the results should agree.

Interpretation of the latent space

A useful way to read \(\mathcal{Z}\) is through three complementary lenses. The geometry describes the shape of the representation, the variables describe what information lives there, and the activities describe what the agent does there.

5. JEPA-style Latent Prediction

JEPA fits naturally into abstract world modeling because it trains prediction where the abstraction lives. Instead of asking the model to reproduce every pixel of a target image or video patch, it asks the model to predict the target's embedding. This encourages the encoder to keep information that is stable and predictive while discarding details that are unpredictable or irrelevant.

The energy-based JEPA objective measures the discrepancy between a predicted target latent representation and the actual target latent representation:

The energy \(E\) is low when the predicted target representation \(Pred(s_x,z)\) is close to the actual target representation \(s_y\). It is high when the prediction misses. Training minimizes this energy, so the model learns representations in which future or missing latent states can be predicted from context.

where:

• \(x\) is the input or context data.
• \(y\) is the target data.
• \(z\) is a latent variable.
• \(s_x\) is the latent representation of the context.
• \(s_y\) is the latent representation of the target.
• \(Pred\) is the prediction model in latent space.

The target estimate is:

This line is the actual prediction step. The context representation \(s_x\), together with latent variable \(z\), is passed through the predictor to produce \(\hat{s}_y\), the model's estimate of the target embedding.

JEPA-style models are especially relevant to abstract world models because they treat latent space as the primary arena of learning and prediction. The model does not need to reconstruct all details of \(y\); it needs to predict the abstract features that matter for future inference, planning, or control.

A practical complication is representation collapse. If the encoder maps many inputs to nearly the same embedding, the prediction task becomes artificially easy but useless. This is why JEPA-style systems often need architectural or regularization mechanisms that keep the latent space informative and well spread out.

6. Reward Hacking

Reward hacking is easiest to understand as a gap between measurement and meaning. The reward function is the measurable training signal; the intended utility is what the designer actually wanted. When those two diverge, a capable optimizer can exploit the measurable signal while violating the broader goal.

Suppose the intended utility is:

\(U\) represents the real objective we care about. It may be difficult to observe, difficult to write down, or too broad to encode directly.

but the agent is trained on the specified reward:

\(R\) is the proxy objective used during training. It may be a hand-designed reward, a learned reward model, a heuristic metric, or a feedback signal from another system.

The learned policy optimizes:

The subscript in \(\pi_R^*\) is a reminder that the policy is optimal for the specified reward \(R\), not necessarily for the intended utility \(U\).

Reward hacking occurs when this policy has high specified reward but low intended utility:

This pair of inequalities is the failure mode in mathematical form. The training metric says the policy is doing well, while the real objective says it is doing poorly.

The core mismatch is:

The equation is short because the problem is conceptual rather than notational. If the reward is not the goal, optimizing the reward more effectively can amplify the mismatch.

Reward hacking in learned world models

In a learned world model, the problem can arise from an inaccurate reward model, transition model, or both:

The transition error means the imagined world does not match the real world. The reward error means the imagined scoring function does not match the intended objective. Either error can create artificial high-reward plans.

The agent may optimize imagined trajectories:

This objective is evaluated inside the learned model. The policy can become excellent at producing trajectories that \(\hat{P}_{\theta}\) and \(\hat{R}_{\theta}\) approve of, even if those trajectories are unrealistic or undesirable outside the model.

while those trajectories perform poorly in the real environment:

The mismatch now has two layers: the policy is optimizing a learned reward inside learned dynamics, but deployment is judged by real dynamics and intended utility.

Reward hacking in latent abstract world models

In an abstract world model, the agent may exploit the learned latent reward model:

or the learned latent transition model:

In latent space, reward hacking can be subtler. The exploit may not look like a strange raw observation; it may look like a latent state that receives high score because the abstraction omitted a variable that mattered for safety, feasibility, or human intent.

The latent-space optimization objective is:

This is the same planning idea as before, but all rollouts occur through \(\tau_{\theta}\) and all rewards come from \(r_{\theta}\). The quality of the abstraction now determines whether imagined success corresponds to real success.

Reward hacking occurs when high latent reward does not correspond to high real-world utility:

This line states the alignment requirement for abstraction. A latent reward model is useful only if scoring \(z=\varphi_{\theta}(s)\) preserves what matters about the original state \(s\).

In short:

7. Comparison and Unifying View

The concepts above can be read as successive layers. Reinforcement learning defines the decision problem. Bellman equations define how long-term value decomposes over time. World models learn a simulator for the decision problem. Abstract world models move that simulator into a structured latent space. JEPA-style prediction is one way to train such latent spaces. Reward hacking is the warning that none of these layers guarantee alignment by themselves.

Unified mathematical picture

The full agent can be seen as combining representation, prediction, reward estimation, and policy optimization:

First, the observation \(s_t\) is encoded into the latent state \(z_t\). This is the abstraction step.

Second, the model predicts how the latent state changes after action \(a_t\). This is the dynamics step.

Third, the model estimates the reward or utility associated with that latent transition. This is the evaluation step.

Finally, the policy chooses actions from the latent state. This is the control step. The policy parameters \(\psi\) may be trained using real rollouts, imagined rollouts, or both.

The policy objective in latent space is:

This objective says: choose a policy that obtains high predicted reward when futures are rolled out through the latent transition model. It is powerful because planning can happen in a smaller, structured space. It is risky because errors in \(\varphi_{\theta}\), \(\tau_{\theta}\), or \(r_{\theta}\) can be optimized by the policy.

with the important alignment condition:

The alignment condition says the latent reward should preserve the intended utility of the original state-action pair. If the encoder drops a variable that matters to \(U\), then \(r_{\theta}\) may score the abstract state incorrectly.

and the important dynamics condition:

The dynamics condition says that predicting after encoding should match encoding after the real transition. In words: the abstract model should commute with the real world dynamics. This is a concise test of whether the latent transition is faithful to the environment.

Thus, an abstract world model should preserve the task-relevant structure of the world, not merely compress observations: