Chain of Thought Imitation with Procedure Cloning
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Imitation learning aims to extract high-performance policies from logged demonstrations of expert behavior.
It is common to frame imitation learning as a supervised learning problem in which one ﬁts a function approximator to the input-outputmapping exhibited by the logged demonstrations (input observations to outputactions ).
While the framing of imitation learning as a supervised input-outputlearning problem allows for applicability in a wide variety of settings, it is also anoverly simplistic view of the problem in situations where the expert demonstrationsprovide much richer insight into expert behavior.
For example, applications such aspath navigation, robot manipulation, and strategy games acquire expert demonstrations via planning, search, or some other multi-step algorithm, revealing not just theoutput action to be imitated but also the procedure for how to determine this action.
While these intermediate computations may use tools not available to the agentduring inference (e.g., environment simulators), they are nevertheless informativeas a way to explain an expert’s mapping of state to actions.
To properly leverageexpert procedure information without relying on the privileged tools the expertmay have used to perform the procedure, we propose procedure cloning , whichapplies supervised sequence prediction to imitate the series of expert computations.
This way, procedure cloning learns not only what to do (i.e., the output action),but how and why to do it (i.e., the procedure).
Through empirical analysis onnavigation, simulated robotic manipulation, and game-playing environments, weshow that imitating the intermediate computations of an expert’s behavior enablesprocedure cloning to learn policies exhibiting signiﬁcant generalization to unseenenvironment conﬁgurations, including those conﬁgurations for which running theexpert’s procedure directly is infeasible.
Introduction
The idea of learning by imitation in autonomous agents closely resembles how humans (especiallychildren) learn in real life — by watching and mimicking how someone else performs a certain task [1].
While humans are able to generalize exceedingly well from a small number of demonstrations, today’s imitation-learned autonomous agents often struggle in situations that only slightly differ from thedemonstrations, e.g., opening a door in a different shape or color [2].
One explanation for such a difference is that while humans imitate, we understand the task at a high level as opposed to only remembering a mapping from images to actions [3], and indeed, generalization failures can also occurin humans when there is a lack of understanding of the underlying reasons for the behavior, such assolving a complex math problem.
As a result, students are told to “learn principles, not formulas” and“understand, do not memorize” to encourage better generalization — to be able to solve problems that are similar but different from the ones taught in lectures [4, 5, 6].
Just as students are taught the step-by-step derivations of a math problem during a lecture, is it possible to have an equivalent “chain of thought” supervision for training autonomous agents
While solving complex math derivations may not be the predominant application of today’s imitation- learned autonomous agents, the tasks they are commonly applied to — e.g., path navigation, robot
Training Datasets
[Figure]
manipulation, and strategy games — often do employ chain-of-thought reasoning procedures in theform of planning, search, or other multi-step algorithms when collecting expert data [7, 8, 9].
Sincesuch multi-step algorithms are application-speciﬁc and therefore lack a common representation whichcan be systematically characterized, it is typical for imitation learning to assume access to only loggeddemonstrations of state-action pairs, leaving out the much richer insight into expert behavior providedby the algorithm’s intermediate computations.
In addition, the planning or search procedures used bythe expert may rely on tools not available to the agent during inference (e.g., environment simulators),so
how these intermediate computations should be used to facilitate imitation learning is not clear.
In this work, we formulate chain of thought imitation as an extension of traditional imitation learning.
Different from the traditional imitation learning setup where an agent only has access to expertstate-action pairs as demonstrations, chain of thought imitation also has access to the intermediate computations that generated the expert state-action pairs in the training data (Figure 1).
Wethen propose procedure cloning (PC), an alternative to behavioral cloning (BC) [10], which appliessupervised sequence prediction to imitate the complete series of of expert computations beforeoutputting an expert action (Figure 1).
Procedure cloning learns a policy by maximizing the likelihood of the joint distribution of procedure observations and expert actions, which can be modeledautoregressively using a transformer-like architecture [11].
During inference, a procedure clonedpolicy autoregressively generates procedure observations from a given input state, mimicking thecomputations of a search or planning algorithm before outputting the ﬁnal action, thus avoidingany reliance on privileged tools or information used by the expert’s procedure.
From a modelingperspective, procedure cloning can employ a more expressive model (e.g., transformer) trained onmore data (i.e., procedure observations), which leads to better generalization according to the new“scaling law” of large language models [12].
Intuitively, procedure cloning learns not only what to do(i.e.
, the output action), but how and why to do it (i.e., the procedure), which further resembles howhumans learn complex tasks.
We demonstrate how to conduct procedure cloning and leverage the generalization ability of procedurecloned agents on a variety of path navigation, robotic manipulation, and strategy game tasks.
Forexample, in path navigation an expert trajectory may be determined by using a BFS search algorithmon an annotated map (e.g., x, y coordinates of obstacles), which are expensive to obtain [13].
Weshow that a procedure cloning agent can successfully learn to imitate BFS on a previously unseentest map without requiring additional annotations, achieving 100% test accuracy while a BC agentcompletely fails to navigate (0% accuracy) when the maze layout changes.
Similarly, in an imagebased robotic manipulation task, we observe that BC quickly overﬁts to the set of training images,whereas procedure cloning learns to predict the intermediate computation outcomes of a scriptingpolicy, thus generalizing much better and achieving a success metric of 83.9% compared to 78.2%from the previous state-of-the-art [14].
Finally, in strategy games expert trajectories are collectedby running MCTS [15], which requires access to a simulator and can be extremely slow [16].
Weshow that procedure cloning can effectively learn from the path traversed by MCTS collected in a deterministic environment to then successfully generalize in a zero-shot manner to stochasticenvironments and more difﬁcult game settings, where running MCTS, even if given access to asimulator, performs poorly.
Related Work
Generalization in sequential decision making.
Learning decision making policies with goodgeneralization properties has a long-standing history in bandit [17, 18, 19], imitation learning [20], andreinforcement learning (RL) [21, 22, 23, 24] settings in the form of regret or Probably Approximately
Correct (PAC) bound analysis.
These studies are concerned with generalization in a continuallearning setup without explicit separation of training and testing stages.
We are more interested inan agent’s ability to generalize to a separate “test” environment whose conﬁguration is differentfrom the training environment, which is often what happens when an agent is deployed to the realworld after being trained in a simulator.
With high-capacity neural network parametrized policiesbeing the norm in training image-based agents, the risk of overﬁtting to the training environmentis nontrivial [25, 26].
Various works have injected stochasticity into the training process includingusing stochastic policies [27], random starts [28, 29], sticky actions [30], and frame skipping [31]as a means to prevent overﬁtting to the speciﬁc environment dynamics experienced by the agentduring training.
A variety of image-based regularization techniques have also been applied to deep
RL including dropout and l -2 regularization [32], data augmentation and batch normalization [33],domain randomization [34], and network randomization [35].
Instead of improving generalizationthrough regularization or data augmentation like in existing work, we instead ask the question ofwhether learning the sequence of computations as opposed to only the ﬁnal expert action duringtraining can help an agent generalize better.
Access to additional task information Many previous works in imitation learning and RL have observed that access to additional task information can help an agent learn better policies.
For instance, [36] relies on access to the simulator for collecting more expert trajectories to reduce the quadratic error in task horizon to linear.
[37] assumes that expert demonstrations contain bothhigh-level and low-level trajectories, and learns a hierarchical policy explicitly.
While our procedureobservations can appear to be high-level labels or sub-goals similar to [37, 38], we do not makeassumptions about the structure of procedure observations, which can simply be scalar variablevalues computed during program execution.
There is also a large body of literature that assumesaccess to a suboptimal ofﬂine dataset in addition to the expert demonstration data collected fromthe same environment, and conduct representation learning [39, 40, 41, 42, 43], hierarchical skillextraction [44, 45, 46], or dynamics model learning [47, 48, 49] on the suboptimal ofﬂine datafollowed by imitation learning from an expert.
These works commonly assume good coverage inthe ofﬂine data, which requires large amounts of state-action pairs being collected from runningadditional policies in the environment.
We instead propose observing additional information fromonly running the expert policy, which requires many fewer state-action pairs being collected fortraining.
Another line of existing work in the direction of multi-task learning suggests that learning anauxiliary task enhances imitation learning or RL [50, 51, 2, 52, 2, 53, 54].
For instance, [55] showedthat predicting internal features of the emulator (e.g., enemy on the screen) is beneﬁcial, [56] foundpredicting scene depth helps navigation, and [57] found predicting explanations helps relational taskswith causal structure.
Chain of thought imitation differs from multi-task learning in that proceduresdirectly inﬂuence the action outcomes through computations, in which case learning the proceduredirectly (as opposed to treating it as an auxiliary objective similar to existing work) is more beneﬁcial.
The graphical model in Figure 2 visualizes this distinction.
Chain of thought sequence modeling.
The idea of decomposing multi-step problems into intermediate steps (the so-called chain of thought [58]) and learning the intermediate steps using a sequencemodel has been applied to domain speciﬁc problems such as program induction [59], learning tosolve math problems [60], learning to execute [61], learning to reason [62, 63, 64, 65, 66, 67], andlanguage model prompting [58].
The chain of thought imitation learning problem we formulateis domain agnostic and applicable to many sequential decision making task traditionally solvedby imitation learning in a Markovian setting such as robot locomotion, navigation, manipulation,and strategy games.
Unlike language-based tasks, problems in decision making have only recentlystarted being explored by language models [68, 69, 70, 71, 72, 73], as the Markovian nature of theseproblems brings the value of sequence modeling into question.
Our work contributes to bridging thegap between learning memoryless policies in Markovian environments and the intuition that largesequence models should help in reasoning-based decision making.
Preliminaries
In this section, we deﬁne relevant notations and introduce the problem of imitation learning.
We alsoreview behavioral cloning (BC), which our work aims to improve.
MDP notations.
Consider a Markov decision process (MDP) [74] M := h S, A, R , T , µ,  i , 139 consisting of a state space S , an action space A , a reward function R : S ⇥ A ! R , a transition 140 function
T : S ⇥ A !  ( S ) 1 , an initial state distribution µ 2  ( S ) , and a discount factor 141  2 [0 , 1) .
A policy ⇡ : S !  ( A ) interacts with the environment starting at an initial state s {+
+}
⇠ µ .
At each timestep t  0 , an action a t ⇠ ⇡ ( s t ) is sampled and applied to the environment, andthe environment transitions into the next state s t +1 ⇠ T ( s t , a t ) while producing a scalar reward
R ( s t , a t ) .
The state visitation distribution d ⇡ ( s ) induced by a policy ⇡ is deﬁned as d ⇡ ( s ) : =(1   ) P 1 t =
t · Pr [ s t = s | ⇡ , M ] .
We use ( s, a ) ⇠ d ⇡ to denote s ⇠ d ⇡ , a ⇠ ⇡ ( s ) .
Imitation learning.
In the standard imitation learning setting, R , T , and µ are unknown to the agent.
Learning solely takes place on a ﬁxed dataset of expert state-action pairs D ⇤
= { ( s ( i ) , a ( i ) ) } ni =1collected by an expert policy ⇡ ⇤ by sampling s ( i )
⇠ d ⇡ ⇤ and a ( i ) ⇠ ⇡ ⇤ ( s ( i ) ) .
The goal of imitationlearning is to train a policy ⇡ : S !  ( A ) on D ⇤ , such that ⇡ closely approximates ⇡ ⇤ according 150 to some metric such as the Kullback–Leibler (KL) divergence of their state visitation distributions
D KL ( d ⇡ ⇤ ( s ) k d ⇡ ( s )) .
Regardless of the choice of this metric, it is computed over states includingthose outside of the training data D ⇤ , and requires an agent to generalize over unseen states to achievegood test performance.
Behavioral cloning (BC).
The nature of learning from an expert dataset D ⇤ leads to the commonframing of imitation learning as a supervised problem of learning state to action mappings.
Under thisframing, behavioral cloning (BC) [10] proposes to solve the imitation learning problem by learning ⇡that minimizes
[Equation]
using samples from the training data ( s, a ) ⇠ D ⇤ .
Equation 1 can also be viewed as the crossentropy loss for multi-class classiﬁcation of state to action mappings.
Agents trained with thisvanilla BC objective (i.e., direct max-likelihood on the training dataset D ⇤ ) can quickly fail togeneralize to unseen states when the expert data size is small compared to the state-action space ofthe environment [36, 75], which is often the case when the states are based on images.
Image-basedinputs also pose the challenge of learning the right invariances for the agent to generalize fromtraining to testing environments with different conﬁgurations (e.g., maze layouts and environmentstochasticity).
Procedure cloning
In this section, we ﬁrst observe that in many imitation learning situations, expert demonstrationscan provide much richer insight into the desirable behavior than just the ﬁnal optimal action.
Weformulate learning under these situations as chain of thought imitation , where an agent can learn fromnot just the optimal action but the “thought process” an expert goes through before arriving at theﬁnal decision.
We then formalize such thought process as procedures , and propose procedure cloningfor learning such thought process through supervised sequence prediction.
4.1 Chain of thought imitation
Depending on the form of the expert policy ⇡ ⇤ used to generate the training data D ⇤ , an agent 175 potentially has access to a rich set of information about a task which can facilitate learning.
Forinstance, when ⇡ ⇤ is some scripting policy following a ﬁxed set of rules [76], the rules (e.g., “ifclose to enemy then ﬁre”) reveal causal information between ﬁring and the enemy disappearing.
Inother settings when ⇡ ⇤ is a search algorithm, the induced search tree exposes the path for ﬁnding theoptimal action.
When ⇡ ⇤ is a multi-step hand-coded algorithm, the breakdown of the steps (e.g., “ﬁrstmove to objects then sweep”) reveals important ordering information about the task.
Even in the casewhere ⇡ ⇤ is a human demonstrator, we can ask the human to explain the thought process that led totheir decision.
We refer to learning from the procedures (e.g., planning, search, multi-step algorithm)that generated the ﬁnal action as chain of thought imitation .
a x 0 s a s a x s a
Auxiliary BC
Procedure cloning Vanilla BC x L s
Autoregressive
Conditional independence
[Figure]
4.2 Procedures and procedure observations
To formulate chain of thought imitation as a learning problem, we deﬁne a procedure ⇧ : S !  ( A )as a sequence of computations ( ⇧ 0 , ⇧ 1 , ..., ⇧ L , ⇧ L +1 ) that ﬁrst transform an input state s 2 S intosome computation state (e.g., a variable value) using ⇧ 0 : S ! R d , followed by repeatedly applyingeach subprocedure, ⇧ ` : R d ! R d , 8 ` 2 [1 , L ] , to the computation state before mapping the lastcomputation state back to the MDP action space using ⇧ L +1 : R d !  ( A ) , in a sequential order.
Aprocedure can be broken down to subprocedures at any user-deﬁned granularity (e.g., a function ora for loop), depending on how frequently the computation state is retrieved for procedure learning.
Each pair of training data, ( s, a ) ⇠ D ⇤ , is acquired through executing such a procedure, i.e., ⇡ ⇤ =
⇧ .
We deﬁne procedure observations as the sequence of computation states captured from a procedure:x =
( x 0 , ...
x L ) 2 R ( L +1) ⇥
d where x 0 = ⇧ 0 ( s ) and x ` = ⇧ `  ⇧ l  1  ...  ⇧ 1  ⇧ 0 ( s ) , 8 ` 2 [1 , L ] .
The training data with procedure observations is denoted as D ⇧ = { ( s ( i ) , x ( i ) , a ( i ) ) } ni =1 .
Note thatwe do not assume any procedure or procedure observation is available at test time, hence we stillneed to learn a policy ⇡ that only takes s as input.
4.3 Procedure cloning
With the procedure observations deﬁned above, we are now ready for learning procedures throughprocedure cloning (PC).
We model a PC policy by estimating the joint distribution of the procedureobservations and the ﬁnal action conditioned on the input state p ( a, x | s ) , which we can factorizeautoregressively as:
[Equation]
Under this factorization, estimating p ( a, x | s ) reduces to estimating each conditional factor, which canbe parametrized using a transformer model [11].
The autoregressive factorization is highly ﬂexible,but if the amount of expert demonstrations is small, and each procedure observation x ` only dependson the previous procedure observation x `  1 (i.e., the computation states are fully observed), and theﬁnal action only depends on the last computation state, a conditionally independent factorization canbe more desirable: [Equation]
[Equation]
The graphical models of the two factorizations of PC policies are shown in Figure 2.
To learn a PCpolicy, we maximize the empirical likelihood of the joint distribution in Equation 2 on given samples
D ⇧ = { ( s ( i ) , x ( i ) , a ( i ) ) } ni =1 :
[Equation]
Connection to BC with auxiliary tasks.
The PC objective in Equation 5 can be reduced to thevanilla BC objective in Equation 1 by discarding the second and third term inside the expectationof Equation 5 and setting q ( a | x , s ) =
⇡ ( a | s ) .
We note that several previous works [51, 50, 2,52, 55, 56] have used procedure information as auxiliary tasks to BC, which may be interpreted aslearning p ( a, x | s ) under the assumption that a and x are independent conditioned on s : p ( a, x | s ) =
⇡ ( a | s ) · p ( x | s ) .
Procedure cloning instead focuses on situations where such a conditional independenceassumption does not hold (i.e., a is directly computed by the procedure represented by x ), and in
Action = R
State
U R
U R R R U R R R R
U R R R
R U
R R R
R R R
R
U R x 0 =  0 ( s )  5  6 xxxx 4 x 5 x
BFS procedure execution
Run BFS expand (keep track of actions to each cell).
Run BFS backtrack.
Datasets
State
X
Action = R
Procedure cloning p
( x 1 | x 0 ) p  ( x 2 | x 1 ) q  ( a | x 6 ) p  ( x 0 | s ) p  ( x 3 | x 2 ) p  ( x 4 | x 3 ) p  ( x 5 | x 4 ) p  ( x 6 | x 5 )        PC = these situations, as we will show in our experiments, treating the procedure information as a precursorto a can perform better than using it as an auxiliary task.
The graphical models of vanilla BC and BCwith auxiliary task objective are shown in Figure 2.
Proof of concept: Synthetic maze navigation
In this section, we study a tabular maze navigation task with synthetically generated maze layouts (see
Figure 3).
We describe the task setup, followed by how to extract procedure data from a breadth-ﬁrstsearch (BFS) path planning algorithm to train a procedure cloning agent, and empirically show thatprocedure cloning indeed generalizes much better to unseen maze layouts than other BC baselines.
While this proof of concept illustration might seem domain-speciﬁc to BFS in mazes, we will see in
Section 6.3 that procedure cloning can be applied to many types of search or multi-step algorithms.
Task description and evaluation protocol.
We use a gridworld maze environment in which anagent seeks to navigate to a goal location in a maze from a random starting location using 4 discreteactions including up ( U ), down ( D ), left ( L ), and right ( R ).
The input to the agent is a multi-channel“image” with the maze wall, goal location, and agent location encoded in separate channels.
Themaze layout is algorithmically generated with random internal walls that form a tunnel-shaped map(see example maze in Figure 4 and details in Appendix A.1).
We generate a set of mazes S [-
 ⊂-]0 ⇢ Sand split S 0 into disjoint training S trainand testing S test 0 sets.
We then generate expert trajectories byrunning BFS on only the training set of mazes S train 0 .
At test time, the agent is evaluated on S test 0 bytheir success rate of navigating to the goal.
Figure 3 visualizes the data and training pipeline.
Procedure data collection.
BFS is a common path planning algorithm for navigation [77, 78,79, 80], which we use to generate expert trajectories for training an imitation learning agent.
Tocompute the optimal action at each time step, BFS keeps track of a visited 2D array (colored cellsin Figure 3) that marks whether each position (1) has been visited by the search, (2) if so, whichaction visited it, and (3) has a position been backtracked.
We simply take a snapshot of the entirevisited array as procedure observations x ` every time BFS expands the search perimeter, resultingin a series of procedure data x = ( x 1 , ..., x L ) as shown in Figure 3.
⇧ 0 is the identity map andx
= s .
See pseudocode for collecting procedure observations in Appendix A.2.
Procedure learning.
Since each of the visited 2D arrays x ` only depends on the previousvisited array x `  1 , and the ﬁnal visited array after backtracking uniquely identiﬁes the expertaction on its own, we choose the conditionally independent factorization of p ( a, x | s ) (Equation 3)described in Section 4.3.
Speciﬁcally, we parametrize p ✓ ( x ` | x `  1 ) , 8 ` 2 [1 , L ] using a deep convolutional neural network that takes in the current visited array x `  1 as input and produces thenext visited array x ` as output.
We optimize ✓ using the cross-entropy loss between the predictedand true next visited array.
Since the procedure observations x ` and the original input s are in thesame image space,  ( x ) shares the same parameters as p ✓ ( x ` | x `  1 ) .
During inference when a newtest maze layout s is given, we apply ˆ x 0 =  ( s ) and ˆ x ` ⇠ p ✓ ( ·| ˆ x `  1 ) repeatedly until an array ˆ x Lis predicted for which the entry in ˆ x L corresponding to the agent’s current location is labelled as“backtracked”, and we return the backtracked action as the ﬁnal output action.
[Figure]
Generalization results.
We compare procedure cloning to applying data augmentation with randomcrop, translation, and zoom (Aug BC) or auxiliary objective of predicting the visited array (Aux
BC) to vanilla BC.
BC policies are parametrized with convolutional neural networks (CNN) andmulti-layer perceptrons (MLPs) (see hyperparameters in Appendix A.5).
Figure 4 (Disc.16 × 16and Disc.32 × 32) shows the average success rate (over 5 trajectories) of reaching the goal fromrandom start locations on 10 test maze layouts unseen during training.
Procedure cloning successfullygeneralizes to test mazes, whereas vanilla BC completely fails to learn (0% sucess rate) in biggermaze 32 × 32.
Aux BC and Aug BC help in the smaller maze but not in the bigger maze.
The poortest performance of BC is due to generalization failure as opposed to insufﬁcient model capacity as
BC’s success rate on the training mazes are close to 100% (see Figure 10 in Appendix B.2).
We alsoinclude additional experiments using VIN [81] as a baseline in Appendix B.1, showing that even anapproach specialized to the structure of gridworld tasks exhibits poor generalization.
Experiments
We now evaluate procedure cloning in larger scale settings on tasks including simulated roboticnavigation [82] and manipulation [83, 14], and learning to play MinAtar [84] (a miniature versionof Atari [85]).
Procedure cloning exhibits signiﬁcant generalization to previously unseen mazelayouts, positions of objects being manipulated, and environment conﬁgurations such as transitionstochasticity and game difﬁculty in each of the tasks, respectively.
See Appendix B for more results.
6.1 Evaluating continuous robot navigation in AntMaze
Task description.
Following the discrete maze navigation task in Section 5, we now consider amore realistic setting of mimicking real-world robotic navigation.
We adopt the AntMaze environmentfrom D4RL [82], where an 8-DoF “Ant” quadruped robot with continuous state and action spacesis placed in a 2D maze environment.
The original task designed for ofﬂine RL only has one mazelayout which is not revealed to the agent (the agent only sees its current position and joint-based statemeasurements).
To adopt the task for evaluating generalization across maze layouts, we also pass thealgorithmically generated maze layouts from Section 5 to the agent as inputs.
During evaluation, theagent needs to navigate to the goal in previously unseen maze layouts.
PC implementation.
The procedure that generated the expert training trajectories for goal-reachingin AntMaze involves a low-level PID controller that is good for navigating to local locations closeto the robot and a high-level waypoint generator which uses BFS search to ﬁnd the next localwaypoint [82].
We apply procedure cloning to the high-level waypoint generator following the samesteps as in discrete maze described in Section 5.
The predicted waypoint is then passed to a Gaussianparametrized policy (optimized with max-likelihood) together with the agent’s joint measurements.
For BC and variants, we use a convolutional neural network to embed the maze layout into a ﬁxeddimensional vector and concatenate the robot joint measurements together with the agent and goallocations to a Gaussian policy.
Generalization results.
Figure 4 (Ant 10 × 10 and Ant 13 × 13) shows the average (over 5trajectories) success rate of the robot reaching the goal from random starting locations on 10 testmazes.
Procedure cloning consistently provides signiﬁcant beneﬁts over vanilla BC, whereas dataaugmentation on the maze layout and predicting the visited array as an auxiliary loss are lesshelpful.
We know the poor test performance of BC (and other baseline methods) is more due togeneralization failure as opposed to insufﬁcient model capacity because the success rate on thetraining mazes between procedure cloning and auxiliary BC are similar (Figure 11 in Appendix B.3).
[Figure]
PC implementation and results.
The scripted policy ﬁrst scoops the particles by moving thespatulas to their computed geometric center and then moves the spatulas to a bowl before releasingthe particles.
At each state of a trajectory, we collect the Cartesian coordinates of one of these twogoal points – geometric center of the particles or position of the bowl – depending on the expert’sbehavior mode at that state, and use these as procedure observations.
We supervise the PC agent toﬁrst predict these coordinates, then predict the actions conditioned on the predicted coordinates andend-effector positions and orientations.
For both BC and PC we parameterize log-likelihood usingenergy-based models, also known as an implicit loss, which is state-of-the-art for this task [14].
Wecompare PC to BC trained directly on image inputs and auxiliary BC which learns to predict oraclecoordinates as an auxiliary objective.
We see that the BC variants quickly overﬁt to the training set ofobservations, whereas PC generalizes much better (Figure 5, left).
Despite early stopping (taking themaximum evaluation success rate over all training steps), PC still outperforms all of the BC
variants(Figure 5, right), improving signiﬁcantly over the state-of-the-art.
6.3 Evaluating strategy games in MinAtar
Task description.
MinAtar is a minature version of the Atari Arcade Learning Environment [85]consisting of 5 games with simpliﬁed 10 × 10 multi-channel images as inputs.
To generate the experttrajectories for training, we run a AlphaZero-style [15] Monte-Carlo tree search (MCTS) algorithmon the deterministic version of the environments to collect expert trajectories (see Appendix A.3 fordetails).
During evaluation, we test the imitation-learned agents on a set of test environments withdifferent seeds from the training environments where expert trajectories are collected.
To furtherevaluate generalization, we apply sticky actions with probability 0.1 and game difﬁculty ramping tothe test environments where such an option is available.
PC implementation.
In contrast to running BFS in maze navigation where the visited arraycleanly captures the entire state of the search, running MCTS in MinAtar is more convoluted,involving a number of MCTS simulation runs each with selection, expansion, rollouts, and backtracksteps and different tree structures, making capturing the full search state difﬁcult.
Fortunately,procedure cloning with autoregressive factorization (Equation 2) allows procedure data to be partialobservations of the computation state, and so we elect to use only a subset of the MCTS computationstates as supervision for PC.
Namely, we record the optimal action sequence after the last MCTSsimulation run (from the ﬁnal search tree) and the goal image at the tree leaf as procedure observations(highlighted in blue in Figure 6); this is effectively the optimal future trajectory determined by MCTS.
A PC policy is trained to use the input state to ﬁrst predict the goal image using a CNN and then use
State
MCTS root L L L
R R
R  0 ...  L MCTS Procedure execution [Figure]
(left
[Figure]
State
Goal
X =
[Goal, R, L]
Action (left)
Datasets p  ( x 0 | s )  p  ( x l | x < l , s )  q  ( a | x , s ) State Goal R L
Goal
R
L Action Procedure cloning Game envs GPT  PC = the goal as input to an autoregressive action sequence model p ( x ` | x < ` , s ) where x ` is the optimalaction that is ` -steps away from the goal (i.e., predicting the optimal action sequence backwards fromthe goal).
Figure 6 illustrates the data and training pipeline of procedure cloning.
Generalization results.
Figure 7 shows the average reward (over 50 episodes) collected by running
PC and BC policies in 3 test environments with different environment seeds than the environmentsused to collect training trajectories.
Figure 7 (left) evaluates generalization to stochastic environments,where we found sticky actions caused MCTS (gray) to struggle to search for good actions withoutextensive tuning.
Figure 7 (right) evaluates generalization to more difﬁcult game settings (availablein 3 out of 5 MinAtar games).
Autoregressive PC policy generalizes the best across all games and allsettings.
Conclusion
We have identiﬁed a major gap between the imitation learning problem setup and how training datafor imitation are often obtained in practice.
Speciﬁcally, applications that use planning, search, orother multi-step algorithms reveal the procedure for determining an expert action in addition tooutputting the expert action itself.
This led us to formulate chain of thought imitation learning, wherean agent is also given access to the intermediate computations that generated the expert state-actionpairs.
In order to generalize to test environments where intermediate computation is not available, wepropose procedure cloning, which learns to predict the intermediate computation outcomes step-bystep using a sequence model, and emulates the intermediate computations during inference throughautoregressive generation.
Evaluation on a variety of navigation, manipulation, and game play tasksshows that procedure cloning generalizes much better to test environments of different conﬁgurationsthan alternative improvements over behavioral cloning.
We hope the success of applying sequencemodeling to imitation learning achieved by procedure cloning opens up new research directions inthe intersection of large-scale sequence modeling and sequential decision making.
Limitations.
One major limitation of procedure cloning compared to traditional BC is in thecomputational overhead, since PC needs to predict intermediate procedures.
Furthermore, the choiceof how to encode the expert’s algorithm into a form amenable to PC is up to the practitioner.
Whilewe have presented ways to encode a variety of policies here (BFS, MCTS, scripted robotic policies),applying PC to other domains may require some amount of trial-and-error in designing the idealcomputation sequence for PC.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .•
Did you include the license to the code and datasets?
[No] The code and the data areproprietary.•
Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes](b) Did you describe the limitations of your work?
[Yes] See last paragraph in Section 6.3.(c)
Did you discuss any potential negative societal impacts of your work?
[N/A](d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A](b)
Did you include complete proofs of all theoretical results?
[N/A]3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See Appendix Aand supplementary material.(b)
Did you specify all the training details (e.g., data splits, hyperparameters, how they [Table] were chosen)?
[Yes] See Appendix A.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- [Table] ments multiple times)?
[Yes] See all plots in Section 6.
[Table] (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix A.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] See Section 6.
(b) Did you mention the license of the assets?
[Yes] The licenses are Apache License 2.
and GNU general public license.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
See supplementary material.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[Yes]
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A](b)
Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]