Chain of Thought Imitation with Procedure Cloning

1 Introduction

• Expert demonstrations from planning, search, or other multi-step algorithms contain intermediate computations, but imitation learning typically uses only state-action pairs and those procedures may rely on unavailable inference-time tools.The procedures are application-specific and lack a common representation.Source: §1 · p. 2
• Procedure cloning applies supervised sequence prediction to imitate the complete series of expert computations before producing an action, without relying on the expert procedure’s privileged tools at inference.It learns the procedure as well as the output action, offering a form of chain-of-thought supervision for autonomous agents.Source: §1 · p. 2
• PC achieves 100% test accuracy on unseen maze maps, while BC achieves 0% accuracy when the maze layout changes.The task uses BFS-derived expert paths and avoids additional test-map annotations.Source: §1 · pp. 2–3
• In image-based manipulation, PC reaches an 83.9% success metric versus 78.2% for the previous state-of-the-art, and it generalizes to stochastic, more difficult games.For games, PC learns from MCTS paths collected in a deterministic environment; MCTS performs poorly in those test settings even with simulator access.Source: §1 · p. 3

2 Related Work

• Prior generalization work often targets continual learning, whereas this paper studies transfer to separate test environments with configurations different from training.Existing approaches also use stochasticity, regularization, or data augmentation to reduce overfitting.Source: §2 · p. 3
• Unlike methods using extra offline data or auxiliary tasks, chain-of-thought imitation observes information from expert runs and learns procedures that directly influence action outcomes.The paper distinguishes procedure learning from treating procedure-like information as an auxiliary objective.Source: §2 · p. 3
• Prior chain-of-thought sequence modeling has focused on domain-specific problems, while this work formulates a domain-agnostic approach for sequential decision-making tasks.It aims to bridge memoryless policies in Markovian environments with sequence models for reasoning-based decision making.Source: §2 · pp. 3–4

3 Preliminaries

• An MDP specifies states, actions, rewards, transitions, initial-state distribution, and discount factor, with a policy selecting actions as the environment evolves.The state visitation distribution describes the states induced by a policy.Source: §3 · p. 4
• Standard imitation learning trains a policy on a fixed dataset of expert state-action pairs, aiming to approximate expert behavior beyond the training states.The evaluation metric is computed over states that may not appear in the dataset, so good test performance requires generalization.Source: §3 · p. 4
• Behavioral cloning frames imitation learning as supervised state-to-action mapping and minimizes the negative log-likelihood of expert actions on the training data.With small datasets relative to the state-action space, vanilla BC can fail to generalize, especially for image-based inputs and changed environment configurations.Source: §3 · p. 4

4 Procedure cloning

• Procedure cloning trains sequence prediction over intermediate computations, so the learner models both the final action and the procedure leading to it.Source: §4 · p. 4

4.1 Chain of thought imitation

• Scripts can expose causal relations, search trees the path to an optimal action, and multi-step algorithms task ordering; learning these procedures is chain-of-thought imitation.Source: §4.1 · pp. 4–5

4.2 Procedures and procedure observations

• A procedure is an ordered sequence of computations from input state to action; procedure observations capture its intermediate computation states, while the learned test policy takes only the state as input.Source: §4.2 · p. 5

4.3 Procedure cloning

• PC estimates p(a, x|s) by factorizing procedure and action predictions; its empirical likelihood can use a conditionally independent factorization when traces are fully observed and actions depend on the last computation state.Source: §4.3 · p. 5
• Unlike auxiliary BC, PC does not assume procedure observations and actions are conditionally independent given the state, instead treating the procedure as a precursor to the action.Source: §4.3 · p. 6

5 Proof of concept: Synthetic maze navigation

• The experiment uses generated gridworld mazes with four discrete navigation actions, training experts on training layouts and evaluating agents on disjoint test layouts.BFS generates the expert trajectories; success is measured by reaching the goal from random starting locations.Source: §5 · p. 6
• BFS procedure observations are snapshots of its visited array, recorded each time the search expands its perimeter.The array records visited positions, the action that reached them, and whether they have been backtracked.Source: §5 · p. 6
• On unseen 32 × 32 mazes, vanilla BC achieves 0% success, while procedure cloning generalizes; augmentation and auxiliary prediction help only on smaller mazes.BC's training-maze success is close to 100%, indicating that its test failure reflects generalization rather than model capacity.Source: §5 · p. 7

6 Experiments

• Procedure cloning generalizes across unseen maze layouts, manipulated-object positions, and game configurations with changed stochasticity or difficulty.The experiments cover simulated robotic navigation and manipulation, and MinAtar strategy games.Source: §6 · p. 7

6.2 Evaluating image-based robot manipulation

• With 1,000 expert demonstrations and early stopping, procedure cloning achieves 83.9% success versus 78.2% for state-of-the-art implicit BC.The bimanual task measures the proportion of particles in bowls; BC variants quickly overfit, while procedure cloning generalizes better.Source: Figure 5 · p. 8, §6.2 · p. 8

6.1 Evaluating continuous robot navigation in AntMaze

• On unseen AntMaze layouts, procedure cloning consistently improves success over vanilla BC; data augmentation and an auxiliary visited-array objective help less.Similar training-maze success for procedure cloning and auxiliary BC suggests the test gap is a generalization failure rather than insufficient model capacity.Source: §6.1 · p. 8

6.2 Evaluating image-based robot manipulation

• For bimanual sweeping, procedure cloning predicts the active geometric-center or bowl goal coordinates before predicting actions conditioned on those coordinates.Inference uses images and end-effector positions and orientations rather than the oracle's privileged object poses and contact points.Source: §6.2 · p. 8

6.3 Evaluating strategy games in MinAtar

• For MinAtar, procedure cloning uses the MCTS leaf-goal image and optimal action sequence as partial procedure observations, predicting the goal first and actions backward.The full MCTS search state is difficult to capture; the policy predicts the expert's output action last.Source: §6.3 · pp. 8–9, Figure 6 · p. 9
• Autoregressive procedure cloning generalizes best across MinAtar games and test settings, including sticky actions and increased game difficulty.Under sticky actions, MCTS struggles to search for good actions without extensive tuning.Source: §6.3 · p. 9, Figure 7 · p. 9

7 Conclusion

• Procedure cloning predicts intermediate computation outcomes step by step and generalizes better to differently configured test environments than improvements over behavioral cloning.At inference, it emulates those computations through autoregressive generation.Source: §7 · p. 9
• Procedure cloning adds computational overhead and requires practitioners to choose an encoding for expert procedures, potentially through trial and error in new domains.The conclusion identifies BFS, MCTS, and scripted robotic policies as examples of procedures encoded in the paper.Source: §7 · pp. 9–10

A.1 Maze generation details

• Generated 2D grid mazes contain random goals and internal walls forming tunnel-like maps, with at least one valid path from each empty location to the goal.Long corridors make learning a good policy with little data difficult, especially in larger mazes.Source: §A.1 · p. 16, Figure 8 · p. 16

A.2 Pseudocode for BFS and MCTS procedure data extraction

• BFS procedure data records the visited array, which tracks the action that visits each cell and supports backtracking from the start location.The pseudocode initializes a queue at the start, expands neighboring cells, and then backtracks.Source: §A.2 · p. 16
• MCTS procedure data records the search trajectory and rollout goal as intermediate observations for training procedure cloning.The pseudocode adds the trajectory during selection and the goal after simulation.Source: §A.2 · p. 16

A.3 Details of training the AlphaZero-style MCTS expert

• The MinAtar MCTS expert uses the environment simulator for rollouts rather than learning an environment model.Its implementation modifies an AlphaZero-style MCTS agent from Acme.Source: §A.3 · p. 16
• During MinAtar MCTS simulations, action selection is similar to UCT and uses a learned value, node visit counts, a search-greediness hyperparameter, and a guiding prior.The prior is updated by minimizing KL divergence to the softmax distribution of child visit counts.Source: §A.3 · p. 17

A.4 Details for training PC policies

• For MinAtar procedure data, the authors fix the predicted goal at 15 steps into the future and report that predicting multiple goals and backward action sequences improves procedure cloning performance.The model is trained on sequences of 15 actions.Source: §A.4 · p. 17
• The experiments use 10 expert trajectories per MinAtar game in the main-text results, while Freeway uses only the first 50 samples of one trajectory.During evaluation, PC stops at the environment timeout and outputs a uniformly random action if it times out.Source: §A.4 · p. 17

A.5 Architecture and hyperparameters

• Discrete maze, AntMaze, and MinAtar use five-layer CNNs and two-hidden-layer MLPs; larger tested model capacities did not improve generalization.The CNNs use channels 128, 128, 256, 256, and 256; the MLPs use 256 units per hidden layer.Source: §A.5 · p. 17
• MinAtar uses a four-layer, eight-head transformer, while bimanual sweep incorporates procedure observations as predictions preceding output actions.Data-augmentation baselines apply random crops, translations, and zooms.Source: §A.5 · p. 17
• Maze, AntMaze, and MinAtar models use Adam at learning rate 3e-4 for 500k steps with batch size 32, and are evaluated every 10k steps.Reported line plots average the last three evaluations; bimanual sweep follows the training procedure of [14] using TPU v2.Source: §A.5 · p. 18

B.1 Results on value iteration network

• The evaluation varies training-maze count across 5, 10, 20, and 40 mazes, maze size across 8 × 8, 16 × 16, and 32 × 32, and trajectories per maze across 1 and 16.The original VIN paper used 5,000 training mazes with seven trajectories each, substantially more data than this evaluation setting.Source: §B.1 · p. 18
• VIN does not generalize to the unseen tunnel-maze layouts in this evaluation, where its grid-structure reliance and the limited training data are identified as likely factors.The cited figure compares PC, BC variants, and VIN across 10 test maze layouts.Source: Figure 9 · p. 18, §B.1 · p. 18

• BC variants can achieve high training performance when the number of expert trajectories is large.Source: Figure 10 · p. 19
• BC variants still exhibit poor generalization performance on test mazes.Source: Figure 10 · p. 19
• Figure 10 reports average success rates from random start locations on training and testing discrete mazes.Source: Figure 10 · p. 19

B.3 Additional results on AntMaze

• Aux BC can achieve training performance similar to PC but exhibits poorer generalization on test mazes.Source: Figure 11 · p. 20

B.4 Additional results on MinAtar

• PC generally provides gains over BC across MinAtar test environments with training-matched configurations, sticky actions, and difficulty ramping.Source: Figure 12 · p. 21
• MinAtar results report average episode reward across three test environments and vary the number of expert trajectories by row.The rows use 50, 10, and 20 expert trajectories.Source: Figure 12 · p. 21