Synthesizing Programmatic Policy for Generalization within Task Domain

1 Introduction

Deep reinforcement learning has made significant breakthroughs across many control tasks. However, the ability to generalize across diverse tasks in a task domain remains a challenge, even for state-of-the-art deep reinforcement learning algorithms [Packer et al., 2018]. Specifically, consider the 'Tower of Hanoi' task, characterized by three pillars and a set of discs of varying sizes. The primary objective is to transfer all the discs from the source pillar to the target pillar while adhering to specific constraints. The variations of the task occur when the number of discs changes; for example, a control policy learned from a 'Tower of Hanoi' task with three discs. When evaluating the policy in the 'Tower of Hanoi' with four discs, it exhibits poor performance under these conditions.

∗Corresponding author.

The primary distinction in the 'Tower of Hanoi' task domain is related to the number of discs. Through summarizing the similar states across the tasks, it is found that the 'Tower of Hanoi' with different numbers of discs shares the same substructure. It is feasible to obtain a general solution applicable across tasks based on identifying state patterns and reusing subroutines. Thus, we focus on the problem of learning such a policy capable of generalizing in a task domain with different variations such as differences in number or position, where the underlying dynamics are the same.

When examining a neural process in humans and animals, they possess the capability not only to acquire novel solutions but also to deduce subroutines and apply them to similar tasks. These subroutines can be applied to similar tasks to improve efficiency. Inspired by this phenomenon, we consider synthesizing a programmatic policy characterized by a conditional branch structure, which is capable of capturing shared subroutines and state patterns across different tasks in one domain. Such a policy consists of a set of conditional branch structures. Each condition is designed to capture the state patterns. Each branch is applied to computes an action as return for the agent to execute, expected to deduce a subroutine. By doing so, the programmatic policy is expressive to capture tasks of interest, e.g., subroutines used to improve decision making and state patterns facilitating the selection of subroutines. This enables the learned policy to generalize to unseen tasks, which involves shared subroutines and similar states.

Algorithms have been proposed to learn programmatic policies that generalize better than traditional neural network policies [Verma et al., 2018; Verma et al., 2019], while necessitating user demonstrations. A programmatic state machine policy [Inala et al., 2020] that inductively generalizes to the tasks requiring repetition. A programmatic policy with a conditional branch structure is proposed in [Qiu and Zhu, 2022], but the generalization across different tasks is not evaluated. The architecture of the programmatic policy is synthesized based on a context-free grammar, which is a tailored domain-specific language. The DSL describes the definition of a nested If-Then-Else structure. We also incorporate Recurrent Neural Network (RNN) blocks to synthesize the architecture of the programmatic policy. The hidden states of the RNN are retained across different episodes throughout the agent-environment interactions. This enables the agent to remember past knowledge and apply it to new tasks. A task distribution can be obtained by varying the certain aspects of the environment settings from a task domain. We utilize a meta-learning algorithm to train the programmatic policy across the task distribution.

We benchmark our method against the state-of-the-art reinforcement learning baseline methods and ablation methods in seven benchmarks, which are adapted from PDDLGym for task planning and Pybullet for robotic manipulation. The results demonstrate that the programmatic policy trained via our method exhibits generalization across different tasks. It also performs well in the more complex task that have not been encountered before.

Our method draws inspiration from a real-world scenario involving real robot tasks. We operate under the assumption that the value of the variations within an environment are unlikely to exhibit large fluctuation. For instance, consider the 'Tower of Hanoi' environment, where complexity escalates exponentially as the number of discs increases. Completing such a task with thousands of operations becomes nearly impossible for a robot, resulting from the risk of operational failure in real-world environment. Therefore, the the values of the variations we consider are all adhere to a reasonable range in this paper.

2 Method Overview

2.1 ProblemFormalization

We consider a task distribution denoted as p(TH), in which H represents a collection of tasks. An agent with a programmatic policy is in the form of PE,θ, where the θ presents the parameter of the policy. The parameter E denotes the program architecture which can be defined by Domain Specific Language (DSL). S and A represent the states and actions across the tasks. The problem is to train a policy PE,θ(at | st), that can generalize within TH. Formally, we model each task τi ∈ TH as a Markov Decision Process defined by a tuple {Si, Ai, Ti, Ri} where Si and Ai denote the environmental observation and action spaces, respectively. Furthermore, Ti : Si ×Ai ×Si → [0, 1] represents state transition probabilities, and Ri : Si ×Ai → R quantifies the corresponding rewards when transitioning between states. The parameter of the policy ̂ θi is estimated by maximizing the cumulative discounted reward Es0,a0,s1···∼PE,θi [∑∞0 γ t · Ri(st, at)] where γ ∈ (0, 1]. Subsequently, we update the policy PE,θ by the θi learned from each task.

2.2 Architecture Synthesis

A programmatic policy processes an environmental state as input and produces an action for the agent. Drawing inspiration from the program architecture search framework [Qiu and Zhu, 2022], we infer the architecture of a programmatic policy denoted as PE,θ by DSL.

The DSL is represented in Backus-Naur form [Winskel, 1993]. A context-free grammar, depicted in Figure 1, is crafted to define the programs to be learned. The E is nonterminals, representing the program. The C is an affine transformation to compute an action value. B is evaluated by RNN layers which acts as a condition. This design facilitates

Figure 2: Illustration of training programmatic policy

the retention of knowledge across various tasks. X ∈ Rm represents the state observation, m denoting the dimension. This DSL allows the derivation of programs. As an example, we can deduce a program as if B1 then C1 else E1 to if B1 then C1 else (if B2 then C2 else E2). The semantics of the program, such as if B1 then C1 else E1, are computed by a function denoted as �if B1 then C1 else E1�(x), where the variable x serves as input to the If-Then-Else program. yielding a real-valued vector as output. C ::= θc + θT · X, yielding a real-valued vector as output, is an affine transformation, where θ ∈ Rm×n and n represents the dimension of the action spaces.

To ensure differentiability, the program derived from the DSL can be interpreted as a numerical approximation:

, where σ represents the sigmoid function. By utilizing the sigmoid function, the If-Then-Else program is transformed into a differentiable expression with binary branch structure. The output of the sigmoid function represents the probability of selecting a particular branch. The definition of the policy architecture aids to capture the subroutines and state patterns.

2.3 Training of Programmatic Policy

The left part of Figure 2 illustrates a program derivation tree with a depth of three. The blue circle represents E which can be expanded to a if-then-else branch, and orange represents C which is to output the action value. According to the DSL expression, we can systematically expand a program into a program derivation tree. The program derivation tree represents all possible program derivations within a specified depth limit for program abstract syntax trees. From Figure 2, W is a vector containing the knowledge of the selection of Algorithm 1: Algorithm for training programmatic policy

Input: Distribution over tasks p(TH), Learning rate α, Meta Learning rate β, DSL E, Depth d Output: Trained policy Pθ 1 Derive Programmatic Policy Pθ via E and d 2 Initialize θ = {W, φ} 3 while not done do 4 Sample task Ti ∼ p(TH) 5 Sample Trajectories with Pθ 6 Store in ReplayBuffer Bi 7 Using PPO to estimate ̂ θi with learning rate α 8 θi+1 ← θi + β(θi - ̂ θi) 9 end 10 Extract Pθ={φ} by fixing an optimal W 11 while not done do 12 Sample task Tj ∼ p(TH) 13 Sample Trajectories with Pθ 14 Store in ReplayBuffer Bj 15 Using PPO to estimate ̂ θj with learning rate α 16 θj+1 ← θj + β(θj - ̂ θj ) 17 end

each layer's architecture. Each digit in W models a binary selection. For instance, probC1 represents the probability of expanding E1 into C1 and 1 - probC1 signifies the probability of expanding it into if B1 then C2 else E2. The value of probC1 is calculated using a Softmax function. Besides W, φ is another learnable parameter of the programmatic policy. According to the definition of DSL, φ encompasses the parameters from both the RNN layer denoted as B and the linear layer denoted as C. Thus, the parameter θ can be viewed as a combination of W and φ, denoted as θ = {W, φ}. As E derived from a DSL, E can be regarded as a constant. Then we focus on optimizing the θ of policy P.

In this paper, we consider Proximal Policy Optimization [Schulman et al., 2017] as the foundational reinforcement learning algorithm to train the policy as depicted in the right part of Figure 2. The actor P i represents the programmatic policy we defined, while the critic Qi is actually a neural network. We utilize a ReplayBuffer Bi for storing trajectories during training.

We adopt Reptile [Nichol and Schulman, 2018] as the learning algorithm for gradient update, which is shown in the bottom part of Figure 2. It operates by performing a stochastic gradient descent on the sampled tasks and updating the initial parameters toward achieving the final learned parameters specific to the given task. It is worth noting that we omit the fine-tuning process. The algorithm solely needs a black-box optimizer such as SGD or Adam and offers good computational efficiency and performance.

We propose an algorithm to demonstrate the synthesis of the programmatic policy, as in Algorithm 1. The algorithm takes as input training hyperparameters, a DSL description, the maximum depth of the program derivation tree denoted as d, and a task distribution. In line 1, a programmatic policy Pθ is deduced in the form of a program derivation tree with depths ranging from 1 to d based on the input DSL. In line 2, we initialize the parameters θ = {W, φ} to be learned. Specifically, we initialize W with a 50% probability for each branch selection. At the training stage, in line 4, a task is sampled from the task distribution. The trajectories are obtained by agent-environment interactions, in which we design a mechanism to help the agent utilize the RNN blocks. It is elaborated upon in the following section. Subsequently, the trajectories derived from these interactions are stored in a ReplayBuffer for training. PPO is employed to optimize the parameters. We update both {W, φ} sequentially using a bilevel optimization technique. According to Reptile, the formula β(θi - ̂ θi) can be considered as a gradient and subsequently utilized in a more advanced optimizer, such as Adam. After a certain number of training steps, the W and φ of the policy tend to stabilize. At this point, we can determine the specific architecture of the programmatic policy fixing the W, as depicted in line 10. Specifically, since W models a series of binary selections, the better architecture of the policy is selected with the maximum likelihood. Since the current policy has only one parameter φ, repeat the above steps to train the policy until it converges, in line 11 to 17.

2.4 Using RNN to Improve Agent-environment Interactions

The process of the agent interacting with the environment, as described in both line 6 and line 16 of Algorithm 1, is visualized in Figure 3. For each task obtained from a task distribution, multiple episodes are generated through agent's exploration. During each episode, the agent engages with a task environment. Once the agent generates an action at, the environment provides the corresponding reward rt, advances to the next state st+1, and determines if the episode terminates. This termination status is recorded using the flag dt, which is set to 1 if the episode ends or left at a default value of 0 otherwise. The input is constructed by combining the following elements: the next state st+1, action at, reward rt, and termination flag dt.

As the task changes, the agent must adjust its actions according to the MDP in which it believes it is currently located.

Thus, the agent aggregates all available information on past rewards, actions, and termination flags and continually adapts its policy. To facilitate the agent in learning from prior experience, the programmatic policy incorporates RNN. It utilizes actions and rewards from preceding time steps as training inputs. According to the DSL in Figure 1, certain blocks of the policy comprise recurrent neural network cells. The hidden state ht is a vector summarized from the programmatic policy. Using the hidden state ht+1 and input state st+1 as inputs, the policy generates the subsequent action at+1 and updates the subsequent hidden state ht+2. The policy's hidden state is retained across episodes but is not carried over between distinct tasks.

3 Experiments and Evaluation

We evaluate the effectiveness of our approach on two groups of challenging benchmarks 1 , as illustrated in Figure 4. One group comprises three benchmarks-Hanoi, Stacking and Hiking adapted from [Silver and Chitnis, 2020], where the action space is discrete. These benchmarks primarily serve to evaluate our approach in the context of task planning.

• Hanoi: Involves various-sized discs stacked in a pyramid formation on the source pillar. The goal is to methodically transfer all the discs, one at a time, to target pillar in Figure 4a. The primary variation in this environment is determined by the number of discs involved. In Hanoi, we encode the observation state as a 1×9 vector.
• Hiking: The character is required to traverse pathways and collect all stars on the map. Variations in this environment pertain to the quantity and positioning of these stars, as depicted in Figure 4b. The entire state of the map is encoded as the state of the environment, with the dimension of 1 × 80.
• Stacking: Multiple plates of varying sizes are scattered on a tabletop. The task is to systematically pick up these plates one by one using a gripper and to assemble them in descending order of size, shown in Figure 4c. In this environment, we encode the table status, the stacked plates and the location of the gripper as the environment state, with a dimension of 1 × 35. Variations in this environment arise from different quantities of plates and their respective positions.

Another group comprises four tasks which are PandaReach, PandaPush, PandaSlide, PandaStack. These benchmarks are developed within the Pybullet environment [Gallou´edec et al., 2021]. Each environment consists of a Panda robot and a set of manipulable objects. The benchmarks in this group aim to assess our approach in continuous action spaces such as motion control. The variations in the four benchmarks involve differences in both number and position of the manipulable objects, indicated by distinct colors. The target object needs to be manipulated to reach the target point with the corresponding color. When the number and position of the manipulable objects change, each task domain becomes very large.

1 Code and benchmarks:

https://github.com/V0idwu/

meta-prl-code

• PandaReach: This task involves manipulating the robot's end-effector to reach the target positions, as illustrated in Figure 4d.
• PandaPush: The objective is to use the robot's endeffector to push the blocks to the target positions of the corresponding colors, as shown in Figure 4e.
• PandaSlide: The goal is to use the robot's end-effector to slide the cylinders to the target positions of the corresponding colors, as depicted in Figure 4f.
• PandaStack: The aim of the task is to place the blocks in the correct positions of the corresponding colors and stack them in order, as shown in Figure 4g.

We encode the observation state as a 1×32 vector. The action of the robot is a 4-dimension vector, which contains a 3D coordinate for the end-effector and a finger control.

To highlight the advantages of our approach, we conduct a comparative analysis with the following baselines/ablations:

• PPO, SAC: PPO [Schulman et al., 2017], SAC [Haarnoja et al., 2018] serve as baselines in reinforcement learning, trained within a specific environment with fixed variations.
• HIRO: HIRO [Nachum et al., 2018] is a general offpolicy algorithm for hierarchical reinforcement learning, trained under identical conditions as PPO.
• PRL: PRL, the Programmatic policy [Qiu and Zhu, 2022], is learned within the framework of PPO and trained under the same conditions.
• ReptilePPO, ReptilePRL: Reptile serves as the foundational meta-learning algorithm. ReptilePPO and ReptilePRL are trained in a task domain with the PPO and PRL respectively.
• PPO-N: Trained within the same meta-learning algorithm of ReptilePPO and ReptilePRL. However, the parameter of the model is not updated in a meta-learning manner.
• Ours: Our approach, based on the ReptilePRL, incorporates an RNN structure and a designed agentenvironment interaction.

Performance Table 1 presents an evaluation of the performance of the methods PPO, SAC, HIRO, PRL, ReptilePPO, ReptilePRL, and our approach. We trained the three methods-PPO, SAC, HIRO and PRL in a fixed task, such as Hanoi-3. ReptilePPO, ReptilePRL and our approach are trained across tasks, e.g., the number of the discs in Hanoi ranging from 1 to 4, denoted as Hanoi-1234. Subsequently, we evaluated their performance in the trained tasks and unseen tasks respectively. Evaluation involves measuring the mean and variance of the required action steps for agents to complete tasks. A smaller number of action steps indicates better performance. In the Hanoi benchmark, a single test suffices because the initial state is the same under the different random seeds.

Overall, our method exhibits superior performance in the majority of benchmarks. Moreover, as the environment's complexity increases, the corresponding rise in the required number of steps remains within acceptable limits. In a few benchmarks, RetilePRL achieved the best results and ReptilePPO showed the best performance in Hiking-3. These results illustrate that employing a meta-learning training approach and a designed programmatic policy enables the learned policy to generalize effectively when facing unseen tasks. In contrast, PPO, SAC, HIRO and PRL exhibit poor performance in most benchmarks, with their performance varying significantly as the complexity of the benchmarks increases.

Table 1: Performance comparison for evaluating policies in the benchmarks is averaged over 10 random seeds. The performance is measured by mean number of steps an agent takes to achieve goals, along with standard deviations. PPO, SAC, HIRO, PRL are trained in specific environments with fixed variations. For example, Hanoi-3 denotes a 3-disc 'Tower of Hanoi' task with three discs. ReptilePPO, ReptilePRL and our approach are trained in environments with a range of variations, e.g., Hanoi-1234. The test domain contains the tasks during training as well as the unseen tasks.

The performance of these algorithms in benchmarks characterized by greater complexity is also illustrated in Table 1. The threshold for the maximum number of agentenvironment interactions is set to 500. For instance, when PRL is applied to Hanoi-7, the number of action steps required reaches 500, which means the failure of the task. Policies trained by PPO and HIRO demand approximately 400 steps to complete the task. When confronted with previously unseen tasks during training, policies trained with ReptilePPO or ReptilePRL demonstrate a degree of generalization by completing the task in roughly 100 steps. Evidently, these methods outperform the PPO, HIRO and PRL. Nonetheless, it exhibits subpar performance in intricate scenarios such as hiking-10 or PandaReach-8.

Policies trained using our method consistently exhibit commendable performance. With the exception of PandaSlide8 and PandaStacking-8, the performance of our method remains stable despite changes in the value of the variations in the benchmarks. Given the exponential complexity of Hanoi, the optimal number of actions for hanoi-7 is 127. In our method, the solution for this task requires 151 action steps, showcasing commendable performance compared to other methods. These results indicate that our approach has the ability to capture the state patterns when trained across the tasks in one domain. The designed policy architecture aids in summarizing useful subroutines. By capturing the state patterns and shared subroutines, our method showcases generalization when facing unseen tasks.

Ablation Study We investigate the impact of two ideas: the combination of a meta-learning framework with a designed programmatic policy and the utilization of RNN blocks in policy architecture. The convergence curves of PPO-N, Reptile, ReptilePRL, and our method (Ours) are compared, as illustrated in Figure 5 Each policy is trained in the task domain and is evaluated in a fixed environment respectively. The vertical axis represents the number of steps required for the agent to achieve its objective, while the horizontal axis denotes the number of iterations involving agent-environment interaction episodes.

First of all, the policies learned by PPO-N struggle to achieving convergence in all benchmarks. Combining the results from Table 1, policies trained across tasks tend to exhibit better performance. It suggests that programmatic policies tend to be more effective in learning sharing subroutines for tasks.

According to Figure 5, ReptilePRL and ReptilePPO achieve similar results in most benchmarks. However, compared to ReptilePPO, ReptilePRL performs more consistent performance across all benchmarks. For example, ReptilePPO struggles to learn to finish the tasks in PandaPush-5, PandaPush-7 and PandaPush-8. This suggests that programmatic policies tend to be more effective to learn sharing subroutines of a task domain. This confirms that a well-designed programmatic policies may be more effective in handling generalization.

In the comparison between our approach and ReptilePRL, both approaches attain commendable performance in most tasks. However, as scenario complexity increases, ReptilePRL struggles to maintain an optimal solution in complex scenarios, whereas our approach consistently converges to a satisfactory solution. This is likely due to the fact that when the agent interacts with environments and preserves the RNN's hidden states between each episode, the agent trained by our method effectively leverages prior knowledge and transfers it to new tasks. These results signify that the incorporation of RNN blocks into programmatic policies effectively improve the ability to generalize across tasks.

4 RelatedWork

Generalization in RL Generalizing out of data distribution is challenging for reinforcement learning algorithms. A set of research focus on vision-based models to enhance sample efficiency or generalization. One main idea is data augmentation for visual observations such as, random translate, crop, color jitter, patch cutout [Laskin et al., 2020; Yarats et al., 2021]. Moreover, [Hansen and Wang, 2021] proposes SOft Data Augmentation which decouples data augmentation from policy learning.

In robotics, models that perform well in simulators tend to exhibit reduced performance in the real world. Environment randomization is explored to bridge this gap [Tobin et al., 2017]. [Packer et al., 2018] claims that environment randomization is the most effective method so far to improve generalization ability based on experiments on several sets of MuJoCo. [Akkaya et al., 2019] introduces the Automatic Domain Randomization (ADR) algorithm to address the issue of models trained in simulated environments performing poorly in real environments. [Tzeng et al., 2020] treats the transition from simulator to reality as a transfer learning problem. The real world robotic controller is learned by the ideas of domain adaptation and paired image alignment. These methods use environment randomization for sim2real problem. CoinRun, designed to test the generalization performance of deep reinforcement learning algorithms, is introduced by [Cobbe et al., 2019]. An approach, randomizing convolutional neural netowork, is evaluated in CoinRun [Lee et al., 2019]. A regularization parameter is proposed as a positive role to improve the model's generalization. However, there are potential issues with increasing environment randomization, including: increased complexity of the environment, increased complexity of training and dramatically increased variance.

Another way to improve generalization is regularization. According to [Liu et al., 2019; Farebrother et al., 2018], it is claimed that L2 regularization can produce better results than entropy regularization, and L2 regularization can find a good balance point for model's ability and generalization. [Lu et al., 2020] considers deep reinforcement learning models as two parts: the perception layer and the decision-making layer. An information Bottleneck approach is proposed to constrain the information transmitting, due to the perception layer being more prone to overfitting to the current training environment.

Programmatic Policy There existing prior work that mentions programmatic policies to address motion control problems in reinforcement learning. [Trivedi et al., 2021] proposes that combine learning a program embedding space with unsupervised searching to yield a program that maximizes the return of a given task. Besides, the searching part is improved in [Liu et al., 2023]. Though, the program in this work is in the format of a procedure instead of a policy in reinforcement learning, which can not interact with environments.

Some approaches are proposed to learn programmatic policies from neural network based policies, such as learning decision tree policy by model compression [Bastani et al., 2018], learning a program by Bayesian optimization [Verma et al., 2018] or generating symbolic policy by an autoregressive recurrent neural network [Landajuela et al., 2021]. The main idea of these work is to learn programmatic policies by referring to a neural network.

An algorithm for learning programmatic state machine policies that can capture repeating behaviors is proposed in [Inala et al., 2020]. This method follows a teacher-student learning paradigm. Specifically, a programmatic state machine policy is guided by the supervision of a neural network policy trained by RL. Besides the training algorithm, we focus on the generalizability of the model when the variations of the environment change, such as in number or position. [Liu et al., 2018; Cui and Zhu, 2021] propose to encode program architecture search as learning the probability distribution over all possible program derivations induced by a context-free grammar. This allows the search algorithm to efficiently prune away unlikely program derivations to synthesize optimal program architectures. In [Qiu and Zhu, 2022], a programmatic policy is learned by architecture search for control problems without demonstrations. Follow these lines of work, we utilize meta-learning algorithm to synthesize a programmatic policy for generalization in a task domain.

5 Conclusion and FutureWork

We introduce an approach aimed at synthesizing programmatic policies that can capture subroutines across tasks in one task domain. In order to improve the effectiveness of training, we utilize a meta-learning algorithm to synthesize the programmatic policy. Additionally, we incorporate RNN blocks for the policy architecture, enabling the utilization of previously acquired knowledge. Experimental results demonstrate the effectiveness of our method in synthesizing a programmatic policy capable of generalizing across tasks, even extending to unseen tasks during training.

Building upon our current method, two avenues for future research emerge. The first involves expanding the applicability of our method to more intricate scenarios, such as the scope of the task domain getting larger. The second direction entails exploring the inclusion of additional grammars into the DSL, resulting in a programmatic policy endowed with a more intricate architecture capable of addressing more complex tasks.

Acknowledgements

This work was supported by the National Key R&D Program of China under Grant No. 2022YFB4501803.