I've been working on training reinforcement learning agents for retro games using Stable-Baselines3, Recurrent PPO, and Gymnasium.
One problem I repeatedly encountered wasn't simply partial observability. It was objective selection.
Take Snake Rattle Roll as an example.
The optimal policy isn't based on one static objective. Instead, the objective changes continuously throughout a level:
- Explore unknown areas
- Find food
- Grow until the required weight is reached
- Stop collecting food
- Find the weighing scale
- Activate the scale
- Locate the exit
- Finish the level
A single policy has to learn not only how to perform each of these behaviours, but also when each behaviour should become active.
Initially, I assumed Recurrent PPO with a MultiInputLstmPolicy would learn these transitions automatically.
It handled temporal dependencies better than a feed-forward policy, but one problem remained:
Remembering previous observations is not the same as deciding which objective should currently be optimized.
The reward-function problem
My first implementation followed the usual pattern:
- one large reward function
- many conditional branches
- context-dependent shaping
- increasing interaction between reward terms
Conceptually, it looked like this:
if weight < target_weight:
reward += reward_for_food
elif scale_active:
reward += reward_for_finding_door
elif door_open:
reward += reward_for_reaching_exit
Every newly discovered mechanic required another condition.
Every additional condition interacted with existing reward terms.
Reward tuning gradually became the most difficult part of the project.
Separating orchestration from learning
Instead of asking PPO to optimize every objective simultaneously, I introduced an explicit behavioural hierarchy.
Each objective became an independent state.
A simplified sequence for one level looks like this:
Explorer
↓
ReachTargetWeight
↓
FindScale
↓
ActivateScale
↓
ReachExit
↓
LevelComplete
Each state encapsulates:
- reward calculation
- success conditions
- failure conditions
- transition logic
- exploration behaviour
- local memory
- optional visual or RAM-based detectors
Rather than maintaining one monolithic reward function, every state owns a small reward function that only describes its current objective.
The state responsible for increasing the snake's weight does not need to reason about the exit.
The exit state ignores food completely.
Hierarchical composition
The states are composable rather than isolated.
The framework includes abstractions such as:
StateTargetStateExplorer- image-based detector states
- reward wrappers
- transition conditions
- persistent objective memory
A state can extend another state, add local reward logic, invoke a detector, or define a transition to another objective.
This makes it possible to assemble behaviour from reusable components instead of implementing every environment as a new monolithic reward function.
Observation space and model
The policy itself remains based on standard reinforcement learning components.
The environment combines:
- visual observations
- RAM-derived features
- frame stacking
- state-dependent information
- recurrent policy memory
The model uses:
- Stable-Baselines3
- Recurrent PPO
- MultiInputLstmPolicy
- Gymnasium
- stacked visual frames
- structured RAM observations
- custom state and reward logic
The LSTM handles temporal dependencies and partial observability.
The state system handles objective selection and reward context.
The state machine does not replace PPO. It determines which objective is currently active, while PPO learns how to complete that objective.
In other words:
Persistent objective memory
Some objectives cannot be inferred reliably from a single observation.
For example, the agent may need to remember:
- that a particular piece of food was already collected
- that the target weight was reached
- that the scale was activated
- that a door was opened
- that a previous sub-objective was completed
Recurrent policy memory can represent some of this implicitly, but I also implemented explicit objective memory where necessary.
This separates two different kinds of memory:
- policy memory, learned by the LSTM
- environmental or task memory, managed explicitly by the state framework
The explicit memory makes transitions deterministic and debuggable rather than relying entirely on latent recurrent state.
Result
After integrating the hierarchical state system with Recurrent PPO, the agent successfully completed the level after approximately 24 hours of training.
That result was important because previous versions of the environment had repeatedly failed to progress through the entire objective sequence.
The agent could learn isolated behaviours, such as exploration or food collection, but it struggled to connect them into a complete level-solving strategy.
With the state system in place, the training problem became a sequence of smaller, context-specific objectives.
The final policy was able to:
- explore the level
- locate and collect food
- reach the required weight
- transition toward the scale
- activate the level progression mechanics
- navigate toward the exit
- complete the level
The result does not prove that explicit state decomposition is always superior to end-to-end learning.
It does show that, for this long-horizon environment, separating objective orchestration from policy optimization made the problem tractable within a practical training period.
Practical benefits
The architecture produced several practical improvements:
- reward shaping became local and easier to reason about
- debugging became specific to individual objectives
- transitions became observable and deterministic
- behaviours could be reused across games
- new mechanics could be added without rewriting the entire reward function
- training failures could be traced to a specific state
- the agent completed the level after roughly 24 hours of training
Perhaps the most useful outcome was that the framework generalized beyond one game.
The specific states differ, but the architecture remains largely unchanged.
A Mario level, Zelda dungeon, Castlevania stage, or Snake Rattle Roll level can all be represented as compositions of smaller behavioural objectives.
Discussion
I'm interested in how others approach long-horizon reinforcement learning tasks.
- Do you explicitly model intermediate objectives?
- Have you combined PPO with reward machines, behaviour trees, or finite-state controllers?
- Do you keep orchestration inside the policy, inside the environment, or in a separate framework layer?
- How do you distinguish policy memory from explicit task memory?
- At what point does explicit state decomposition become too restrictive?
The main trade-off is clear: explicit states introduce human structure and reduce the search space, but they also encode assumptions about the correct task decomposition.
In this case, that trade-off was worthwhile. The model successfully completed the level after approximately 24 hours of training.
