Learning to explore with Thinker

AI · Reinforcement learning Thinker · NeurIPS 2023

Learning to explore with Thinker

Develop a reinforcement learning agent that learns to explore and clears the first level of the supplied simplified Montezuma’s Revenge environment. Build on Stephen Chung’s Thinker: Learning to Plan and Act (NeurIPS 2023) and the supplied extensions to learned planning and novelty memory. The broader goal is a general exploration method that can be applied to much harder exploration games.

Stay within the overall Thinker framework, with broad flexibility in how you extend it. You may change architectures, learning objectives, planning representations, exploration mechanisms and training procedures, and incorporate ideas from other approaches. Thinker must remain a substantive part of the resulting method; replacing it with a separate RL approach is outside this task.

Montezuma’s Revenge rewards long sequences of useful actions, with little feedback between discoveries. The challenge is to learn where to explore, how to return to promising states and how to use a learned world model when its predictions are still imperfect. The baseline already leaves the first room reliably, but has not cleared the first level.

Environment and task materials

The baseline ZIP contains a standalone JAX environment for fast, batched learning, adapted from JAXtari. It includes all 24 first-level rooms, game dynamics, collision maps, the complete training and evaluation code, pinned dependencies, reference results and the upstream license. No pretrained checkpoint or Atari ROM is required. Start with README.md in the ZIP.

This is a symbolic version of the game. The agent receives a 38 × 40 grid with 18 object channels and 42 context features, including room identity, inventory and movement information. The baseline pads the grid to 40 × 40 for its actor and pools it to 20 × 20 for its world model. Actions use the 18 Atari action IDs, with four simulation frames per real action.

The benchmark is the supplied environment and its first-level completion signal. A game starts from the normal reset and lasts at most 27,000 real actions, equivalent to 108,000 simulation frames. Leaving the first room is an intermediate milestone; completing the level requires reaching the environment’s terminal level exit. Keep the game dynamics, observations, external rewards, starting conditions and evaluation horizon fixed for comparable results.

Compute: Computational speed is not a main concern in this task. The priority is learning to explore and complete the first level with a task-general algorithm. A complete training and evaluation run taking a reasonable amount of time, for example around 24 hours, is acceptable. The baseline’s two-hour runtime is a reference, not a time limit or a requirement to match its speed. Report the hardware, training time and evaluation time used.

The reference run used one NVIDIA RTX 5090 and took about two hours, including final evaluation and independent action replay. Its recipe allows 6,000 seconds for the training loop; you may increase this budget. It uses Python 3.11 and pinned JAX, Flax and Optax versions, with a CUDA 13-compatible driver.

How the baseline works

Learn to plan and act. Before each real action, the actor takes ten imagined actions using a learned world model. It can reset an imagined trajectory to the current real state. A second learned model supplies value estimates, policy predictions and state embeddings. The actor uses the real observation and the resulting planning features to choose its real and imagined actions. Separate rollout statistics track external reward and curiosity, while an imagination reward encourages improvements in the best predicted rollout return. Planning uses learned predictions; it does not query the game simulator for hypothetical outcomes.

Retrieve planning experience from a small table. The supplied Thinker v2 extension replaces an explicit search-tree representation with a 20-slot memory of learned state embeddings, actions and temporal-difference statistics. A query retrieves the five closest entries by cosine similarity. The actor queries both the current real state and the current imagined state. Each retrieved row contains similarity, an action indicator, two accumulated temporal-difference quantities, a mean, a maximum and a visit count. A shared feed-forward network encodes each five-row table. This gives the actor a compact summary of relevant planning experience without maintaining a growing explicit tree.

Make novelty sensitive to revisits. The baseline pools the full observation grid to 10 × 10, appends the context vector scaled by 1/10, then uses a fixed random projection to produce a normalized 128-dimensional fingerprint. A cosine-similarity lookup matches this fingerprint to stored observations. New entries receive novelty bonus 1. Familiar entries receive

where age is the number of real collection steps since the last matching visit and visits is the previous visit count. Immediate repeated visits become less attractive, while returning after time away can become useful again. In the main training phase, groups of 32 environments share an 8,192-entry novelty memory. A learned prediction head supplies fingerprints for imagined observations; imagined queries do not change real visit counts. This novelty memory is separate from the planning table.

Learn from the agent’s own discoveries. Self-imitation replays the agent’s positive-return experience and reinforces actions whose recorded return exceeds the current value estimate. Training combines normalized PPO, learned-model updates and a staged schedule for novelty, self-imitation and life-loss boundaries. The exact schedule is in recipe.json and the supplement README. No demonstrations or scripted routes are used.

The planning-table and novelty designs are Stephen Chung’s unpublished extensions supplied with this task. They are part of the starting method, and are not fully described in the public Thinker repository. Read the bundled query.py, tree_v2.py, novelty.py, novelty_group.py and runner.py when extending them, and distinguish these starting ideas from your own contributions.

Reference performance and evaluation

| Measurement | Recorded RTX 5090 baseline | | --- | ---: | | Mean raw game reward | 2,107.8125 | | Maximum raw game reward | 2,500 | | First-room exits | 57 / 64 games | | First-level completions | 0 / 64 games | | Real training transitions | 27,494,400 | | Training, including checkpoint validation | 100.53 minutes | | Total, including final evaluation and CPU replay | 114.13 minutes |

These results use one training seed and two fresh evaluation batches of 32 games. Novelty memory starts empty and is shared within each evaluation batch. Report first-level completion rate as the main outcome, alongside raw reward, first-room exits, real training transitions, training time and evaluation time. Use the same evaluation protocol, disclose differences in training resources, and confirm promising gains across independent training seeds. When attributing a gain to a new mechanism, include a comparison that distinguishes its effect from simply training longer; matching the baseline’s runtime is not an acceptance requirement.

Run python run.py --out runs/seed1 --save-traces with one allocated GPU selected through CUDA_VISIBLE_DEVICES. This trains from scratch, evaluates the final policy and independently replays its action traces on CPU. The default evaluator’s success field refers to first-room exits; the replay reports’ level_completions field checks the actual task target. Use those replay reports for completion claims. Default runs without --save-traces keep only the final weights checkpoint and small result files; a supplied checkpoint is not needed to start.

Generality and research direction

You may use environment knowledge for diagnostics, but the final algorithm must be task-general. Inspecting the simulator, examining room layouts or testing a known situation can help explain a failure. Such diagnostic knowledge must not become a hand-coded route, room-specific policy, reward rule, privileged training target or demonstration for the final learner. The supplied symbolic observations are allowed inputs; learn how to use them instead of assigning special behavior to particular room IDs or coordinates. Planning must use a learned model rather than simulator lookahead, and final evaluation must begin from normal resets without state edits.

Within the overall Thinker framework, freely extend, revise or replace individual baseline components where the evidence supports it. Promising questions include how to learn useful novelty representations, preserve discoveries over long horizons, choose informative imagined trajectories, retrieve past planning experience and remain effective when model predictions are inaccurate. Explain how your method uses Thinker and why the resulting mechanism should remain useful in other hard-exploration games. Transfer experiments are welcome as the method matures.

A successful solution should demonstrate repeatable first-level completion by a trained, task-general agent using the Thinker framework, with reproducible code and independently replayable evaluation actions. Useful partial advances are also welcome when they reveal a new mechanism or substantially improve exploration. Explain what changed, why it helps and which experiments separate that contribution from the supplied baseline.

Task materials

Type

closed-ended

Solve status

open

Access

Invite-only