S YNTHESIZING P ROGRAMMATIC P OLICIES THAT I NDUCTIVELY G ENERALIZE
Anonymous authors
Paper under double-blind review
A BSTRACT
Deep reinforcement learning has successfully solved a number of challenging control tasks.
However, learned policies typically have difﬁculty generalizing to novel environments.
We propose an algorithm for learning programmatic state machine policies that can capture repeating behaviors.
By doing so, they have the ability to generalize to instances requiring an arbitrary number of repetitions, a property we call inductive generalization .
However, state machine policies are hard to learn since they consist of a combination of continuous and discrete structure.
We propose a learning framework called adaptive teaching , which learns a state machine policy by imitating a teacher; in contrast to traditional imitation learning, our teacher adaptively updates itself based on the structure of the student.
We show how our algorithm can be used to learn policies that inductively generalize to novel environments, whereas traditional neural network policies fail to do so.
I NTRODUCTION
Existing deep reinforcement learning (RL) approaches have difﬁculty generalizing to novel environments (Packer et al., 2018).
More speciﬁcally, consider a task that requires performing a repeating behavior—we would like to be able to learn a policy that generalizes to instances requiring an arbitrary number of repetitions.
We refer to this property as inductive generalization .
In supervised learning, specialized neural network architectures have been proposed that exhibit inductive generalization on tasks such as list manipulation (Cai et al., 2017), but it is not obvious how those techniques would generalize to the control problems discussed in this paper.
Alternatively, algorithms have been proposed for learning programmatic policies that generalize well (Verma et al., 2019), but existing approaches have focused on simple stateless policies that cannot represent repeating behaviors.
We propose an algorithm for learning programmatic state machine policies .
Such a policy consists of a set of internal states, called modes , each of which is associated with a controller that is applied while in that mode.
The policy also includes transition predicates that describe how the mode is updated.
These policies are sufﬁciently expressive to capture tasks of interest—e.g., they can perform repeating tasks by cycling through some subset of modes during execution.
Additionally, state machine policies are strongly biased towards policies that inductively generalize, that deep RL policies lack.
In other words, this policy class is both realizable (i.e., it contains a “right” policy that solves the problem for all environments) and identiﬁable (i.e., we can learn the right policy from limited data).
However, state machine policies are challenging to learn because their discrete state transitions make it difﬁcult to use gradient-based optimization.
We use a standard solution where we “soften” the state transitions by making them probabilistic.
However, these techniques alone are insufﬁcient; they still run into local optima due to the the constraints on the structure of the policy function, as well as the relatively few parameters they possess.
To address this issue, we propose an approach called adaptive teaching , where we alternatingly learn a teacher , which is an overparameterized version of a student , which is a state machine policy trained, in turn, to mimic the teacher.
Because the teacher is overparameterized, it can more easily accomplish the task compared to the student (but does not generalize as well as the student).
Furthermore, the teacher is regularized to favor strategies similar to the ones taken by the student, to ensure the student can successfully mimic the teacher.
As the student improves, the teacher improves as well.
This alternating optimization can naturally be derived within the framework of variational inference, where the teacher encodes the variational distribution (Wainwright et al., 2008).
[Figure]
We implement our algorithm and evaluate it on a set of reinforcement learning problems focused on tasks that require inductive generalization.
We show that traditional deep RL approaches perform well on the original task, but fail to generalize inductively, whereas our state machine policies successfully generalize beyond the training distribution.
Our algorithm solves all of our benchmarks using state-machine policies with at most 4 modes.
We emphasize that we do not focus of problems that require large state-machines, which is a qualitatively different problem from ours and would require different algorithms to solve.
We believe that state-machines are most useful when only a few modes are required.
In particular, we are interested in problems where a relatively simple behavior must be repeated a certain number of times to solve the given task.
The key premise behind our approach, as shown by our evaluation, is that in these cases compact state-machines can represent policies that both have good performance and are generalizable.
When many modes are needed, then the number of possible transition structures grows exponentially, making it difﬁcult to learn the “true” structure without having an exponential amount of training data.
For these cases, we believe recurrent neural network policies may work better.
Example.
Consider the autonomous car in Figure 1, which consists of a blue car (the agent) parked between two stationary black cars.
The system state is ( x, y, θ, d ) , where ( x, y ) is the center of the car, θ is the orientation, and d is the distance between the two black cars.
The actions are ( v, ψ ) , where v is velocity and ψ is steering angle (we consider velocity control since the speed is low).
The dynamics are standard bicycle dynamics.
The goal is to drive out of the parked spot to an adjacent lane while avoiding collisions.
This task is easy when d is large (Figure 1a).
It is somewhat more involved when d is small, since it requires multiple maneuvers (Figures 1b and 1c).
However, it becomes challenging when d is very small (Figure 1d).
A standard RL algorithm will train a policy that performs well on the distances seen during training but does not generalize to smaller distances.
In contrast, our goal is to train an agent on scenarios (a), (b), and (c), that generalizes to scenario (d).
In Figure 1e, we show a state machine policy synthesized by our algorithm for this task.
We use d f and d b to denote the distances between the agent and the front and back black cars, respectively.
This policy has three different modes (besides a start mode m s and an end mode m e ).
Roughly speaking, this policy says (i) immediately shift from mode m s to m 1 , and drive the car forward and to the left, (ii) continue until close to the car in front; then, transition to mode m 2 , and drive the car backwards and to the right, (iii) continue until close to the car behind; then, transition back to mode m 1 , (iv) iterate between m 1 and m 2 until the car can safely exit the parking spot; then, transition to mode m 3 , and drive forward and to the right to make the car parallel to the lane.
This policy inductively generalizes since it captures the iterative behavior of driving forward and then backward until exiting the parking spot.
Thus, it successfully solves the scenario in Figure 1d.
Related work.
State machines have been used represent policies that have internal state (typically called memory ).
To learn these policies, gradient ascent methods assume a ﬁxed structure and optimize over real-valued parameters (Meuleau et al., 1999; Peshkin et al., 2001; Aberdeen & Baxter, 2002), whereas policy iteration methods uses dynamic programming to extend the structure (Hansen, 1998).
Our method combines both, but similarly to Poupart & Boutilier (2004), the structure space is bounded.
In addition, programmatic state machines use programs to represent state transitions and actions rules, and as a result can perform well while remaining small in size.
Hierarchies of Abstract
Machines (HAM)s also use programmatic state machines for hierarchical reinforcement learning, but assumed a ﬁxed, hand-designed structure (Parr & Russell, 1998; Andre & Russell, 2002).
Next, there has been growing interest in using program synthesis to aid machine learning (Lake et al., 2015; Ellis et al., 2015; 2018; Valkov et al., 2018;
Young et al., 2019).
Our work is most closely related to recent work using imitation learning to learn programmatic policies (Verma et al., 2018;
Bastani et al., 2018; Zhu et al., 2019; Verma et al., 2019).
However, these approaches assume a domain-speciﬁc program synthesizer that can learn programmatic policies given a supervised dataset of state-action pairs.
Building such a synthesizer for state machine policies is challenging since they contain both discrete and continuous parameters and internal state.
Thus, our student does not learn based on examples provided by the teacher, but is trained to mimic the teacher’s internal structure.
Our inductive generalization goal is related to that of meta-learning (Finn et al., 2017); however, whereas meta-learning trains on a few examples from the novel environment, our goal is to generalize without additional training.
Our work is also related to guided policy search, which uses a teacher in the form of a trajectory optimizer to train a neural network student (Levine & Koltun, 2013).
However, training programmitic policies is more challenging since the teacher must mirror the structure of the student.
Finally, it has recently been shown that overparameterization is essential in helping neural networks avoid local minima (Allen-Zhu et al., 2019).
Relaxing optimization problems by adding more parameters is a well established technique; in many cases, re-parameterization can make difﬁcult non-convex problems solve efﬁciently (Carlone & Calaﬁore, 2018).
P ROBLEM F ORMULATION
Dynamics.
We are interested in synthesizing control policies for deterministic, continuous-time dynamical systems with continuous state and action spaces.
In particular, we consider partially observable Markov decision processes (POMDP)
(cid:104)X , A , O , F, Z, X 0 , φ S , φ G (cid:105) with states X ⊆ R
d X , actions
A ⊆
R d A , observations O ⊆
R d O , deterministic dynamics F : X × A → X (i.e., ˙ x =
F ( x , a ) ), deterministic observation function Z : X → O , and initial state distribution x 0 ∼ X 0 .
We consider a safety speciﬁcation φ S : X → R and a goal speciﬁcation φ G : X → R .
Then, the agent aims to reach a goal state φ G ( x ) ≤ 0 while staying in safe states φ S ( x ) ≤ 0 .
A positive value for φ S ( x ) (resp., φ G ( x ) ) quantiﬁes the degree to which x is unsafe (resp., away from the goal).
Policies.
We consider policies π : O × S
→ A × S that keep internal memory; we assume the memory is initialized to a constant s 0 .
Given such a policy π , we sample a rollout (or trajectory) τ
= ( x 0 , x 1 , ..., x N ) with horizon N ∈ N by sampling x 0 ∼ X
and then performing a discrete-time simulation x n +
= x n + F ( x n , a n ) · ∆ , where ( a n , s n +1 )
=
π ( Z ( x n ) , s n ) and ∆ ∈ R > 0 is the time increment.
Since F , Z , and π are deterministic, τ is fully determined by x 0 and π ; τ can also be represented as a list of actions combined with the initial state i.e τ = (cid:104) x 0 , ( a 0 , a 1 , · · · , a N ) (cid:105) .
The degree to which φ S and φ G are satisﬁed along a trajectory is quantiﬁed by a reward function
R ( π, x 0 )
= − φ G ( x N ) + − (cid:80)
Nn =
φ S ( x n ) + , where x + = max (0 , x ) .
The optimal policy π ∗ in some class Π is one which maximizes the expected reward E x 0 ∼ X
[ R ( π, x 0 )] .
Inductive generalization.
Beyond optimizing reward, we want a policy that inductively generalizes to unseen environments.
Formally, we consider two initial state distributions: a training distribution
X train, and that includes initial states drawn from similar distribution seen during training whereas the test set a test distribution X test 0 that includes the extreme states never encountered during training.
P ROGRAMMATIC S TATE M ACHINE P OLICIES
To achieve inductive generalization, we aim to synthesize programmatic policies in the form of state machines.
At a high level, state machines can be thought of as compositions of much simpler policies, where the internal state of the state machines (called its mode) indicates which simple policy is currently being used.
Thus, state machines are capable of encoding complex nonlinear control tasks such as iteratively repeating a complex sequence of actions (e.g., the car example in Figure 1).
At the same time, state machines are substantially more structured than more typical policy classes such as neural networks and decision trees.
More precisely, a state machine π is a tuple (cid:104)M , H , G , m s , m e (cid:105) .
The modes m i ∈ M of π are the internal memory of the state machine.
Each mode m i ∈ M corresponds to an action function
H m i ∈ H , which is a function H m i : O → A mapping observations to actions.
When in mode m i , the agent takes action a
= H m i ( o ) .
Furthermore, each pair of modes ( m i , m j ) corresponds to a switching condition G m j m i ∈ G , which is a function G m j m i : O → R .
When an agent in mode m i observes o such that G m j m i ( o ) ≥ 0 , then the agent transitions from mode
m i to mode m j .
If there are multiple modes m j with non-negative switching weight G m j m i ( o ) ≥ 0 , then the agent transitions to the one that is greatest in magnitude; if there are several modes of equal weight, we take the ﬁrst one according to a ﬁxed ordering.
Finally, m s , m e ∈ M are the start and end modes, respectively; the state machine mode is initialized to m s , and the state machine terminates when it transitions to m e .
Formally, π ( o n , s n ) =
( a n , s n +1 ) , where a n = H s n ( o n ) , s
= m
s and s n
[Equation]
1 =
[Equation]
otherwise (1) where d arg max is a deterministic arg max that breaks ties as described above.
Action functions and switching conditions are speciﬁed by grammars that encode the space of possible functions as a space of programs.
Different grammars can be used for different problems.
Typical grammars for action functions include constants { C α : o
(cid:55)→ α } and proportional controls { P iα 0 ,α 1 : o (cid:55)→ α 0 ( o [ i ] − α 1 ) } .
A typical grammar for switching conditions is the grammar
[Equation]
of Boolean predicates over the current observation
o , where o
[ i ] is the i th component of o .
In all these grammars, α i ∈ R are parameters to be learned.
The grammar for switching conditions also has discrete parameters encoding the choice of expression.
For example, in Figure 1, the action functions are constants, and the switching conditions are inequalities over components of o .
F RAMEWORK FOR S YNTHESIZING P ROGRAMMATIC P OLICIES
We now describe our adaptive teaching framework for synthesizing state machine policies.
In this section, the teacher is abstractly represented as a collection of trajectories τ x 0 (i.e., an open-loop controller consisting of a ﬁxed sequence of actions) for each initial state x 0 .
A key insight is that we can parameterize τ x 0 in a way that mirrors the structure of the state machine student.
As we discuss in Section 4.2, we parameterize τ x 0 as a “loop-free” state machine.
Intuitively, our algorithm efﬁciently computes τ x 0 (from multiple initial states x 0 ) using gradient-based optimization, and then “glues” them together using maximum likelihood to construct a state machine policy.
4. A DAPTIVE T EACHING VIA V ARIATIONAL I NFERENCE
We derive the adaptive teaching formulation by reformulating the learning problem in the framework of probabilistic reinforcement learning, and also consider policies π that are probabilistic state machines (see Section 4.3).
Then, we use a variational approach to break the problem into the teacher and the student steps.
In this approach, the log-likelihood of a policy π is deﬁned as follows:
[Equation]
where p ( τ | π ) is the probability of sampling rollout τ when using policy π from a random initial state x 0 , λ ∈ R ≥ 0 is a hyperparameter, and R ( τ ) is the reward assigned to τ .
We have
[Equation]
where q ( τ ) is the variational distribution and the inequality follows from Jensen’s inequality.
Thus, we can optimize π by maximizing the lower bound Eq (3) on (cid:96) ( π ) .
Since the ﬁrst and third term of
Eq (3) are constant with respect to π , we have
[Equation]
Next, the optimal choice for q (i.e., to minimize the gap in the inequality in Eq (3)) is
[Equation]
where Z is a normalizing constant.
We choose q to have the form q ( τ ) =
p ( x 0 ) · δ ( τ − τ x 0 ) where δ is the Dirac delta function, p ( x 0 ) is the initial state distribution, and τ x 0 are the parameters to be optimized, where τ x 0 encodes a trajectory from x 0 .
Then, up to constants, the objective of Eq (5) equals
[Equation]
The ﬁrst term is constant; the second term is degenerate, but it is also constant.
Thus, we have
[Equation]
Thus, we can optimize Eq (2) by alternatingly optimizing Eq (4) and Eq (6).
We interpret these equations as adaptive teaching.
At a high level, the teacher (i
[Equation]
e., the variational distribution q ∗ in Eq (6)) is used to guide the optimization of the student (i
[Equation]
e., the state machine policy π ∗ in Eq (4)).
Rather than compute the teacher in closed form, we approximate it by sampling ﬁnitely many initial states x k 0 ∼ X 0 and then computing the optimal rollout from x k 0 .
Formally, on the i th iteration, the teacher and student are updated as follows:
[Equation]
where
[Equation]
The teacher objective Eq (7) is to both maximize the reward R ( τ ) from a random initial state x 0 and to maximize the probability p ( τ | π, x 0 ) of obtaining the rollout τ from initial state x
according to the current student π .
The latter encourages the teacher to match the structure of the student.
Furthermore, the teacher is itself updated at each step to account for the changing structure of the student.
The student objective Eq (3) is to imitate the distribution of rollouts according to the teacher.
Figure 2 shows the different components of our algorithm.
4. T EACHER : C OMPUTING L OOP -F REE P OLICIES
We begin by describing how the teacher solves the trajectory optimization problem Eq (7)—i.e., computing τ k for a given initial state x k 0 .
Parameterization.
One approach is to parameterize τ as an arbitrary action sequence ( a 0 , a 1 , ... ) and use gradient-based optimization to compute τ .
However, this approach can perform poorly—even though we regularize τ towards the student, it could exhibit behaviors that are hard for the student to capture.
Instead, we parameterize τ in a way that mirrors the student.
In particular, we parameterize τ like a state machine, but rather than having modes and switching conditions that adaptively determine the sequence of action functions to be executed and the duration of execution, the sequence of action functions is ﬁxed and each action function is executed for a ﬁxed duration.
More precisely, we represent τ as an loop-free policy τ = (cid:104)H , T (cid:105) .
To execute τ , each action function
H i ∈ H is applied for the corresponding duration T i ∈ T , after which H i +1 is applied.
The action functions are from the same grammar of action functions for the student.
The obvious way to represent a duration T i is as a number of time steps T i ∈ N .
However, with this choice, we cannot use continuous optimization to optimize T i .
Instead, we ﬁx the number of discretization steps P for which
H
i is executed, and vary the time increment ∆ i
= T i /P —i.e., x n +1 ≈ x n + F ( x n , H i ( o )) · ∆ i .
We enforce ∆ i ≤ ∆ max for a small ∆ max to ensure that the discrete-time approximation of the dynamics is sufﬁciently accurate.
Figure 3 shows examples of two loop-free policies and their corresponding trajectories for the example in Figure 1.
Optimization.
We use model-based trajectory optimization to compute loop-free policies.
The main challenge is handling the term p ( τ | π, x 0 ) in the objective.
Symbolically computing the this probability is hard because of the discrete-continuous structure of π .
Another alternative is to precompute the probabilities of all the trajectories τ that can be derived from π .
However, this is also infeasible because the number of trajectories is unbounded.
Thus, we perform trajectory optimization in two phases.
First, we use the a sampling-based optimization algorithm to obtain a set of good trajectories τ 1 , ..., τ L .
Then, we apply gradient-based optimization, replacing p ( · | π, x 0 ) with a term that regularizes τ to be close to { τ (cid:96) } L(cid:96)
=1 .
The ﬁrst phase proceeds as follows:
(i) sample
τ 1 , · · · , τ L using π from x 0 , and let p (cid:96) be the probability of τ (cid:96) according to π , (ii) sort these samples in decreasing order of objective p (cid:96) · e λR ( τ (cid:96) ) , and (iii) discard all but the top ρ samples.
This phase essentially performs one iteration of CEM (Mannor et al., 2003).
Then, in the second phase, we replace the probability expression with p ( τ | π, x 0 )
≈ (cid:80) ρ(cid:96)
=1 p (cid:96) ·
e − d ( τ,τ(cid:96) )
(cid:80) ρ(cid:96)
=1 p (cid:96) , which we use gradient-based optimization to optimize.
Here, d ( τ, τ (cid:96) ) is a distance metric between two loop-free policies, deﬁned as the L 2 distance between the parameters of τ and τ (cid:96) .
We chose the number of samples, ρ = 10 .
For our benchmarks, we did not notice any improvement in the number of student-teacher iterations by increasing ρ above 10.
So, we believe we are not losing any information from this approximation.
4. S TUDENT : L EARNING S TRUCTURED S TATE M ACHINE P OLICIES VIA I MITATION
Next, we describe how the student solves the maximum likelihood problem
Eq (8) to compute π ∗ .
Probabilistic state machines.
Although the output of our algorithm is a student policy that is a deterministic state machine, our algorithm internally relies on distributions over states induced by the student policy to guide the teacher.
Thus, we represent the student policy as a probabilistic state machine during learning.
To do so, we simply make the action functions H m j and switching conditions G m j 2 m j 1 probabilistic—instead of constant parameters in the grammar for action functions and switching conditions, now we have Gaussian distributions
N ( α, σ ) .
Then, when executing π , we obtain i.i.d.
samples of the parameters H (cid:48) m j ∼ H m j and { ( G m (cid:48) j m j ) (cid:48) ∼ G m (cid:48)
j m j } m (cid:48)
j every time we switch to mode m j , and act according to H (cid:48) m j and { ( G m (cid:48)
j m j ) (cid:48) } until the mode switches again.
By re-sampling these parameters on every mode switch, we avoid dependencies across different parts of a rollout or
[Figure]
Swimmer
[Figure]
T r a i n p e r f o r m a n c e
Ours
RL
RL-LSTM
Direct-opt
Car
Quad QuadPO PendulumCartpole Swimmer 0.
[Figure]
T e s t p
e r f o r m a n c e
[Figure]
RL
[Figure]
LSTM
[Figure]
different rollouts.
On the other hand, by not re-sampling these parameters within a mode switch, we ensure that the structure of π remains intact within a mode.
Optimization.
Each τ k can be decomposed into segments ( k, i ) where action function H k,i is executed for duration T k,i .
Furthermore, for the student π , let H m j be the action function distribution for mode m j
and G m j 2 m j 1 be the switching condition distribution for mode m j 1 to mode m j 2 .
Note that H m j and G m j
m j 1 are distributions whereas H k,i and T k,i are constants.
We have
[Equation]
For each ( k, i ) , let µ k,i be the latent random variable indicating the i th mode used by π starting from x k 0 ; in particular, µ k,i is a categorical random variable that takes values in the modes { m j } .
And µ k,i = m j means that H k,i ∼ H m j and T k,i is determined by the sampled switching conditions from distributions { G m (cid:48)
j m j } .
Assuming the latent variable µ k,i allows the student to compute π ∗ by computing H ∗ m j
and G m j 2 ∗
m j 1 separately.
Since the directly optimizing the maximum likelihood π is hard in the presence of the latent variables µ k,i , we use the standard expectation maximization (EM) approach to optimizing π , where the E-step computes the distributions p ( µ i,k
= m j ) assuming π is ﬁxed, and the M-step optimizes π assuming the probabilities p ( µ i,k
= m j ) are ﬁxed.
See Appendix A for details.
E XPERIMENTS
Benchmarks.
We use 6 control problems, each with different training and test distributions (summarized in Figure 8 in Appendix C): (i) Car, the benchmark in Figure 1, (ii) Quad, where the goal is to maneuver a 2D quadcopter through an obstacle course by controlling its vertical acceleration, where we vary the obstacle course length, (iii) QuadPO, a variant where the obstacles are unobserved but periodic (so the agent can perform well using a repeating motion), (iv) Pendulum, where we vary the pendulum mass, and (v) Cart-Pole, where we vary the time horizon and pole length.
(vi) Swimmer, where the goal is to move the swimmer forward through a viscous liquid, where we vary the length of the segments comprising the robot swimmer.
Baselines.
We compare against: (i) RL: PPO with a feedforward neural network policy, (ii) RLLSTM:
PPO with an LSTM, (iii) Direct-Opt: learning a state machine policy directly via numerical optimization.
Hyper-parameters are chosen to maximize performance on the training distribution.
Each algorithm is trained 5 times; we choose the one that performs best on the training distribution.
Results.
Figure 4 shows results on both training and test distributions.
We measure performance as the fraction of rollouts (out of 1000) that both satisfy the safety specﬁcation and reach the goal.
Inductive generalization.
For all benchmarks, our policy generalizes well on the test distribution.
In three cases, we generalize perfectly (all runs satisfy the metric).
For Quad and QuadPO, the policies result in collisions on some runs, but only towards the end of the obstacle course.
Comparison to RL.
The RL policies mostly achieve good training performance, but generalize poorly since they over-specialize to states seen during training.
The exceptions are Pendulum and
[Figure]
Swimmer.
Even in these cases, the RL policies take longer to reach the goals than our state machine policies (see Figure 10 and Figure 11 in Appendix C).
For QuadPO, the RL policy does not achieve a good training performance since the states are partially observed.
We may expect the LSTM policies to alleviate this issue.
However, the LSTM policies often perform poorly even on the training distribution and also generalize worse than the feedforward neural network policies.
We empirical analyze the policies.
Figure 5 shows the trajectory taken by (a) the RL policy, compared to (c) our policy, from a training initial state.
The RL policy does not exhibit a repeating behavior, which causes it to fail on the trajectory from a test state shown in (b).
Similarly, Figure 6 (Right) compares the actions taking by our policy to those taken by the RL policy on Quad and QuadPO.
Our policy produces smooth repeating actions, whereas the RL policy does not.
Comparison to direct-opt.
The state machine policies learned using direct-opt perform poorly even in training, illustrating the need to use adaptive teaching to learn state machine policies.
Varying the training distribution.
We study how test performance changes as we vary the training distribution on the Car benchmark.
We vary X trainas d ∼ [ d min , 13 . 5] , where d min = { 13 , 12 . 5 , 12 , 11 .
5 , 11 . 2 , 11 } , but ﬁx X test 0 to d ∼ [11 , 12] .
Figure 4 (right) shows how test performance varies with d min for both our policy and the RL policy.
Our policy inductively generalizes for a wide range of training distributions.
In contrast, the test performance of the RL policy initially increases as the train distribution gets bigger, but it eventually starts declining.
The reason is that its training performance actually starts to decline.
Thus, in some settings, our approach can outperform
RL policies even on the training distribution.
Interpretability.
An added beneﬁt of our state machine policies is interpretability.
In particular, we demonstrate the interpretability of our policies by showing how a user can modify a learned state machine policy.
Consider the policy from Figure 1e for the autonomous car.
We manually make the following changes: (i) increase the steering angle in H m 1 to its maximum value 0 .
5 , and (ii) decrease the gap maintained between the agent and the black cars by changing the switching condition G m 2 m 1 to d f ≤ 0 .
1 and G m 1 m 2 to d b ≤ 0 . 1 .
Figure 5 demonstrates these changes—it shows trajectories obtained using (c) the original policy, (d) the ﬁrst modiﬁed policy, and (e) the second modiﬁed policy.
There is no straightforward way to make these kinds of changes to a neural network policy.
Conclusion.
We have proposed an algorithm for learning state machine policies that inductively generalize to novel environments.
Our approach is based on a framework called adaptive teaching that alternatively learns a student that imitates a teacher and a teacher who adapts to the structure of the student.
We demonstrate that our policies inductively generalize better than RL policies.
A E XPECTATION M AXIMIZATION FOR S TUDENT O PTIMIZATION A. C OMPUTING p ( τ | π, x 0 )
First, note that we have
[Equation]
Similarly, the duration T k,i is determined both by the current mode
µ
i,k
= m j 1 , and by the switching conditions
G − m j 1 = { G m j 2 m j 1 } m j 2 from the current mode m j 1 into some other mode m j 2 .
More precisely, let γ k,i denote the trajectory on the ( k, i ) segment of τ k , and let ζ ( γ k,i , G − m j ) denote the earliest time at which a switching condition G ∈ G − m j becomes true along γ k,i .
Since G ∈ G − m j are distributions, ζ ( γ k,i , G − m j ) is a distribution on transition times.
Then, we have
[Equation]
In other words, T k,i is the duration until G m j 2 m j 1 triggers, conditioned on none of the conditions G m j 3 m jtriggering (where m j 3 (cid:54) = m j 2 ).
A. O PTIMIZING THE S TUDENT P OLICY
Numerically optimizing this maximum likelihood objective is hard because it requires integrating over all possible choices for the latent variables.
For example, if the teacher generates 10 loop-free policies every iteration and there are 10 modes in each loop-free policy, and 4 modes in the state-machine, the number of choices for the latent variables is 4 100 , which makes the enumeration infeasible.
The expectation-maximization method provides an efﬁcient way for computing the maximum likelihood, by alternatingly optimizing for the latent variables and the state-machine parameters.
The E-step computes the probability distributions p ( µ k,i = m j ) for a ﬁxed π , and the M-step optimizes H m j and G m j 2 m j 1 given p ( µ k,i = m j ) .
E-
step.
Assuming π is ﬁxed, we have
[Equation]
M-step.
Assuming p ( µ i,k
= m j ) is ﬁxed, we solve
[Equation] [Equation]
For G m j 2 m j 1 , the ﬁrst term handles the case µ k,i +1 =
m j 2 , where we maximize the probability that
G m j 2 m j 1 makes the transition at duration T k,i , and the second term handles the case µ k,i +1 (cid:54)
= m j 2 , where we maximize the probability that G m j 2 m j 1 does not make the transition until after duration T k,i .
We brieﬂy discuss how to solve these equations.
For action functions, suppose that H encodes the distribution N ( α H , σ 2 H ) over action function parameters.
Then, we have
[Equation]
Solving for the parameters of G m j 2 m j 1 is more challenging, since there can be multiple kinds of expressions in the grammar that are switching conditions, which correspond to discrete parameters, and we need to optimize over these discrete choices.
To do so, we perform a greedy search over these discrete choices (see Section A.3 for details on the greedy strategy).
For each choice considered during the greedy search, we encode Eq (11) as a numerical optimization and solve it to compute the corresponding means α ∗ G mj 2 mj 1 and standard deviations σ ∗ G mj 2 mj
Then, we choose the discrete choice that achieves the best objective value according to Eq (11).
Computing the optimal parameters for switching conditions is more expensive than doing so for action functions; thus, on each student iteration, we iteratively solve Eq (9) and Eq (10) multiple times, but only solve Eq (11) once.
The EM method does not guarantee global optima but usually works well in practice.
In addition, since computing the switching conditions is expensive, we had to restrict the number of EM iterations.
However, note that even if the EM algorithm didn’t converge, our overall algorithm can still recover by using additional teacher-student interactions.
The alternate method would be to run the EM algorithm multiple times/longer to get better results per student iteration, and “maybe” reduce the total number of teacher-student iterations.
We say “maybe” because the EM algorithm might have already converged to the global optima, making the extra EM iterations useless.
The tradeoff between our approach and this alternative depends on whether the teacher’s algorithm or the student’s algorithm is expensive for a particular benchmark.
However, from Figure 11, we can see that some of our benchmarks already use very few (< 5) teacher-student iterations (Car, QuadPO, Pendulum, Mountain car, and Swimmer).
Of the other three benchmarks that needed many iterations, for two of them (Cartpole and Acrobot), the student’s algorithm is as expensive as the teacher’s algorithm.
This justiﬁes our decision to not run the EM algorithm multiple times/longer.
A. S YNTHESIZING S WITCHING C ONDITIONS
Next, we describe how we search over the large number of discrete choices in the grammar for switching conditions.
It is not hard to show that in Eq (11), the objectives for the switching condition parameters G m j 2 m j 1 corresponding to different transitions ( m j 1 , m j 2 ) decompose into separate problems.
Therefore, we can perform the search for each transition ( m j 1 , m j 2 ) separately.
For each transition, the naïve approach would be to search over the possible derivations in the context-free grammar for switching conditions to some bounded depth.
However, this search space is exponential in the depth due to the productions B ::= B ∧ B and B ::=
B ∨ B .
Thus, we employ a greedy search strategy to avoid the exponential blowup.
Intuitively, our search strategy is to represent switching conditions as a kind of decision tree, and then perform a greedy algorithm to search over decision trees.
1 Our search strategy is similar to (but simpler than) the one in Bielik et al.
In particular, we can equivalently represent a switching condition as a decision tree, where the internal nodes have the form o
[ i ] ≤ α or o
[ i ] ≥ α (where i ∈ { 1 , ..., d O } and α ∈ R are parameters), and the leaf nodes are labeled with “Switch” or “Don’t switch”—e.g., Figure 7 shows two examples of switching conditions expressed as decision trees.
Then, our algorithm initializes the switching condition to a single leaf node—i.e.,
G cur ← “Switch”.
At each step, we consider switching conditions G ∈ next ( G ) that expand a single leaf node of G cur ; among these, we choose G cur to be the one that minimizes a loss cost ( G ) .
[Figure]
More precisely, to construct next ( G cur ) , we iterate over all leaf nodes L ∈ leaves ( G cur )
[Figure]
and all expressions E ∈ E , where
E
[Equation]
if o [ i ] ∼ α then “Switch” else “Don’t Switch”
[Table]
Here, ∼∈ {≥ , ≤} is a inequality relation, i ∈ { 1 , ...,
d O } is a component of o , and α ∈ R is a threshold.
For each pair L and E , we consider the decision tree G obtained by replacing L with E in
G cur .
The set next ( G cur ) contains all G constructed in this way.
Next, the loss function cost ( G ) is given by Eq (11).
In each iteration, our algorithm optimizes cost ( G ) over G ∈ next ( G cur ) , and updates G cur ← G .
To solve this optimization problem, we enumerate the possible choices ∼ and i and use numerical optimization to compute α (since α is a continuous parameter).
An example of a single iteration of our algorithm is shown in Figure 7.
In particular, letting G be the tree on the left and G (cid:48) be the tree on the right, the left-most leaf node of G is expanded to get G (cid:48) .
Our algorithm is summarized in Algorithm 1.
Overall, our algorithm searches over N · ( N − 1) ·
d O different discrete structures, where N is the number of nodes in the decision tree and d O is the length of the observation vector o .
B I MPLEMENTATION D ETAILS B. B ENCHMARKS AND S TATE -M ACHINES S TATISTICS
Figure 8 shows the statistics regarding the benchmarks such as the number of action variables and observation variables, and the set of initial states used for training and testing.
Figure 8 also shows the different aspects of the grammar used to describe the space of possible state-machine policies.
We are able to learn policies for these benchmarks using 2 to 4 distinct modes in the state-machine with either a constant or a proportional grammar for the action functions.
We use the Boolean tree grammar of depth 1 or 2 for all the switching conditions.
For the Quad benchmark, the action variable is the acceleration of the quadcopter in the vertical direction.
The observations include the position x, y , the velocities v x , v y , and the four sensors ox l , ox u , oy l , oy u to describe the obstacle course in the near neighborhood.
The QuadPO benchmark has the same action space as the Quad benchmark, but can only observe x, y, v x , and v y .
The synthesized state-machine policies for these benchmarks are shown in Figure 14 and Figure 15.
The action functions used for these benchmarks choose the acceleration to be proportional to v y .
The goal for the Pendulum benchmark is to control the force (continuous) at the actuated link in order to invert the link.
The observation space includes the full state, i.e., the angle θ and the angular velocity ω .
Figure 16 shows the synthesized state-machine policy for the pendulum benchmark.
The Cartpole benchmark consists of a pole attached to a cart.
The goal is to keep the pole upright by applying a continuous force to move the cart to the right or to the left.
The observations include the position x , the velocity v of the cart, the angle θ , and the angular velocity ω of the pole.
The synthesized solution is shown in Figure 17.
The Acrobot benchmark is similar to the Pendulum benchmark but with two links; only the top link can be actuated, and the goal is to drive the bottom link above a certain height.
The observations are the angles θ 1 , θ 2 and the angular velocities ω 1 , ω 2 of the two links.
For this benchmark, we vary the mass of the links between the training and the test distributions.
The synthesized solution is shown in
Figure 18.
For the Mountain car benchmark, the goal is to drive a low powered car to the top of a hill.
An agent has to drive back and forth to gain enough momentum to be able to cross the hill.
The agent controls the force (continuous) to move the car to the right or left and observes the position x and the velocity v at every timestep.
We vary the power of the car between the training and the test distributions.
The synthesized solution is shown in Figure 19.
The Swimmer benchmark is based on the Mujoco’s swimmer.
To make this benchmark more challenging, we use 4 segments instead of 3.
There are three actions that control the torques at the joints and the goal is to make the swimmer move forward through a viscous liquid.
The agent can observe the swimmer’s global angle θ , the joint angles ( θ 1 , θ 2 , θ 3 ) , the swimmer’s global angular velocity ω , the angular velocities of the joints ( ω 1 , ω 2 , ω 3 ) , and the velocity of the center of mass ( v x , v y ) .
We vary the length of the segments between the training and the test distributions.
The actions are chosen to be proportional to their corresponding angles.
The synthesized state machine policy is shown in Figure 20.
B. H YPER - PARAMETERS
There are three main hyper-parameters in our algorithm:
• The maximum number of segments/modes in a loop-free policy.
A large number of segments makes the teacher’s numerical optimization slow, while a small number of segments might not be sufﬁcient to get a high reward.
• The maximum time that a segment can be executed for in a loop-free policy.
This maximum time constraint helps the numerical optimization to avoid local optima that arise from executing a particular (non-convex) action function for too long.
• The parameter λ in Section 4.1.
This parameter strikes a balance between preferring highreward loop-free policies versus preferring policies that are similar to the state-machine learned so far.
The ﬁrst two parameters solely affect the teacher’s algorithm; thus, we choose them by randomly sampling from a set and select the one that produces high-reward loop-free policies.
We use λ = 0 .for
all our experiments.
B. T HE D IRECT -O PT B ASELINE
For this baseline, we convert the problem of synthesizing a state machine policy into a numerical optimization problem.
To do this, we ﬁrst encode the discreteness in the grammar for switching conditions into a continuous one-hot representation.
For example, the set of expressions o
[ i ] ≤ α 0 or o
[ i ] ≥ α 0 are encoded as α s ( α 1 o [1] + α 2 o [2] + · · · α n o [ n ]) ≤ α 0 with constraints − 1 ≤ α s ≤ 1 , ∀ i ∈ { 1 , ..., n } .
≤ α i ≤ 1 and (cid:80) ni
=
α i = 1 .
The choices between the leaf expressions, conjunctions, and disjunctions are also encoded in a one-hot fashion.
We also tried an encoding without the extra constraints on α —i.e., the switching conditions are linear functions of the observations.
We would expect the linear encoding to be less generalizable than the one-hot encoding.
However, we found that it is hard to even synthesize a policy that works well on the training set with either of the encodings.
Another difﬁculty with direct optimization is that we need to optimize the combined reward from all the initial states at once.
In contrast, the numerical optimization performed by the teacher in our approach can optimize the reward for each initial state separately.
To deal with issue, we use a batch optimization techinique that uses 10 initial states for every batch and seeds the starting point of the numerical optimization for each batch with the parameters found so far.
It also restarts the process with a random starting point if the numerical optimization stalls.
We carryout this process in parallel using 10 threads until either a solution is found or the time exceeds 2 hours.
B. RL B ASELINES
We use the PPO2 implementation from OpenAI Baselines (Dhariwal et al., 2017) with the standard “mlp” and “lstm” networks for our RL baselines using 1e7 timesteps for training.
Environment featurization.
We used the same action spaces, observation spaces, and the set of initial states that we used for our approach.
One exception is the Car benchmark, for which we appended the observation vector with observations from the previous timestep.
This modiﬁcation was essential for the RL baseline to achieve a good performance on the training dataset.
Designing reward functions.
While our approach takes in a safe speciﬁcation φ S ( x ) ≤ 0 and a goal speciﬁcation φ G ( x ) ≤ 0 , the RL baselines need a reward function.
For the classic control problems such as cartpole, pendulum, acrobot, mountain car and swimmer, we used the standard reward functions as speciﬁed by their OpenAI environments.
For Quad and QuadPO benchmarks, since the goal is to avoid collisions for as long as possible, we use a reward of 1 for every timestep that the agent is alive and the agent is terminated as soon as it collides with any of the obstacles.
Designing the reward function for the Car benchmark was tricky, because this benchmark has both a goal and a safety speciﬁcation, and ﬁnding a right balance between them is crucial for learning.
We tried various forms of rewards functions and ﬁnally, found that the following version achieves better performance on the training distribution (on the metric that measures the fraction of roll-outs that satisfy both the goal and the safety property):
[Equation]
otherwise which adds the numerical error for not satisfying the goal with a constant negative error ( − L ) if the safety speciﬁcation is violated at any time-step.
We tried different values for L ∈ { 0 . 1 , 1 , 2 , 10 , 20 } and found that L = 10 achieves the best performance on the training distribution.
Hyper-parameters search.
We performed a search over the various hyper-parameters in the PPO algorithm.
We ran 10 instances of the PPO2 algorithm with parameters uniformly sampled from the space given below, and chose the one that performs well on the training distribution.
This sampling is not exhaustive, but our results in Figure 4 (left most) show that we did ﬁnd parameters that achieve good training performances for most of our benchmarks.
• The number of training minibatches per update, nminibatches = {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048}.
For the lstm network, we set this hyper-parameter to 1.
• The policy entropy coefﬁcient in the optimization objective, ent_coef = {0.0, 0.01, 0.05, 0.1}.
• The number of training epochs per update, noptepochs = {3..36}.
• The clipping range, cliprange = {0.1, 0.2, 0.3}.
• The learning rate, lr =
[5e-6, 0.003].
[Table]
[Table]
C A DDITIONAL R ESULTS C. A DDITIONAL P ERFORMANCE R ESULTS
Figure 9 shows the training and test performance of the acrobot and mountain car benchmarks.
We can see that both our approach and RL generalizes for these benchmarks.
Figure 10 qualitatively analyzes the policies learned by our approach versus RL for the Pendulum benchmark.
We can see that the RL policy performs slightly sub-optimally compared to our policy.
Figure 11 shows the trajectories from the learned state machine policy and RL policy on Swimmer for a train environment and a test environment.
While both policies generalize, the swimmer with the state machine policy is slightly faster (it takes about 35s to cover a distance of 10 units while the RL policy takes about 45s).
C. A NALYSIS OF R UNNING T IME
Figure 12 shows the synthesis times for various benchmarks.
It also shows the number of studentteacher iterations, and the time spent by the teacher and the student separately.
The teacher optimizes the loop-free policies for different initial states in parallel.
The student optimizes the switching conditions between different pairs of modes in parallel.
We used a parallelized implementation with 10 threads, and report the wall clock time.
C. E XAMPLES OF STUDENT - TEACHER INTERACTIONS
Figure 13 shows two consecutive iterations of the teacher-student algorithm for the example in Figure 1.
[Figure]
While both policies generalize, the swimmer with the state machine policy is slightly faster (it takes about 35s to cover a distance of 10 units while the RL policy takes about 45s)
[Figure]
• In the ﬁrst iteration, the teacher learns loop-free policies for set of initial states that maximizes the reward.
In the ﬁrst row of the ﬁgure, we can see that most of the trajectories reach the goal, but the trajectories are unstructured.
Then, the student learns a probabilistic state-machine that best captures the loop-free policies.
The second row shows the most probably trajectories from this state-machine for the different initial states.
We can see that these trajectories are more structured, but some of them (columns 1 and 2) do not reach the goal.
In the next iteration (row 3), the teacher now learns loop-free policies that both reach the goal as well as match the trajectories from the state-machine policy learned in iteration 1.
• The next state-machine policy (row 4), learned based on these more structured loop-free policies, satisﬁes the goal for all initial states in the training set; hence, our algorithm terminates.
Iteration 1:
Teacher
[Figure]
Iteration 2:
[Figure]