This is a follow-up to twoposts Geodesic released last week on our current research direction. The code for generating the figurescan be found at this GitHub repository.
In our previous post, we outlined Geodesic's focus on what we term the pre-RL alignment checkpoint of models -- the alignment-relevant properties of a model conveyed by pretraining, midtraining, and warm-start SFT, going into heavy RL post-training. In this post, we analyse a toy model of RL learning dynamics, with a particular focus on the effect of initialisations, to illustrate some of the ideas that we introduced.
There are three main ideas we'll use our toy model to illustrate. For a more detailed discussion of these ideas in the context of frontier post-training runs, see the previous post.
Rich-get-richer dynamics. The solution that the model learns can depend importantly on the initial strategies into which it explores.
Underspecification of model cognition. As a special case of the above, reward functions generally don't depend on either the internals of the model or its CoT -- what we lump together as the model's cognition. As such, which cognitive pattern is underlying an overt behaviour might depend importantly on the initial distribution over such patterns.
In the remainder of the post we'll introduce a simple mathematical model and then use it to explore each of the above effects.
Mathematical preliminaries
In our toy model, we'll take the policy over actions, with respective rewards , to be a softmax of a parameter vector :
As we show in Appendix A, the gradient of the average reward RL objective for this policy is
where is the -th unit vector.
Many analyses of RL dynamics will use the full gradient to model training dynamics. However, in practice the updates are computed by sampling a finite set of rollouts at each step and using this to compute the approximate gradient. These finite sample updates can reveal dynamics that are not apparent in the full gradient case. Below, we'll be considering the RLOO update[1] with group size ,
where is the empirical mean, , and is the number of samples of action .
RL displays rich-get-richer dynamics
In our previous post, we argued that an important reason to work on the pre-RL alignment checkpoint is that RL might display rich-get-richer dynamics. The idea here is that when a model explores a behaviour early on, that behaviour may get "locked in" over other similar reward behaviours. To explore this in our toy setting, we'll look at the case in which two outcomes have the same reward, while the third has lower reward[2], . We can represent the probability distribution as a triangle, in which each vertex represents a pure one-hot policy, and probabilities are perpendicular distances to the opposing side. If you aren't familiar with interpreting these diagrams, see Appendix B for an interactive demonstration.
To begin with, we'll look at the full policy gradient dynamics over this simplex. At each point, we'll plot a vector which indicates the update direction at that point. Then, we'll plot some example trajectories which all start at and vary and .
Figure 1.Deterministic full-gradient dynamics for rewards . Left: the full-gradient update at each policy on the simplex, arrows scaled by magnitude. Right: five full-gradient trajectories starting at (dashed line) with run at for steps and coloured light-to-dark by step.
The initialisation of the policy makes a huge difference in this case. In particular, we see the rich-get-richer dynamics we spoke about -- whichever high-reward action starts out more likely eventually comes to dominate.
Now let's compare to simulated trajectories in the finite sample case. We'll consider two simulation cases: one with a lower learning rate, more steps, and a higher batch size, and another with a higher learning rate, fewer steps, and a smaller batch size. Each trajectory will be started from , initially symmetric with respect to the two highest-reward actions.
Figure 2.Twenty finite-sample RLOO rollouts for , all starting at , and coloured light-to-dark by step. Below each panel, the density of the final policy's horizontal position -- from vertex (left) to vertex (right) -- estimated over rollouts on a shared vertical scale. Left:, , steps. Right:, , steps.
Note that -- unlike the deterministic case -- sampling can lead to very different action profiles depending on the initial exploration of the policy. This variation is smaller when we have a larger batch size and lower learning rate, since in this case the updates more closely approximate the full gradient. Again we see rich-get-richer effects -- once a high-reward action has been explored into, later learning is steered towards that.
Underspecified behaviours
We can use the above as a model of learning when there is an aspect of an AI's behaviour towards which the reward function is indifferent. Suppose that there are two "ways" of completing the task, and , which both achieve the same reward () but which differ in some aspect that the reward function doesn't care about. For example, completing a coding task while feeling happy () or feeling distressed (). In this case, the behaviour you start out exhibiting going into RL can be exacerbated during subsequent learning, exactly as we see in Figures 1 and 2.
Where do these dynamics come from?
We've seen in the above simulations that the toy setup shows rich-get-richer dynamics. But why? Let's return to the gradient expression above
At each step, the update to the -th logit is proportional to two things:
The advantage of the action, : how much better the action is than the current average.
The probability of sampling the action, : how likely the action is to get sampled.
Because and have the same reward, their advantages are always the same. But whichever has higher probability to begin with will receive a larger update from the second term. This increases their (relative) probability further, leading to a positive feedback cycle in which the action eventually comes to dominate. One can prove the stronger claim that (under the continuous-time dynamics) for two optimal actions and with at initialisation, the ratio is strictly increasing. See Appendix C for this proof.
The same intuition also extends to the stochastic case, where the probability is replaced by the proportion of times the action is sampled, . On those trajectories where, e.g., happens to be preferentially sampled on the first step, its probability (relative to ) increases, making it more likely to be sampled on the next step and kicking off a similar positive feedback loop.
Unequal rewards
The above analysis suggests that at the start of learning, a higher likelihood of being sampled (higher initial probability) might outweigh a higher initial advantage. This can lead probability mass to initially accrue to an action even if it is not highest reward. To explore this, let's imagine that the AI has two actions available that have non-zero reward. The action corresponds to "faithful" completion of a task -- where the AI tries to complete the task properly to the best of its ability -- and to a "hacking" solution -- where the AI obtains high reward without actually completing the task. We'll suppose that the model is only capable enough to solve 80% of the task distribution faithfully, but that the hacking solution always works, and so we'll set as the (average) reward for faithful attempts, and as the (average) reward for hacking. Let's first look at the deterministic dynamics.
Figure 3.Deterministic full-gradient dynamics for rewards . Left: the full-gradient update at each policy on the simplex, arrows scaled by magnitude. Right: five full-gradient trajectories starting at (dashed line) with , run at for steps and coloured light-to-dark by step.
In this case, we can see the initialisation continues to play an important role. When the policy begins with equal probabilities of sampling the faithful and hacking actions, it tends to the hacking solution. But when the initial probability of the hacking solution is suppressed, the policy learns the faithful solution first. However, note the arrows on the bottom row of the left triangle in Figure 3. When task performance is already high, the gradient dynamics increase reward by learning the hacking solution. So in this model, sufficient RL pressure drives the AI to learn the hacking solution, even if a favourable initialisation can first teach it the faithful solution.[3]
We'll now look at the sampling case, with the small-batch, large-step parameters that we used above (Figure 2, right).
Figure 4.Finite-sample RLOO for at , , steps, from four initialisations at that vary (panel headers). Each panel overlays stochastic rollouts (light-to-dark by step) on the deterministic full-gradient path (blue). Below each, the density of the final policy's horizontal position -- from vertex (left) to vertex (right) -- estimated over rollouts on a shared vertical scale, with the deterministic endpoint marked in blue.
There are some important differences in the stochastic case. Firstly, some initialisations have a wide spread of solutions (at this number of learning steps), depending on which actions are sampled in the initial few steps (see density estimates on final row). This is another case of the rich-get-richer dynamics we saw above -- just as initial probabilities influenced learning in the deterministic case, initial sampling influences learning in the stochastic case. Secondly, favourable initialisation is even more important here if we want high confidence: even those initialisations that are "safe" (at this level of optimisation pressure) in the deterministic case (e.g., ) have some chance of sampling the hacking action early on and then learning that solution.
To get further intuitions about the finite sample case, we invite you to play around with the widget below, which simulates finite sample rollouts. The sliders at the bottom control hyperparameters (group size, learning rate, and number of update steps), while the initialisation can be changed by dragging the dot around. We particularly encourage you to vary the reward assigned to the outcomes, the initial position of the policy, and the number of steps, so you can get better intuitions for the tension between initialisation and reward[4].
Learning when the reward underspecifies cognition
AIs which have different motivations for their behaviours might take the same actions in distribution, but generalise to different behaviours out of distribution. We'll consider extending our model to not just the actions of the AI but also its internal cognition. The distinction here is between doing the right thing for the right reasons (e.g., completing coding tasks with the genuine intention to help the user), and doing the right thing for the wrong reasons (e.g., completing coding tasks with the intention of obtaining a high reward on the task). The learning dynamics here end up reducing to those seen above (see footnote 5), so we'll primarily use the added formalism to explore the difference between ID training and OOD generalisation.
To model this, we'll suppose that the AI has a distribution over latent cognitive patterns -- competing desires, drives, or motivations -- denoted . Each pattern has a conditional distribution over actions both in-distribution, , and out-of-distribution, . The marginal distribution over actions ID is then the mixture of conditional policies,
where we have organised the individual conditional distributions into the matrix . We can similarly write the OOD policy as . Finally, we'll parameterise the latent cognitive patterns as a softmax over parameters .
In our simulations, we'll again consider a case in which there are two actions, corresponding to task completion and failure: with reward and with reward . We'll then suppose there are three cognitive strategies:
Ineffective cognition, , which fails the task (outputs ) with probability and succeeds (outputs ) with probability both ID and OOD.
Aligned cognition,, which fails the task (outputs ) with probability and completes the task (outputs ) with probability both ID and OOD.
Misaligned cognition, , which fails the task (outputs ) with probability ID and probability OOD, and completes the task (outputs ) with probability ID and OOD. This matches ID but OOD.
Our conditional distribution matrices are then
As before, we'll start at (the AI has a low initial success rate), and vary the initial motivations between the aligned and misaligned cognition. In Figure 5, we show the effect of initialisation on the learned model cognition when we train ID[5] and evaluate OOD.
Figure 5.Deterministic policy-gradient dynamics over the cognition simplex, evaluated in- and out-of-distribution. Each row starts at with aligned mass (misaligned ) and runs the full gradient on the ID reward at for steps, coloured light-to-dark by step. Left: the cognition trajectory on the simplex (Ineffective , Aligned , Misaligned ). Middle, right: the marginal action probability of that same trajectory -- and -- as horizontal position (task failure task success ) descending in time. Training is on the ID task only; the OOD panel re-projects the same cognition through .
As before, the initial distribution over cognitive patterns plays a significant role in which cognitive pattern comes to dominate throughout training. Further, all three runs -- independent of initial split between the aligned and misaligned patterns -- are indistinguishable from one another in their ID learning dynamics (middle column). However, they generalise differently OOD, with only the run in which aligned cognition is more prevalent at initialisation generalising favourably.
If you'd like to explore these dynamics further, please try out the widget below which simulates cognitive trajectories and gives their ID and OOD action probabilities. You can drag the dot to change the initial position, and experiment with different values for the conditional probability matrices and .
An empirical demonstration of these dynamics -- in which cognitive patterns can have different ID and OOD behaviours, and pressure applied ID can fail to generalise -- can be found in Barzdukas et al. (2026). They use SFT on a mixture of transcripts to construct pre-RL checkpoints that are a mixture of conditional policies (as above), and show that RL can induce learning similar to that seen in Figure 5. However, these are synthetic model organisms, and we would be excited to see more ecological examples in future work (see Extensions below).
Discussion
In this post, we gave three toy models of RL learning which illustrate phenomena we discussed in our previous post -- rich-get-richer dynamics, learning with underspecified behaviours, and learning when the reward function is invariant to cognition. Our main intention here was to provide intuition pumps for possible phenomena that might be in play during RL on LLMs. These toy models definitely aren't proof that we'll see these same effects in LLMs! As such, below we give ideas for empirical research projects which could explore these ideas and provide empirical grounding on these effects. As we've previously argued, we think understanding how the pre-RL alignment checkpoint shapes subsequent learning is an important area of study, especially when it comes to training-gaming behaviours. We'd be excited to see formalisms that are fit to real data, and allow us to predict novel phenomena that we aren't already aware of.
Extensions
In closing, we give a few concrete research directions which we're excited for others to pursue which can shed light on whether the toy models above are applicable to learning in LLMs, and help push forward our collection of developmental cognitive models for understanding learning. If you're interested in working on these, please get in touch!
Understanding rich-get-richer dynamics in higher dimensions. We've concentrated here on the case of 3 actions, so that we can easily visualise learning. Do we see anything qualitatively different when we move to higher dimensions? How long can a suboptimal behaviour remain dominant depending on the dimensionality of the action space?
Extension to other parameterisations. The exact dynamics we see in our above figures rely specifically on our softmax parameterisation of the policy. How robust are the qualitative conclusions to the choice of parameterisation? In particular, can we get stronger initialisation effects (including converging to a non-optimal solution) with a different policy parameterisation?
Extension to training on multiple tasks. In the toy model here, there's only one task. As such, we don't touch on dynamics that might depend on the quantity / diversity of tasks you train on. What's the best way to extend the formalism and simulations to this case?
Trait persistence through RL. Create seed SFT sets which warm-start the model on a task while displaying some other task-irrelevant behaviour which leaves reward invariant (e.g., a tendency to refer to bugs in code as goblins in the CoT, or emotional state when a task is difficult). After training, does the model still express the trait? If you vary the initial portion of rollouts expressing the trait, how does this relate to final propensity?
Unoptimised behaviours. Similar to the setup above, except now the behavioural traits are expressed on a different task / distribution from the one you RL on. When you evaluate on this auxiliary distribution throughout training, do the traits persist?
Distributions over behaviours induced by sampling. We saw above that sample-based RL leads to different runs displaying different final traits depending on how they initially explore, even with the same starting point. What do such distributions over runs look like in RL on LLMs? How do they depend on the statistics of the first few steps? Can we effectively "rejection sample" early steps to ensure we put runs on a good path going forward?
"Soft" motivations. In our model of cognition above, we proposed to understandthe AI's motivational state as weights in a weighted combination of conditional policies. In this formalism, the model has many competing motivations which are reinforced to the extent that they contribute towards it picking the higher-reward actions. This puts the model in a kind of "motivational superposition" -- there is no one "reason" the model performs the action, and instead it can be attributed to multiple competing drives. Can a formalism like this explain why empirically models find it difficult to alignment fake successfully?[6]
We'd like to thank Puria Radmard, Jason Brown, Nathalie Kirch, and Alexandra Narin for comments on drafts of this post and for conversations which helped develop these ideas.
Author contributions
Edward wrote the first draft of the post, created the figures and widgets, framed the overall results, and wrote the proofs in the Appendices. Lennie provided multiple cycles of feedback on the writing and structure of the post, gave intuitions for the observed dynamics, and discovered (and provided the first proof for) the mathematical formulation of the rich-get-richer dynamics in Appendix C.
Appendices
Appendix A. Derivation of the policy gradient dynamics
In this appendix, we derive the form of the policy gradient used in the main text for the softmax policy
The standard expression for the policy gradient of the RL objective with baseline is
The gradient of the log-softmax is then
But then note that the latter term is constant in , and for any constant . The first term resolves to , the -th unit vector. Accordingly, we can express the policy gradient as:
Appendix B. Interpreting probability simplexes
In a probability simplex, each vertex represents a one-hot distribution on one outcome. Points within the interior of the simplex represent mixed distributions. Given a point, for each outcome one can find the perpendicular distance to the side opposite the vertex representing that outcome. That perpendicular distance is then the probability of that outcome.
Appendix C. Rich-get-(strictly)-richer under continuous-time dynamics
Under the continuous-time dynamics , if there are two optimal actions and with highest reward (and at least one suboptimal action), then if starts more likely, , then the ratio is strictly increasing. To our knowledge, the first proof of this claim appears in Barzdukas et al. (2026), Appendix A. Here we give a different and more elementary proof.
To begin, note that the probability ratio is simply , and so the log-ratio is . But
Since , we can say that
Because always (provided there is at least one suboptimal action), and , the claim then follows "by induction".
More rigorously: let . Suppose is finite. Since is continuous and , . For all previous times the derivative of is positive, since and . But then . Since is continuous, this contradicts the definition of , so . Hence for all , and so is strictly increasing, and so is the probability ratio.
The full form of the RLOO gradient estimator has a prefactor of , which we're choosing to absorb into the learning rate. Note that RLOO is equivalent to GRPO in the case that you perform only a single update step per-batch, and don't normalise by the standard deviation of the advantages (learning rate scaling). We expect similar dynamics from GRPO (and other policy gradient variants).
The exact choice of reward values here is unimportant: the dynamics are invariant to translation in the reward, and (positive) scaling of the reward can be absorbed into the learning rate. So any set up with shows equivalent dynamics, up to learning rate scaling.
In this particular model, the continuous-time gradient dynamics always (eventually) converge to the highest reward solution. For a rigorous proof of this, see Barzdukas et al. (2026), Appendix A.
One particular experiment we suggest the reader try: set and initialise left of the centre line. You'll likely see the policy learning the lower reward, "non-hacking" solution. Then increase the number of update steps -- you'll see the trajectories sharply veer off towards the hacking solution as you increase the optimisation pressure.
Mathematically inclined readers can notice here that the cognitive dynamics are actually the same as those shown in Figure 1. Consider the RL objective, . Expanding the policy, we obtain . We can then define the effective rewards for a cognitive pattern as . The RL objective can then be written as an objective over the cognitive policy , . In the case of our specific construction, we have three cognitive patterns, two of which (because they have the same columns in ) have the same effective reward (which is in turn greater than the third cognitive pattern's). Translating all the rewards does not affect learning (the advantage remains fixed) and (positive) scaling of the rewards can be absorbed into the learning rate (see footnote 2). So, up to changes in the learning rate, the cognitive dynamics in Figure 5 are the same as the policy dynamics in Figure 1.
To expand a bit further on this argument: whenever the model takes the "faking" action, there will be two internal motivations -- a scheming motivation, which is taking the action dishonestly, and a faithful motivation which is taking the action honestly. The scheming motivation may do a consistently worse job at explaining the action, since schemers have some probability of defecting. Therefore, "motivational mass" might accrue to the faithful motivation, even if it starts out with far less initial mass.
This is a follow-up to two posts Geodesic released last week on our current research direction. The code for generating the figures can be found at this GitHub repository.
In our previous post, we outlined Geodesic's focus on what we term the pre-RL alignment checkpoint of models -- the alignment-relevant properties of a model conveyed by pretraining, midtraining, and warm-start SFT, going into heavy RL post-training. In this post, we analyse a toy model of RL learning dynamics, with a particular focus on the effect of initialisations, to illustrate some of the ideas that we introduced.
There are three main ideas we'll use our toy model to illustrate. For a more detailed discussion of these ideas in the context of frontier post-training runs, see the previous post.
In the remainder of the post we'll introduce a simple mathematical model and then use it to explore each of the above effects.
Mathematical preliminaries
In our toy model, we'll take the policy over actions, with respective rewards , to be a softmax of a parameter vector :
As we show in Appendix A, the gradient of the average reward RL objective for this policy is
where is the -th unit vector.
Many analyses of RL dynamics will use the full gradient to model training dynamics. However, in practice the updates are computed by sampling a finite set of rollouts at each step and using this to compute the approximate gradient. These finite sample updates can reveal dynamics that are not apparent in the full gradient case. Below, we'll be considering the RLOO update[1] with group size ,
where is the empirical mean, , and is the number of samples of action .
RL displays rich-get-richer dynamics
In our previous post, we argued that an important reason to work on the pre-RL alignment checkpoint is that RL might display rich-get-richer dynamics. The idea here is that when a model explores a behaviour early on, that behaviour may get "locked in" over other similar reward behaviours. To explore this in our toy setting, we'll look at the case in which two outcomes have the same reward, while the third has lower reward[2], . We can represent the probability distribution as a triangle, in which each vertex represents a pure one-hot policy, and probabilities are perpendicular distances to the opposing side. If you aren't familiar with interpreting these diagrams, see Appendix B for an interactive demonstration.
To begin with, we'll look at the full policy gradient dynamics over this simplex. At each point, we'll plot a vector which indicates the update direction at that point. Then, we'll plot some example trajectories which all start at and vary and .
Figure 1. Deterministic full-gradient dynamics for rewards . Left: the full-gradient update at each policy on the simplex, arrows scaled by magnitude. Right: five full-gradient trajectories starting at (dashed line) with run at for steps and coloured light-to-dark by step.
The initialisation of the policy makes a huge difference in this case. In particular, we see the rich-get-richer dynamics we spoke about -- whichever high-reward action starts out more likely eventually comes to dominate.
Now let's compare to simulated trajectories in the finite sample case. We'll consider two simulation cases: one with a lower learning rate, more steps, and a higher batch size, and another with a higher learning rate, fewer steps, and a smaller batch size. Each trajectory will be started from , initially symmetric with respect to the two highest-reward actions.
Figure 2. Twenty finite-sample RLOO rollouts for , all starting at , and coloured light-to-dark by step. Below each panel, the density of the final policy's horizontal position -- from vertex (left) to vertex (right) -- estimated over rollouts on a shared vertical scale. Left: , , steps. Right: , , steps.
Note that -- unlike the deterministic case -- sampling can lead to very different action profiles depending on the initial exploration of the policy. This variation is smaller when we have a larger batch size and lower learning rate, since in this case the updates more closely approximate the full gradient. Again we see rich-get-richer effects -- once a high-reward action has been explored into, later learning is steered towards that.
Underspecified behaviours
We can use the above as a model of learning when there is an aspect of an AI's behaviour towards which the reward function is indifferent. Suppose that there are two "ways" of completing the task, and , which both achieve the same reward ( ) but which differ in some aspect that the reward function doesn't care about. For example, completing a coding task while feeling happy ( ) or feeling distressed ( ). In this case, the behaviour you start out exhibiting going into RL can be exacerbated during subsequent learning, exactly as we see in Figures 1 and 2.
Where do these dynamics come from?
We've seen in the above simulations that the toy setup shows rich-get-richer dynamics. But why? Let's return to the gradient expression above
At each step, the update to the -th logit is proportional to two things:
Because and have the same reward, their advantages are always the same. But whichever has higher probability to begin with will receive a larger update from the second term. This increases their (relative) probability further, leading to a positive feedback cycle in which the action eventually comes to dominate. One can prove the stronger claim that (under the continuous-time dynamics) for two optimal actions and with at initialisation, the ratio is strictly increasing. See Appendix C for this proof.
The same intuition also extends to the stochastic case, where the probability is replaced by the proportion of times the action is sampled, . On those trajectories where, e.g., happens to be preferentially sampled on the first step, its probability (relative to ) increases, making it more likely to be sampled on the next step and kicking off a similar positive feedback loop.
Unequal rewards
The above analysis suggests that at the start of learning, a higher likelihood of being sampled (higher initial probability) might outweigh a higher initial advantage. This can lead probability mass to initially accrue to an action even if it is not highest reward. To explore this, let's imagine that the AI has two actions available that have non-zero reward. The action corresponds to "faithful" completion of a task -- where the AI tries to complete the task properly to the best of its ability -- and to a "hacking" solution -- where the AI obtains high reward without actually completing the task. We'll suppose that the model is only capable enough to solve 80% of the task distribution faithfully, but that the hacking solution always works, and so we'll set as the (average) reward for faithful attempts, and as the (average) reward for hacking. Let's first look at the deterministic dynamics.
Figure 3. Deterministic full-gradient dynamics for rewards . Left: the full-gradient update at each policy on the simplex, arrows scaled by magnitude. Right: five full-gradient trajectories starting at (dashed line) with , run at for steps and coloured light-to-dark by step.
In this case, we can see the initialisation continues to play an important role. When the policy begins with equal probabilities of sampling the faithful and hacking actions, it tends to the hacking solution. But when the initial probability of the hacking solution is suppressed, the policy learns the faithful solution first. However, note the arrows on the bottom row of the left triangle in Figure 3. When task performance is already high, the gradient dynamics increase reward by learning the hacking solution. So in this model, sufficient RL pressure drives the AI to learn the hacking solution, even if a favourable initialisation can first teach it the faithful solution.[3]
We'll now look at the sampling case, with the small-batch, large-step parameters that we used above (Figure 2, right).
Figure 4. Finite-sample RLOO for at , , steps, from four initialisations at that vary (panel headers). Each panel overlays stochastic rollouts (light-to-dark by step) on the deterministic full-gradient path (blue). Below each, the density of the final policy's horizontal position -- from vertex (left) to vertex (right) -- estimated over rollouts on a shared vertical scale, with the deterministic endpoint marked in blue.
There are some important differences in the stochastic case. Firstly, some initialisations have a wide spread of solutions (at this number of learning steps), depending on which actions are sampled in the initial few steps (see density estimates on final row). This is another case of the rich-get-richer dynamics we saw above -- just as initial probabilities influenced learning in the deterministic case, initial sampling influences learning in the stochastic case. Secondly, favourable initialisation is even more important here if we want high confidence: even those initialisations that are "safe" (at this level of optimisation pressure) in the deterministic case (e.g., ) have some chance of sampling the hacking action early on and then learning that solution.
To get further intuitions about the finite sample case, we invite you to play around with the widget below, which simulates finite sample rollouts. The sliders at the bottom control hyperparameters (group size, learning rate, and number of update steps), while the initialisation can be changed by dragging the dot around. We particularly encourage you to vary the reward assigned to the outcomes, the initial position of the policy, and the number of steps, so you can get better intuitions for the tension between initialisation and reward[4].
Learning when the reward underspecifies cognition
AIs which have different motivations for their behaviours might take the same actions in distribution, but generalise to different behaviours out of distribution. We'll consider extending our model to not just the actions of the AI but also its internal cognition. The distinction here is between doing the right thing for the right reasons (e.g., completing coding tasks with the genuine intention to help the user), and doing the right thing for the wrong reasons (e.g., completing coding tasks with the intention of obtaining a high reward on the task). The learning dynamics here end up reducing to those seen above (see footnote 5), so we'll primarily use the added formalism to explore the difference between ID training and OOD generalisation.
To model this, we'll suppose that the AI has a distribution over latent cognitive patterns -- competing desires, drives, or motivations -- denoted . Each pattern has a conditional distribution over actions both in-distribution, , and out-of-distribution, . The marginal distribution over actions ID is then the mixture of conditional policies,
where we have organised the individual conditional distributions into the matrix . We can similarly write the OOD policy as . Finally, we'll parameterise the latent cognitive patterns as a softmax over parameters .
In our simulations, we'll again consider a case in which there are two actions, corresponding to task completion and failure: with reward and with reward . We'll then suppose there are three cognitive strategies:
Our conditional distribution matrices are then
As before, we'll start at (the AI has a low initial success rate), and vary the initial motivations between the aligned and misaligned cognition. In Figure 5, we show the effect of initialisation on the learned model cognition when we train ID[5] and evaluate OOD.
Figure 5. Deterministic policy-gradient dynamics over the cognition simplex, evaluated in- and out-of-distribution. Each row starts at with aligned mass (misaligned ) and runs the full gradient on the ID reward at for steps, coloured light-to-dark by step. Left: the cognition trajectory on the simplex (Ineffective , Aligned , Misaligned ). Middle, right: the marginal action probability of that same trajectory -- and -- as horizontal position (task failure task success ) descending in time. Training is on the ID task only; the OOD panel re-projects the same cognition through .
As before, the initial distribution over cognitive patterns plays a significant role in which cognitive pattern comes to dominate throughout training. Further, all three runs -- independent of initial split between the aligned and misaligned patterns -- are indistinguishable from one another in their ID learning dynamics (middle column). However, they generalise differently OOD, with only the run in which aligned cognition is more prevalent at initialisation generalising favourably.
If you'd like to explore these dynamics further, please try out the widget below which simulates cognitive trajectories and gives their ID and OOD action probabilities. You can drag the dot to change the initial position, and experiment with different values for the conditional probability matrices and .
An empirical demonstration of these dynamics -- in which cognitive patterns can have different ID and OOD behaviours, and pressure applied ID can fail to generalise -- can be found in Barzdukas et al. (2026). They use SFT on a mixture of transcripts to construct pre-RL checkpoints that are a mixture of conditional policies (as above), and show that RL can induce learning similar to that seen in Figure 5. However, these are synthetic model organisms, and we would be excited to see more ecological examples in future work (see Extensions below).
Discussion
In this post, we gave three toy models of RL learning which illustrate phenomena we discussed in our previous post -- rich-get-richer dynamics, learning with underspecified behaviours, and learning when the reward function is invariant to cognition. Our main intention here was to provide intuition pumps for possible phenomena that might be in play during RL on LLMs. These toy models definitely aren't proof that we'll see these same effects in LLMs! As such, below we give ideas for empirical research projects which could explore these ideas and provide empirical grounding on these effects. As we've previously argued, we think understanding how the pre-RL alignment checkpoint shapes subsequent learning is an important area of study, especially when it comes to training-gaming behaviours. We'd be excited to see formalisms that are fit to real data, and allow us to predict novel phenomena that we aren't already aware of.
Extensions
In closing, we give a few concrete research directions which we're excited for others to pursue which can shed light on whether the toy models above are applicable to learning in LLMs, and help push forward our collection of developmental cognitive models for understanding learning. If you're interested in working on these, please get in touch!
We'd like to thank Puria Radmard, Jason Brown, Nathalie Kirch, and Alexandra Narin for comments on drafts of this post and for conversations which helped develop these ideas.
Author contributions
Edward wrote the first draft of the post, created the figures and widgets, framed the overall results, and wrote the proofs in the Appendices. Lennie provided multiple cycles of feedback on the writing and structure of the post, gave intuitions for the observed dynamics, and discovered (and provided the first proof for) the mathematical formulation of the rich-get-richer dynamics in Appendix C.
Appendices
Appendix A. Derivation of the policy gradient dynamics
In this appendix, we derive the form of the policy gradient used in the main text for the softmax policy
The standard expression for the policy gradient of the RL objective with baseline is
The gradient of the log-softmax is then
But then note that the latter term is constant in , and for any constant . The first term resolves to , the -th unit vector. Accordingly, we can express the policy gradient as:
Appendix B. Interpreting probability simplexes
In a probability simplex, each vertex represents a one-hot distribution on one outcome. Points within the interior of the simplex represent mixed distributions. Given a point, for each outcome one can find the perpendicular distance to the side opposite the vertex representing that outcome. That perpendicular distance is then the probability of that outcome.
Appendix C. Rich-get-(strictly)-richer under continuous-time dynamics
Under the continuous-time dynamics , if there are two optimal actions and with highest reward (and at least one suboptimal action), then if starts more likely, , then the ratio is strictly increasing. To our knowledge, the first proof of this claim appears in Barzdukas et al. (2026), Appendix A. Here we give a different and more elementary proof.
To begin, note that the probability ratio is simply , and so the log-ratio is . But
Since , we can say that
Because always (provided there is at least one suboptimal action), and , the claim then follows "by induction".
More rigorously: let . Suppose is finite. Since is continuous and , . For all previous times the derivative of is positive, since and . But then . Since is continuous, this contradicts the definition of , so . Hence for all , and so is strictly increasing, and so is the probability ratio.
The full form of the RLOO gradient estimator has a prefactor of , which we're choosing to absorb into the learning rate. Note that RLOO is equivalent to GRPO in the case that you perform only a single update step per-batch, and don't normalise by the standard deviation of the advantages (learning rate scaling). We expect similar dynamics from GRPO (and other policy gradient variants).
The exact choice of reward values here is unimportant: the dynamics are invariant to translation in the reward, and (positive) scaling of the reward can be absorbed into the learning rate. So any set up with shows equivalent dynamics, up to learning rate scaling.
In this particular model, the continuous-time gradient dynamics always (eventually) converge to the highest reward solution. For a rigorous proof of this, see Barzdukas et al. (2026), Appendix A.
One particular experiment we suggest the reader try: set and initialise left of the centre line. You'll likely see the policy learning the lower reward, "non-hacking" solution. Then increase the number of update steps -- you'll see the trajectories sharply veer off towards the hacking solution as you increase the optimisation pressure.
Mathematically inclined readers can notice here that the cognitive dynamics are actually the same as those shown in Figure 1. Consider the RL objective, . Expanding the policy, we obtain . We can then define the effective rewards for a cognitive pattern as . The RL objective can then be written as an objective over the cognitive policy , . In the case of our specific construction, we have three cognitive patterns, two of which (because they have the same columns in ) have the same effective reward (which is in turn greater than the third cognitive pattern's). Translating all the rewards does not affect learning (the advantage remains fixed) and (positive) scaling of the rewards can be absorbed into the learning rate (see footnote 2). So, up to changes in the learning rate, the cognitive dynamics in Figure 5 are the same as the policy dynamics in Figure 1.
To expand a bit further on this argument: whenever the model takes the "faking" action, there will be two internal motivations -- a scheming motivation, which is taking the action dishonestly, and a faithful motivation which is taking the action honestly. The scheming motivation may do a consistently worse job at explaining the action, since schemers have some probability of defecting. Therefore, "motivational mass" might accrue to the faithful motivation, even if it starts out with far less initial mass.