HRM-Agent: Training a recurrent reasoning model in dynamic environments using reinforcement learning

TL;DR

Problem: adapt a compact hierarchical reasoning model to dynamic, uncertain, and partially observable environments using reinforcement learning. Approach: train HRM-Agent by replacing the output head with a DQN and decoupling environment time from recurrent time to support iterative planning and reuse of prior computation. Findings: the agent reaches the goal in ≈99% of episodes and achieves mean path lengths near the optimal ~10 steps; Carry Z generally yields faster convergence and clearer reuse of earlier computation. Significance: this work shows recurrent inference can support dynamic planning under changing conditions, motivating future work toward on-policy training, re-enabling ACT, and scaling to more complex, partially observable domains.

Abstract

The Hierarchical Reasoning Model (HRM) has impressive reasoning abilities given its small size, but has only been applied to supervised, static, fully-observable problems. One of HRM's strengths is its ability to adapt its computational effort to the difficulty of the problem. However, in its current form it cannot integrate and reuse computation from previous time-steps if the problem is dynamic, uncertain or partially observable, or be applied where the correct action is undefined, characteristics of many real-world problems. This paper presents HRM-Agent, a variant of HRM trained using only reinforcement learning. We show that HRM can learn to navigate to goals in dynamic and uncertain maze environments. Recent work suggests that HRM's reasoning abilities stem from its recurrent inference process. We explore the dynamics of the recurrent inference process and find evidence that it is successfully reusing computation from earlier environment time-steps.

Figures (7)

• Figure 1: Maze-navigation tasks used to test the HRM-Agent model in the Nethack Learning Environment (NLE) with MiniHack. a) Four-rooms environment from sutton1999options and rauber2019hindsight with added doors (red). One of the four doors is closed at any given time. The closed door changes randomly during episodes. In each episode, the agent must navigate from a random start location to a random goal location. b) The random maze environment. This is a grid of fixed walls (white) and $N=10$ walls added randomly to some of the red squares before each episode. The random walls remain constant during the episode. In addition, a set of $M=5$ doors are enabled at the start of each episode. Each door randomly opens and closes independently during the episode. This creates a random corridor-type maze as shown in sub-figure (c), with viable paths changing as doors open and close. A* path-planning is used during maze generation to ensure that at least one viable path exists when all doors are open.
• Figure 2: Architecture of the HRM-Agent model. The output head has been replaced with a fully-connected two-layer DQN head, which produces $Q$-values (expected rewards) for each action. Another change is the initialization of the recurrent state $z_L$ and $z_H$ from their final, converged values in the previous environment time-step. However, no gradients are propagated across environment time-steps. Like the original HRM paper, gradients are not accumulated during the recurrent forward process.
• Figure 3: Validation of goal-navigation performance: Fraction of episodes in which the goal is reached. The HRM-Agent model is able to reach the goal in approximately 99% of all episodes in both environments. 5 models were trained with the recurrent state $z$ "carried forward" between each environment step. This allows the model to reuse and continue its current plan. 5 models were trained with recurrent state $z$ "reset" to a constant random vector at the start of each environment step. Each model training run is presented as a separate series.
• Figure 4: Validation of efficient path planning: Mean episode length. Since episodes end when the goal is reached, short episodes indicate efficient paths. The theoretical optimum average path length in both environments is approximately 10 steps, when considering the effects of the random doors. The HRM-Agent model achieves efficient paths in both environments. Each series depicts the performance of a separate model training run (there are insufficient runs for a distributional plot).
• Figure 5: Analysis of the convergence of the recurrent latent state $z$ to its final values in the four-rooms environment. The left plot shows $z_L$ and the right plot $z_H$. Each series is shows the median mean-square-error (MSE) between $z^i$ at recurrent time step $i$ and the final $z^{NT}$, drawn from $N$ environment time-steps collected from many validation episodes. Solid lines depict model variants trained and evaluated with the "Carry Z" initialization of $z^0$. Dashed lines depict the "Reset Z" variant. The peak at recurrent step 4 in the $z_H$ plot appears to be a consistent feature of the dynamics in this environment; note that the iteration of both H and L modules are unrolled into a single axis, and this step displays the result of the 2nd update of the H module.