Chain of Thought Imitation with Procedure Cloning
Anonymous Author(s)
AffiliationAddress 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 fits 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 significant generalization to unseenenvironment configurations, including those configurations 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’simitation-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 adifference is that while humans imitate, we understand the task at a high level as opposed to onlyremembering 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 problemsthat are similar but different from the ones taught in lectures [4, 5, 6].
Just as students are taughtthe 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 imitationlearned 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-specific 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 final 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 overfits 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 difficult 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 configuration 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 overfitting 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 overfitting to the specific 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 final expert action duringtraining can help an agent generalize better.
Access to additional task information Many previous works in imitation learning and RL haveobserved that access to additional task information can help an agent learn better policies.
Forinstance, [36] relies on access to the simulator for collecting more expert trajectories to reducethe 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 offline 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 offline datafollowed by imitation learning from an expert.
These works commonly assume good coverage inthe offline 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 beneficial, [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 influence 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 beneficial.
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 specific 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 define 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 := ⟨ S, A, R , T , µ, γ ⟩ ,consisting of a state space S , an action space A , a reward function R : S ×
A → R , a transitionfunction T : S × A → ∆( S ) 1 , an initial state distribution µ ∈ ∆( S ) , and a discount factorγ ∈ [0 , 1) .
A policy π : S → ∆( A ) interacts with the environment starting at an initial state s [-
-]0 ∼ µ .
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 defined as d π ( s ) :=(1 − γ ) (cid:80) ∞ 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 fixed 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 π ∗ accordingto some metric such as the Kullback–Leibler (KL) divergence of their state visitation distributions
D KL ( d π ∗ ( s ) ∥ 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 classification 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 configurations (e.g., maze layouts and environmentstochasticity).
Procedure cloning
In this section, we first observe that in many imitation learning situations, expert demonstrationscan provide much richer insight into the desirable behavior than just the final 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 thefinal 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 agentpotentially has access to a rich set of information about a task which can facilitate learning.
Forinstance, when π ∗ is some scripting policy following a fixed set of rules [76], the rules (e.g., “ifclose to enemy then fire”) reveal causal information between firing and the enemy disappearing.
Inother settings when π ∗ is a search algorithm, the induced search tree exposes the path for finding theoptimal action.
When π ∗ is a multi-step hand-coded algorithm, the breakdown of the steps (e.g., “firstmove 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 final action as chain of thought imitation .
a [Figure] a s a x s a
[Figure]
Procedure cloning [Figure] s
Autoregressive
Conditional
[Figure]
4.2 Procedures and procedure observations
To formulate chain of thought imitation as a learning problem, we define a procedure Π : S → ∆( A )as a sequence of computations (Π 0 , Π 1 , ..., Π L , Π L +1 ) that first transform an input state s ∈ S intosome computation state (e.g., a variable value) using Π 0 : S → R d , followed by repeatedly applyingeach subprocedure, Π ℓ : R d → R d , ∀ ℓ ∈ [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-defined 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 define procedure observations as the sequence of computation states captured from a procedure:x =
( x 0 , ...
x L ) ∈ R ( L +1) × d
where x 0 = Π 0 ( s ) and x ℓ = Π ℓ ◦ Π l − 1 ◦ ... ◦ Π 1 ◦ Π 0 ( s ) , ∀ ℓ ∈ [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 defined 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 final 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 flexible,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 thefinal action only depends on the last computation state, a conditionally independent factorization canbe more desirable:
[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
[Equation]
[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-firstsearch (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-specific 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 {+
 ⊂+} 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 final visited array after backtracking uniquely identifies the expertaction on its own, we choose the conditionally independent factorization of p ( a, x | s ) (Equation 3)described in Section 4.3.
Specifically, we parametrize p θ ( x ℓ | x ℓ − 1 ) , ∀ ℓ ∈ [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 final 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 insufficient 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 significant generalization to previously unseen mazelayouts, positions of objects being manipulated, and environment configurations such as transitionstochasticity and game difficulty 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 offline 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 find 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 fixeddimensional 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 significant benefits 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 insufficient model capacity because the success rate on thetraining mazes between procedure cloning and auxiliary BC are similar (Figure 11 in Appendix B.3). [Figure] 6.2 Evaluating image-based robot manipulation Task description.
The bimanual sweeping task [83, 14] requires two 7-DoF robot arms equippedwith spatula-like end-effectors to sweep a pile of particles evenly into two bowls while avoidingdropping particles between the tips of the spatulas. The scripted oracle for collecting expert trajectoriesuses access to privileged information including object poses and contact points, which are notaccessible at test time. Rather, only high-resolution (96 × 96) images in conjunction with end-effectorpositions and orientations are given to the imitation-learned agent during inference. Actions areend-effector positions for each robot arm (6 dimensions for each arm for a total of 12 dimensionsfor the action); for how to incorporate kinematic actions into this environment, we refer the readerto [83].
Random seeds determining the initial position of the particles are partitioned so that test-timeconfigurations are not seen in training.
PC implementation and results.
The scripted policy first 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 tofirst 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 overfit 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 significantly 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 simplified 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 difficulty 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 difficult.
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 final 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 first predict the goal image using a CNN and then use
[Figure]
Procedure cloning Game envs
GPT
[Figure]
[Figure]
the training environments.
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 difficult game settings (availablein 3 out of 5 MinAtar games).
Autoregressive PC policy generalizes the best across all games and allsettings.
Conclusion
We have identified a major gap between the imitation learning problem setup and how training datafor imitation are often obtained in practice.
Specifically, 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 configurationsthan 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 justification 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 are
proprietary.• 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 reflect 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 to them?
[Yes]
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]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental 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 theywere chosen)?
[Yes] See Appendix A.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See all plots in Section 6.
(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.4.
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
using/curating?
[Yes]
e) Did you discuss whether the data you are using/curating contains personally identifiable
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, ifapplicable?
[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 amountspent on participant compensation? [N/A] A Experimental details
A.1 Maze generation details
We algorithmically generate maze layouts in a 2D gridworld environment. Each maze contains arandomly generated goal location and randomly generated internal walls that form a tunnel-likemap as shown in Figure 8. The long corridors of the maze layouts make it difficult to learn a goodnavigation policy with little data, especially when the maze is large. The maze layout generatorensures that there is at least one valid path from any empty location to the goal.
A.2 Pseudocode for BFS and MCTS procedure data extraction
Algorithm 1 and Algorithm 2 show the pseudocode of the BFS and MCTS procedures for computingthe expert action used in our experiments.
Procedure observations for training a PC agent is coloredin blue (i.e., the visited array for BFS and [ goal , traj ] list for MCTS).
Procedure observationsare added to the training data D Π as annotated.
Algorithm 2 shows the UCT [86] action selection as apart
Algorithm
Inputs start location s 0 , queue =
[ s 0 ]
Initialize visited ▷ Track the action that visits each cell.
while queue do ▷ Expand search.
s = queue.dequeue() for a, s ′ in actions, neighbors do if visited [ s ′ ] then continue visited [ s ′ ]
= a queue.enqueue( s ′ ) end for ▷ “visited” → D Π .
while not visited [ s 0 ] do ▷ Backtrack.
visited [parent[ s ]] = visited [ s ] ▷ “visited” → D Π .
s
= parent[ s ]
Return visited [ s 0 ]
Inputs root node n , prior p
Initialize traj =
[ n ] while node.children do a = UCT ( n ) ▷ Selection.
n = n. children [ a ] traj += [ n ] ▷ Add “traj” to D Π . n = n. expand ( p ) ▷ Expansion.
v, goal = rollout ( n ) ▷ Simulation.
▷ Add “goal” to D Π . while traj do n = traj .pop
() update( n , v )
Return argmax (n.children.
v )
A.3 Details of training the AlphaZero-style MCTS expert
For the expert MCTS policy on MinAtar, we modify the AlphaZero-style MCTS agent implementationfrom the Acme [87] library 2 .
We use the MinAtar environment simulator for rollouts (as opposed tolearning an environment model).
During each MCTS simulation run, an action is selected similar to
UCT [86] according to:
[Equation]
[Equation]
where Q ( s, a ) is a learned value function using TD-learning, N is the visit-count of the tree nodecorresponding to s , β is a hyperparameter for controlling how greedily to search, and π prior is a priorfor guiding the MCTS search and is updated by minimizing D KL ( π prior ( ·| s )
∥ π MCTS ( ·| s )) , whereπ MCTS ( ·| s ) is the softmax of the node s ’s children visit-counts. The hyperparameters for training the MCTS expert is shown in Table 1.
[Table]
A.4 Details for training PC policies
Predicting multiple goals on MinAtar.
Since we do not limit the search depth of the MCTSagent explicitly, the resulting search tree after 500 simulation runs can have arbitrarily long branches.
Therefore, for collecting procedure data, we set a fixed depth (i.e., 15) as the goal which representsthe state 15 steps into the future, and the GPT-like [88] autoregressive prediction model is trained onsequence of length 15 actions.
We further found that predicting more than a single goal along theoptimal path and more than a single backward action sequence during each time step improve PC’sperformance.
Number of expert samples in D Π .
The timeout for collecting expert trajectories and for evaluatingimitation-learned agents are set to 100, 1000, 2500 for discrete maze, AntMaze, and MinAtarrespectively.
Because some MinAtar games are easier to learn than others, i.e., Freeway only requires a single expert trajectory to learn a good BC policy, we take only the first 50 samples of a singletrajectory as expert data for Freeway to increase the challenge of imitation learning.
For the rest ofthe games, we take the full trajectories.
Results in the main text (Figure 7) use 10 expert trajectoriesfor each game.
Results for other number of expert trajectories can be found in Appendix B.4.
PC timeout.
During evaluation, PC needs to predict the stopping condition to output a final actionafter predicting a series of intermediate procedures.
This can result in the intermidiate procedurebeing arbitrarily long.
We therefore use the environment timeout steps as a hard stopping conditionfor PC, and output an action sampled uniformly at random if PC times out.
A.5 Architecture and hyperparameters
Architecture.
The CNNs used in discrete maze, AntMaze, and MinAtar of our experiments have 5layers of convolution with kernel size (3, 3, 3, 3, 3), stride (1, 1, 1, 1, 1), channel (128, 128, 256, 256,256), SAME padding, and without any pooling.
The MLPs in discrete maze, AntMaze, and MinAtarhave 2 hidden layers of 256 units each interleaved with ReLU activation.
We have experimented withlarger CNN and MLP model capacity, but they do not improve the generalization performance of theagents.
The transformer autoregressive model in the MinAtar experiment has 4 self-attention layersof 8 heads with 128 hidden units each, followed by 2 feedforward layers of 256 units interleavedwith ReLU activation.
The bimanual sweep experiment follows the same architecture as [14] – thisoriginal architecture already performs autoregressive modelling on actions, so it is straightforwardto incorporate procedure observations into this model, treating them as additional predictions thatappear before the output actions.
For all data augmentation baseline, we apply random crop of thesize of the image, random horizontal and vertical translation of [-10%, +10%], and random horizontaland vertical zoom of [-10%, +10%] sequentially.
Training.
For discrete maze, AntMaze, and MinAtar, we use the Adam optimizer with learning rate3e-4, train the models for 500k steps with batch size 32, and evaluate every 10k steps.
The line plotswe present show the average results of the the last 3 evaluations.
Each model is trained on a single
NVIDIA P100 GPU.
For bimanual sweep, we follow the same training procedure as [14] using TPUv2.
B Additional experimental results
B.1 Results on value iteration network
We were curious about whether having a planning-like (procedure-like) module in the policyparametrization similar to Value Iteration Network (VIN) [81] can solve the generalization taskin our maze setup.
We therefore evaluate VIN on 5, 10, 20, 40 training mazes of size 8 × 8, 16 × 16,and 32 × 32 with 1 and 16 trajectories each and test if VIN can generalize to the 10 unseen test mazes.
The original VIN paper used k
= 10 (the number of VI blocks) and k = 20 for maze size 8 × 8, 16 ×16, which we follow, and use k = 40 for maze 32 × 32.
Figure 9 shows that VIN does not generalizein our setting, likely due to VIN heavily relying on heavily on the grid structure of the task, where asour tunnel-shaped map is highly convoluted and requires a lot more training data (VIN was originallytrained on 5000 training mazes with 7 trajectories each, which is way more data than our evaluationsetting).
[Figure]
B.2 Additional results on discrete maze
[Figure]
B.3 Additional results on AntMaze
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
B.4 Additional results on MinAtar
[Figure]
[Figure]
Rows have different number of expert trajectories (50, 10, 20).
PC generally provides gains over BC in different evaluation settings
[Figure]
C Further analysis of procedure cloning
[Figure]
C.1 Compounding error of PC during evaluation
[Figure]
C.2 Visualization of learned PC procedure
[Figure]
C.3 Zero-shot generalization to greater distribution shift
[Figure] U D L
R
[Figure]
D Additional related workD.1 Neural program induction
While neural program induction (NPI) may be interpreted to apply to any “program” similar to PC,its demonstrated applications (both in [89] and [90]) are based on problems that exhibit hierarchicaland modular solutions with shared and repeated subprograms (e.g., a “pick” primitive within ablock-stacking task).
Thus, NPI not only advocates to imitate the full program but also parameterizethe agent in a modular way — like a program with function calls and return values — to takeadvantage of this modular and repeated structure.
The main argument in these existing papers hingeson the modular decomposition of both task and agent, which allows for more efficient data sharing,especially in multiple-task settings.
In contrast, our PC evaluations are on environments with muchless modular structures, and we avoid using specialized agent parameterizations in favor of moregeneric architectures (e.g., transformers).
Thus, we believe our work is complementary to NPI,showing that the paradigm of using procedural information – while first proposed by the NPI work tosome extent – applies to much more general settings than initially suggested by NPI.
D.2 Goal-conditional imitation learning.
While some of our experiments may be interpreted as utilizing “goals” as intermediate computation,there is a key difference from goal-conditioned imitation learning [91, 92].
Namely, goal-conditionedimitation learning advocates for using portions of the observation (or learned functions of theobservation) as goals, whereas PC uses information beyond what is available in the observation(e.g.
, coordinate positions of objects).
Thus, the common argument in goal-conditioned imitationcenters around getting more learning signal (goal-reaching) from the same data (s,a,s’ tuples).
Incontrast, PC shows that having richer data in the demonstrations (procedure information) is useful,even if that data is not available to the agent during inference .
We also note that another commonargument made in goal-conditioned imitation learning is attributing their sample efficiency benefits toreduced temporal frequency induced by the hierarchical design, whereas PC shows benefits withoutany change in temporal frequency.
Thus, we believe that PC presents a novel approach.
D.3 Learning-to-plan methods
Learning-to-plan methods (e.g., Universal Planning Networks [93], Value Iteration Networks [81])are distinct from PC in that they train policies end-to-end using state-action tuples.
No additionalsupervision of the intermediate computations of the expert is used; rather, these methods effectivelypropose a different policy parameterization leveraging inductive biases (i.e., with an embeddedend-to-end differentiable planner).
As shown in our evaluations of Implicit BC (Figure 5) and Value
Iteration Networks (Appendix B.1), flexible policy parametrizations and architectures with desiredinductive bias still fail to generalize without proper integration of procedure information duringtraining, showing that PC provides benefits orthogonal to these existing works.