PDSketch: Integrated Domain Programming, Learning, and Planning
Authors: Jiayuan Mao
Tomás Lozano-Pérez
Joshua B. Tenenbaum
Leslie Pack Kaelbling
Venue: NeurIPS 2022 Accept (2022)
Keywords: Model-Based Learning
Neuro-Symbolic Programming
Heuristic Planning

Abstract: This paper studies a model learning and online planning approach towards building flexible and general robots. Specifically, we investigate how to exploit the locality and sparsity structures in the underlying environmental transition model to improve model generalization, data-efficiency, and runtime-efficiency. We present a new domain definition language, named PDSketch. It allows users to flexibly define high-level structures in the transition models, such as object and feature dependencies, in a way similar to how programmers use TensorFlow or PyTorch to specify kernel sizes and hidden dimensions of a convolutional neural network. The details of the transition model will be filled in by trainable neural networks. Based on the defined structures and learned parameters, PDSketch automatically generates domain-independent planning heuristics without additional training. The derived heuristics accelerate the performance-time planning for novel goals.

PDSketch: Integrated Planning Domain Programming and Learning
Jiayuan Mao Tomás Lozano-Pérez! Joshua B. Tenenbaum' 1,2.3 Leslie Pack Kaelbling'
MIT Computer Science & Artificial Intelligence Laboratory MIT Department of Brain and Cognitive Sciences Center for Brains, Minds and Machines
Abstract
This paper studies model learning and online planning approach towards building flexible and general robots. Specifically, we investigate how to exploit the locality and sparsity structures in the underlying environmental transition model improve model generalization, data-efficiency, and runtime-efficieney We present a new domain definition language, named PDSketch. allows users to flexibly define high-level structures in the transition models, such as object and feature dependen- cies, in a way similar to how programmers use TensorFlow or PyTorch to specify kernel sizes and hidden dimensions of a convolutional neural network. The details the transition model will be filled in by trainable neural networks. Based on the defined structures and learned parameters, PDSketch automatically generates Aomain-independent planning heuristics without additional training. The derived heuristics accelerate the performance-timg planning for novel goals.
1 Introduction
long-standing goal in Al is to build robotic agents that are flexible and general, able to accomplish diverse set of tasks in novel and complex environments. Such tasks generally require robot to generate long-horizon plans for novel goals in novel situations, by reasoning about many objects and the ways in which their state-changes depend to combine model learning with online planning: the agent forms an internal representation of the environment's 's states and dynamics by learning from external actively-collected data, and then applies planning algorithms to generate actions, given a new situation and goal at performance time.
another. A promising solution strategy is
There are primary desiderata for system based on model model-learning and planning. First, the learning process should be data efficient, especially because of the combinatorial complexity of possible configuration of in the real world. Second, the learned model should be computationally efficient, making online planning a feasible runtime execution strategy.
A critical strategy for learning models that generalize well from small amounts of data and that can be deployed efficiently at runtime is to introduce inductive biases. In image processing, we leverage translation invariance and equivariance by using convolutions. In graph learning, we leverage permutation invariance by using graph neural networks. Two essential forms structure we can leverage in dynamic models of the physical world are locality and sparsity. Consider robot picking an object up off a table. At the object level, the operation only changes the states of the object being picked up and the robot, and only depends few other nearby objects, such the table (local). At the object feature level, only the configuration of the robot and pose of the object are changed, but their colors and frictional properties unaffected (sparse).
Classical hand-engineered approaches to robot task and motion planning have designed represen- tations that expose and exploit locality through lifting (or object-centrism), which allows relational
Correspondence to: jiayuanm @ edt. Project page: http://pdssetchh.aail.ii..a
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Model
Goul
all blocks yellow)
(saction :paraneters (201 202) ieffects (and (pose:tassign Po1 (?2f (pose Po2))
Policy
Domain- Independent
Planner (A*)
Action: nove (enel,,,,,, Itents)
θ
Primitive Policies (pick, place,
PDSketch (Stracture)
Neural Modules (Panometers)
Model Training
Environment (e-g. PyBullet)
Robot
Data (Offline Demonstration or Online Interaction)
Command
Training Time
Performance
The life cycle of a PDSketch model. A PDSketch model is composed of a model structure definition and a collection of trainable neural modules. The model parameters be learned from data. During performance time, the model is used by a domain- inmain-inddeendent planner to form a policy that directly interacts with the environment.
Figure! 1:
descriptions of objects and abstrac- tion over them, and through factoring, which represents different attributes of an object in a disentangled way [Gar- rett et al., 2021]. These representa- tions powerful and effective artic- ulations of locality and sparsity, but they traditionally laboriously hand- designed in a process that is very dif- ficult to get correct, similar to writing a full state-transition function as in Fig. 2a. This approach is not directly applicable to problems involving perception or environmental dynamics that are unknown or difficult to specify. In this paper, we present PDSketch, a model-specificatte structural sparsity priors and machine learning of continuous and symbolic aspects of the model. Just as human users may define the structure of convolutional neural network in TensorFlow [Abadi et al., 2016] or PyTorch| [Paszke al., 2019], PDSketch allows users to specify high-level structures of the transition model as in Fig. 2b (analogous to setting the kernel sizes), and uses machine learning to fill in the details (analogous learning the convolution kernels).
def nove into(o, c): o.pose c.pose [e, 0.1] if is b_oooo(c))) and is p_..nttee(e)::
def move into(o, c): orpropl (c(ppp))p)) if ??(0, c):
color color
propp (c.ppro)))
(a) fall specification of the transition model. (b) structure-only specification.
Figure 2: Defining both the transition model structure and implementation in Python (a) vs. defining only the structure while leaving details (the ?? functions) to be learned (b).
language that integrates human specification of
Fig. I depicts the life-cycle of a PDSketch model, PDSketch uses an object-centric, factored, symbolic language flexibly describe structural inductive biases planning domains (i.e., the model structure). A PDSketch model is associated with a collection of neural modules whose parameters can be learned from robot trajectory data that are either collected offline by experts or actively-collected by interacting with the environment. During performance time, PDSketch is paired with a domain- iemain-indeeeenent planner, such as A", and as a whole forms goal-conditioned policy. The planner receives the environmental state and the trained PDSketch model, makes plans in an abstract action space, and invokes primitive policies that actually generate robot joint commands.
Compared to unstructured models, such as single multi-layer perceptron that models the complete state transition monolithically, the structures specified in PDSketch substantially improve model generalization and data-efficiency in training. In addition, they enable the computation of powerful cemmain-indepencee planning heuristics: these are estimates of the cost-to-go from each state state satisfying the goal specification, which can be obtained from the structured transition model without any additional learning. They can be leveraged by A" to efficiently plan for unseen goals, specified in a first-order logic language.
We experimentally verify the efficiency and effectiveness of PDSketch in domains: BabyAI, an 2D grid-world environment that focuses on navigation, and Painting Factory, a simulated table-top robotic environment that paints and moves blocks. Our results suggest that 1) locality and sparsity structures, specified economically in a few lines of code, can significantly improve the data efficiency of model learning: 2) the model learning and planning paradigm enables strong generalization to unseen goal specifications. Finally, the domain-independeen heuristics automatically induced from the structures dramatically improve performance-tin time efficiency, especially for novel goal specifications.

---

Itwekt (423 D. 420
Itors1
pose
Itenes
Iteers
Ento( Itones)
(3750, 23300
color
Itori1
(caction iite
color
Iteea2
Ttwent
sparaeattrss (hol
attnes
Towert
Itewe2
ioffects (and
conse:cansipn tos (?ff Epose 2o2)))
Eaput Inage
Itweii (333.,
bpp
Icoril
Syse
20
Itear3
34550,381.5)
(tyne
Esyse (S2)
pose
color
(wetness: sassigo
color
vetnecs
Ttwelii
Metross
we
Segnentation
Itesks S, 27355
cirtiness Itwent
dirciinss
200 1ig
2tames
Theenii
syse
(a) Raw Observations
(b) Object-Caatiic State (Obseraaioo)
<c) Factored State
(d) Factored State (Nest Timestep)
(e) Onject-Cenrtee State
(Iaternul)
(Next Olsecvvtion)
Eacteder
Tramsiion Functius
Superridon
Figure 3: A factorized state representation and transition model for the robot painting domain. The raw observation (a) is first processed by an external perception module into an object-centric representation (b). This representation is further transformed into fine- ane-grained factorization (c). The transition function To can be defined over this factored state representation: each action may only change few factors of the state. Executing specific action move-into (item#1, item#6) produces the predicted factored state at the next timestep. During training time, we will be using the object-centric observation from the next timestep to supervise the learning of Te.
2 PDSketch
We focus on the problem of learning models for robot that operates in a space of world states that plans to achieve goal conditions that are subsets of S. A planning problem is a tuple (S, So.g. J.A. T), where Sp
is the initial state, g goal specification in first-order logic, A that the agent can execute, and T is a environmental transition model T S x A - S. The task of {a: E A} in which the terminal state ST induced
set of actions
planning is to output a sequence of actions by applying sequentially following T satisfies the goal specification g: eval(g, 8T) 1. The function eval(g, s) determines whether state s satisfies the goal condition g by recursively evaluating the logical expression and using learned neural groundings of the primitive terms in the expression.
At execution time, the agent will observe so and be given g from human input, such as a first-order logic expression corresponding to "all the apples are in a blue bowl." However, we do not assume that the agent knows, in advance, the groundings g (i.e. the underlying eval(g. s) function) or the transition model T. Thus, we need to learn g and T from data, in the form of observed trajectories that achieve goal states of interest.
Formally, we assume the training data given to the agent is a collection of tuples (s, a, I. succ). where 5 (st) is sequence of world states, π (ag) is the sequence of actions taken by the robot, is a goal specification, and suce = (suce,) is the "task-success" signal. Each suce; {0, 1} indicates whether the goal g is satisfied at state sy: succ; = eval(g, (v.)(..s;). The data sequences should be representative of the dynamics of the domain but need not be optimal goal-reaching trajectories.
It can be difficult learn transition model that is accurate state representations. For this reason, we generally assume an arbitrary latent space, $ for planning. The learning problem, then, is to find three parametric functions, collectively parameterized by A: state encoder En : S > & goal-evaluation function evals : G x + - {0, 1}. and transition model To : + A Φ. Although the domain might be mildly partially observable stochastic, our goal will be to recover the most accurate possible deterministic model on the latent space.
the long term
types of
Local, sparse structure. We need models that will generalize very broadly to scenarios with different numbers and types of objects in widely varying arrangements. To achieve this, we exploit structure to enable compositional generalization: throughout this work we will be committing an object-centric representation for s, a logical language for goals g. and a sparse, local model action effects.
We begin by factoring the environmental state into set of object states. Each s E S is a tuple (Ms, fa), where 24, is the set of objects in state s, denoted by arbitrary unique names (e.g., item#1, item#2). The object set Ud is assumed to be constant within single episode, but may differ in different episodes. The second component, fs. a dictionary mapping each object name to a fixed- dimensional otject-state representation, such as local image crop of the object and its position. We
3

---

can extend this representation to relations among objects, for example by adding (o 2) y E U as mapping from each object pair to a vector representation. We assume the detection and tracking of objects through time is done by external perception modules (e.g., object detectors and trackers).
We carry the object-centric representation through to actions, goals, and the transition model. Specifi- cally, we define predicate as a tuple (name, args, grounding), where args is list of A arguments and grownding a function from the latent representations of the objects corresponding to its arguments (in "²) into scalar or vector value. (This is a generalization of the typical use of the term "predicate," " which is better suited for use in robotics domains in which many quantities we must reason about are continuous.) For example, as illustrated in Fig. 3c, the predicate wetness takes a single argument as its imput, and returns a (latent) vector representation of its wetness property; its grounding might be a neural network that maps from the visual appearance of the object to the latent wetness value. Given a set of predicates, we define the language of possible goal specifications to be all first-order logic formulas over the subset of the predicates whose output type is Boolean. To evaluate a goal specification in state (4, f), quantification is interpreted in the finite domain U and fs provides an interpretation of object names into representations that can serve as input to the grounded predicates.
The transition model To is specified in terms of set object-parameterized action schema (name, args, precond, effect, T), where name is a symbol, args is a list of symbols, precond and effect are descriptions of the action's ' effects, described in section 2.1, and π is a parameterized primitive policy for carrying out the action in terms of raw perception and motor commands. These local policies can be learned via demonstration or reinforcement learning in a phase prior to the model-learning phase, constructed using principles of control theory, or combination of these methods. The set of concrete actions A available in a state s is formed by instantiating the action scheme with objects in universe U, We assume the transition dynamics of the domain (i.e., the effect of each action schema) are well characterized in terms of the changes of properties and relations of objects and that the transition model is lifted in the sense that it can be applied to domain instances with different numbers and types of objects. In addition, we assume the dynamics are local and sparse, in the sense that effects of any individual action depend on and change only small number attributes and relations of a few objects, and that by default all other objects and attributes are unaffected. Taking again action schema move-into innt an example, shown in Fig. 3d, only the states of object #1 and #6 are relevant to this action (but not #2, #3, etc.), and furthermore, the action only changes the pose and wetness properties of item#1 (but not the color and the type).
The factored representation also introduces a factored learning problem: instead of learning a monolithic neural network for Ta and evalg, the problem is factored into learning the grounding of individual predicates that appear in goal formulas, as well as the transition function for individual factors that were changed by an action.
2.1 Representation Language
The overall specification of evalg and To can be decomposed into two parts: 1) the locality and sparsity structures and 2) the actual model parameters, 0, such as neural network weights. We provide symbolic language for human programmers to specify the locality and sparsity structure of the domain and methods for representing and learning 0. If the human provides no structure, the model falls back to a plain object-centric dynamics model Zhu et al. [2018]. However, we will show that explicit encoding of locality and sparsity structures can substantially improve the data efficiency of learning and the computational efficiency of planning with the resulting models.
PDSketch is an extension of the planning-domain definition language [Fikes and Nilsson, 1971, Fox and Long, 2003], widely used formalism that focuses exposing locality and sparsity structure in symbolic planning domains. The key extensions are 1) allowing vector values in the computation graph and 2) enabling the programmer to use "blanks", which are unspecified functions that will be filled in with neural networks learned from data. Thus, rather than specifying the model in full detail, the programmer provides only a "sketch" [Solar-Lezama, 2008]. The two key representational components of PDDL are predicates and action schemas (operators).
Fig. 4 shows a simple example of PDSketch definition of predicates. All three predicates take single object ?o of type item as their argument, and return either a floating-point vector or scalar value from 0 to 1, indicating the score of a binary classifier. The image predicate simply refers to the raw image crop feature of the object. The is-yellow predicate's 's grounding takes a very simple form "(??f (color ?o))". The term ??f defines a slot whose name is f. It takes only one argument, the

---

(spredicates ; Input features
image 70 (per-obij)
(Inage seturn_typpe-wacton (float32), input] to iton)
(:derived
colorief (ComoNen
(color hetpretyeereeeooflloootao 20 itte)
(?2F (isage 20])
color 70 (per-obij
(iderived
Raw Observation (Image)
Is-yellow::
(1s-yellow Iton) ] pananster type
(ComoNet)
(?f (color Po])
function body
is-yellow 20
(per-obj. score) 0.0
0.9
0.9
8.1
0.0 0.0
(a) PISSketch definition of input featare of each objects: (inage 2o), and two derived prsture/prodisaees: (color ?0) and (is-yellow 70).
(b) The conesponning "computation graph" indiaeed by the definitions. color: : and 1s -yellooww if are customizable CNNs that are appliod identically each object in the input
Figure 4: A minimal example of defining derived features and predicates with blanks "??".
color of the object ?o, and outputs classification score, which can be interpreted as the score of the object ?o being yellow. The actual computation of the yellow predicate from the color value (as well as the computation of the color value from the image value) is instantiated in a neural network with trained parameters. The computation graph for the whole model can be built by recursively chaining the function bodies of predicate definitions.
Next, we illustrate how locality and sparsity structures can be specified for an action schema. Fig. defines an action schema name move-into with two effect components. First, highlighted in blue, the action changes the pose of object ?o1 to a new pose that depends on the current pose of the second object ?o2. Rather than hand-coding this detailed dependence, leave the grounding blank. In addition, in our domain, the wet wetness of an object may be changed when the object is placed into specific type of container. This is encoded by specifying conditional effect using the keyword when, with two parts: 1) a Boolean-valued condition g of some other predicates on the state (in this case, the types of the two objects), and 2) the actual "effect", in this case, change the wetness of ?o1 based on a function that considers the current wetness of ?01. The update will be applied only if the condition is true. To ensure the computation is differentiable, we make this condition "soft": let w be the current wetness, w' be the new wetness computed by function ??h, and c be the scalar condition value computed by function ??g. The updated value of the wetness will be cw' + +1 c)w. Note that all the pose, wetness, and type representations can be arbitrary latent vectors computed by an encoder from the raw input. Thus, the "blanks" ??f, ??g. ??h are indeed general neural networks. The effect definition here induces corresponding computation graph of neural network weights and state representation tensors.
Like PDDL, PDSketch has full support of first-order logic, including Boolean operations (and, or, not) and finite-domain quantifiers (V' and 3). They allow to define more complex structures in the domain of interest. We present our full language and more examples in the supplementary material.
2.2 Model Learning and Planning with PDSketch
Let 0 denote the collection of all learnable parameters required to complete a PDSketch domain definition into full model. This includes the parameters of the state encoder, all predicate groundings, and the definitions of slot-update functions that were left blank in the sketch. Recall that our training data are tuples of three sequences and goal formula (5, a, g, suce). Fundamentally, our objective is to minimize sum of two losses, one related to predicting the truth values of the goal formulas and one related to predicting the next state, given the previous state and action:
{ Σ BCE Eo(s), succ;) + ΣL LI (To(Eo(si) (.) Eoloenil)
[()) = Σ
where BCE is the binary cross-entropy classification loss and L1 is regression loss. To avoid degenerate local optima, we add "lookahead" loss term that combines both aspects, as detailed in the supplement. If the encoder Eg is constrained to be the identity, and there is no predicate-level structure, then the transition model essentially learns to be a per-object next-image predictor. More generally, including the encoder turns this into bisimulation objective [[.i et al., 2006]: we want to uncover a latent transition model that accurately predicts the reward signal (in this case, whether the goal is satisfied not) but does not necessarily reconstruct the input state representation.
5

---

Action
Itome1
Itessi
(Thaa Thveels)
pose
Itores
sove in
pose
Itwsts
(cactionn nove-Into
sparaneteers (7ol ho2)
Color
color
Itest1
seffects (and
(ponec
color
Itorss
rolor
Ttwato
(423.0 420.0)
(27500273.0)
(226 (pose bo2])
aptness
Theest
nove into::h
Accress
1tess1
(when
(Pgg
(1ype tos)
(Aype toa)
(2755.277.5)
dir dirtiness
Itemps
the cuppur of g.
conditioned on.
dintiness Theatl
Itests
(7755.2735)
casssgs 701
dirt
Treead
dintiness Itests
(22h
(ontness Pat)
Itest
Teeem
nove into::g
type Theead
type
1testo
Timestep C
Timestep tel
Figure 5: A computation graph of partial definition of the action "move-into."
In order to adjust 0 to minimize this loss, we must establish a differentiable computation graph. The encoder Eg will generally be relatively standard combination of convolutional and folne-con feed-forward neural network. Fig. 5 illustrates the computation graph associated with To for a particular choice of action a, and objects that serve its arguments. Importantly, note that there a substantial amount of parameter-tying: the same predicate-groundee multiple times, even when characterizing To for a single action, if that predicate appears multiple times (applied to different objects) in the preconditions or effects of the action. The computation graph shown here, as well those necessary to compute evala (g, Es(s; involve Boolean operators, which do not have useful derivatives for optimization. We address this by representing truth values as elements of the interval [0, 1] and approximate logical operations with the differentiable Gödel t-norms: not (p) = 1 and P2) = min (p1. P2). Vz.p(z) = minxp(z).
network might appear in
Once we have estimated 0 from data, can solve any planning problem in the domain given any starting state 80 E S and goal g expressed in terms of predicates for which we have groundings. The resulting transition model T can be used by variety of different planners. We will focus on forward search, constructing a tree rooted at (sp) with branches corresponding to the possible instantiations, a, of action templates A with the objects in the universe U associated with so, and next latent states computed by applying T. The search terminates when it reaches a node n in which evale(n, g) > .5. Unguided forward search can be very slow when the planning horizon is long or the branching factor is large (e.g., when there are many objects in an environment). To address this, we will use the A heuristic search algorithm using domain-indeen independent search heuristic directly derivable from the locality and sparsity structure defined in PDSketch.
2.3 Inducing Domain- Indeendden Heuristics
Because of the rich representational capacity of PDSketch models, in which values can be continuous and multidimensional, we cannot take advantage of the planning algorithms that operate on PDDL input, such Fast Downward [Helmert, 2006), which derive much of their efficiency from domain- independent heuristics. We cannot use their strategies in detail, but we take inspiration from the idea of deriving optimistic estimate of the cost reach the goal from a state by solving "relaxed" version of the problem, which is computationally easier than the original Bonet and Geffner [2001].
One way to construct a relaxed planning problem is to allow each predicate instance to take on multiple values at the same time. For examplc, the robot or an object can effectively be in multiple places at the same time. In this relaxation, computing the number of steps needed to change the value of a predicate instance can be done in polynomial time. We can use such a relaxation in the hFF heuristic [Hoffmann and Nebel, 2001] to get estimate of the cost to goal, by first chaining the actions forward until all the components of the goal condition have been made true, and then searching to recover small set of actions that can collectively achieve the goal under the relaxation.
To use hFF, however, we must reduce our continuous-space problem to a discrete-space one. Specif- ically, we discretizes all continuous state variables (poses, etc.) into designated number of bins (e.g., 128). Next, for each externally-defined functions, we learn a first-order decision tree to approxi- mate the computation (c.g., to approximate the neural network). Coneretely, use VQVAE [Van

---

Inductive Biases
Suec. Rate
Model
Obj- Centric.
Facing
Rob. Dyn.
Obj. Prop.
Basic
#Obj.Gen.
BC
Y
N
N
N
0.93
0.79
DI (S)
Y
N
N
N
0.91
0.82
DT (S+F)
Y
N
N
N
0.32
0.19
DreamerV2
N
N
N
N
0.96
0.79
PDS-Basc
Y
N
N
N
0.82
0.62
PDS-Abs
Y
Y
N
N
0.99
0.98
Figure 6: A screenshot of the BabyAl environment.
PDS-Rob
Y
Y
Y
1.00
1.00
Table 1: The planning success rate of different models on BabyAl.
(6) Date Efficiency (Adistractoes: 0)
(b) Dain Efficieney (Adistracions: 3)
(6) Planning Efficiency
2.5
1.0
2.0
P08-Abs
2.0
P08-Abt
5x. More Etficiest
POS-Rub
PO6-Rub
1.5
1.5
8.6
Succes
1.0
undentand the abstacles.
0.4
Blind
0.2
GNN
AO-Da
0.0
2000
4000
6000
2000
4000
6000
20
40
100
Episodes
Episoles
Espansions
Figure 7: (a) and (b) Data efficiency comparison for model learning. We compare three structures with different levels of abstractness, (c) Planning efficiency, measured as the number of expanded nodes for different heuristic computation methods.
Den Oord et al., 2017] to discretize the feature vectors. For all predicates whose output is latent embedding, we add a vector quantization layer after the encoding layers. We initialize the quantized embeddings by running K-means clustering over the item feature embedding from a small dataset, and finetune the weights on the entire training dataset for one epoch. We describe this in more detail in the supplementary material. Although our discretizations are inherently lossy on non- non-Boolean predicate values, since they are only used for search guidance, and not the forward simulation of the actual actions during planning, this does not affect the correctness of the overall algorithm.
3 Experiments
We evaluate PDSketch in two domains: BabyAI, a 2D grid-world and Painting Factory, a simu- lated tabletop manipulation task. We compare model with two model-free methods: Behavior Cloning [BC; Bain and Sammut, 1995] and Decision Transformer [DT; Chen et al., 2021]. We implement two DT variants: DT-S with only successful demonstrations, and DT-S+F with successful and failed demonstrations. We use graph neural networks [Gori et al., 2005, Battaglia et al., 2018) as their state encoder.
3.1 BabyAI
BabyAI [Chevalier- Boissver al., 2019] is an image- based 2D grid- environment where an agent can navigate around obstacles, pick and place objects. and toggle doors. In this paper, we focus on a specific level of BabyAl, namely ActionObjDoor. At this level, the agent navigates within 7x7 grid. The goals include go to an X, pick up X, open an X, where X is noun phrase, such as "blue key". We train all models on environments with doors and 4 objects. The offline dataset D contains both successful and failure demonstrations obtained by A" search. We extend PDSketch to interactive data gathering in the supplementary material. Additionally, we test generalization to environments with 6 doors and objects. Since objects may block agents, the agent needs to successfully uncover the underlying dynamics and plan to navigate around them.
We study three PDSketch models with various levels of locality and sparsity structures built in. The PDS-Base model has no built-in structures: each object represented as a holistic vector. This falls back to an object-centric transition model [Zhu et al., 201 8). PDS-Abs model that disentangles poses from object appearances. Importantly, it defines predicate facing and uses this concept to write action definitions (c.g., whether robot move will be blocked by the object it is facing). However, the

---

(saction forward-abs
:paraneters (?r robot)
robot pose
iten-feature $1
iten-feature #2
:precondition (and
reffect
2r
robot-dir
facing #2
facing #2
(??6
(robot: pose 2n) (robot-din te)
(foreach (?o itee)
forward::f
dlten-feturee::onnd-sddee 2o (facing 2r 20))
forward:
(??g (robot pase 2n) (robot dir 2n))
robot:
Figure 8: The PDSketch definition of the "forward" action and the corresponding computation graph for the PDS-Abs model. The mentioned predicate facing is an derived predicate represented as neural network that is jointly learned (omitted in the graph for brevity).
grounding of facing to be learned. Figure shows the detail definition of the forward action the PDS-Abs model. PDS-Rot contains predefined rules for robot movements but the object recognition modules are learned. We provide their definitions in the supplementary material.
Results. Table shows the results. Overall, PDSketch with more structure (PDS-Abs and PDS-Rob outperforms baselines by significant margin. From the performance of PDS-Abs, we see that even tiny amount of additional structure (e.g., an ungrounded predicate "facing") significantly improves the performance. especially when comes to generalization to more complex environments. See below for zoomed-in analysis for model learning. Furthermore, we hypothesize that the inferior performance of decision transformer the mixed training data (Succ+Fail) setting is due to the reward sparsity: the agent only gets reward when it reaches the goal. Thus, the failed demonstrations are generally hard to model as they are irrelevant to the goal specification. In addition, we compare models with DreamerV2, state-of otate-of-the-art the- model-based reinforce learning algorithm for image- based environments. Compared with BC and DT, we see DreamerV2 achieves slightly improved performance for the basic task, but does not show stronger generalization environments with more objects. We hypothesize this is because Dreamer still learns a fixed policy for execution.
Data efficiency. Fig. 7a-b shows model learning curves for the three PDSketch models, in the cases where there a single object in the environment and when there are objects. When the number of object is small, the environmental dynamics is easier to learn: both Base and Abs have a similar performance. However, when the number of object increases, PDSketch can leverage the inductive bias to learn the dynamics faster. Shown in the figure, even at the end of training, the model PDS- Basc has not successfully learned the correct movement dynamics, leading to its inferior performance during generalization to more objects.
Planning runtime efficiency. Fig. 7c quantifies the number of nodes expanded by when using different heuristic computations. Specifically, we compare our model with the blind heuristic (i.e., the heuristic value of state is when it satisfies the goal and otherwise 1.) and GNN GNN- parameterized heuristic learned from successful demonstrations. All results are based on the PDS-Abs model. First, learning-basee heuristics (GNN) shows improvement over the haseline "blind" heuristic, especially when the task is easy (requiring a small number of expansions). Second, the heuristic derived from PDSketch significantly improves the search efficiency. At the success rate of 0.8, our method is times more efficient than learning-bei based heuristics.
3.2 Robot Painting
Finally, we extend the framework to a tabletop manipulation task, built based on the tabletop environment of Zeng et al. [2020]. There are three zones, several bowls, and several blocks on the table. The robot can use its suction gripper to pick-and-place objects into a designated location. Blocks have 8 possible colors. Placing objects in bowl will paint the object to be the same color as the bowl. The task to paint the blocks and organize them in the target brown box. Our training time goal requires the robot to paint-and-plact relationship (e.g., a pink block is left of yellow block. Demonstration are collected using hand- crafted oracle policies following Zeng et al. [2020]. The offline dataset contains only successful trajectories. There are two built-in actions that the robot can execute. Fig. 9 shows the computation graph derived from PDSketch.
two objects. The goal contains their colors and their
8

---

Model
Succ. Rate 100
Succ Rate @ 1000
pose
nove-to ?pose)
pose Iterr1
left
state Itondi
state Itess1
1s-yellow
BC
0.70
0.95
pose Itent2
nove-into (?o1 ?o2)
pose Iterz
is brown ?o
DT PDS
0,56
0.97
Itend2
Itemt2
0,91
0.99
Predicted
Predicates
Table 2: Model performance on Paint- ing Factory, measured plan success rate. The number after @ the num- ber of demonstration trajectories. PDS- ketch models show stronger data effi- ciency than baselines.
Figure 9: Our encoding of the painting factory domain. Each
object has a pose and state. Action move-to directly changes the object pose. move-into moves object into a container, which possible changes the state (e.g., color) of
the object. Goals represented using relational predicates.
Training
Generalization (Unseen Goals and More Objects)
(exists 2a Ty (and (purplo 74) (yellow (lett-of by Pe) (in-target hx)
(exists Px 72 Cand
(forall hx (and
(mxists Px Ty 72 (and
[yullow 7x) (cyan Py) (red ?2) }>
(yellow hd) (in-target Px)
(pink 2x) (ourple Ty) (cyan 72) (in-target hx) hx) con hr
Place yellow block of purple block in the target box.
yellow, a eyan, (no need to put in the target bux). them in the targel box.
block
yellow
Point pink, a parple, and eyan black,
and stack them in the target box.
Figure From only one training task (left), model specified in PDSketch learns primitive concepts including object colors and spatial relations. They can be used in planning for unseen tasks given new first-order logic descriptions of the goal. The natural language descriptions are shown for readability.
State and action space. The state space in the painting factory is composed of list of objects, and list of containers. Both objects and containers are represented as a tuple of a 3D xyz location, and image crop. The image crop is generated camera plain, cropping out the image patch, and resizing it to 32 by 32. The action space contains two primitives. First, move-into is an action defined for each pair of item and container. changes the pose of the item (now the item will be inside the container). When the item is placed into bowl, will be painted into the same color as the bowl. The second action takes an object and 3D location as input, and moves the object directly to designated position (ignoring the rotation).
first computing the 2D bounding box of the object in the
Results. Table 2 shows the planning success rate on the training task with different amounts of training data. Overall, PDSketch is more data efficient than both baselines and achieves strong overall performance in this task. Much more importantly, Fig. 10 shows our generalization to novel task specifications, which involve more objects or new specifiers (e.g., forall). Our model generalizes directly these novel scenarios without any additional training. The quantitative performance for these three generalization tasks are: 0.99, 0.98, and 0.87, respectively, measured as success rate after executing the plan. The last task has lower success rate because stacking objects may fail due to controller and physical noises. Future work may consider building closed-loop controller that can recover such failures. We specify implementation details the supplementary material.
Related Work
Integrated model learning and planning is a promising strategy for building robots that can gener-
alize to novel situations and novel goals. Specifically, Chiappa et al. [2017]. Zhang et al. [2021].
Schrittwieser et [2020] learn dynamics from raw pixels; Jetchev al. [2013], Pasula al. [2007].
Konidaris et al. [2018], Chitnis et al. [2021]. Bonet and Geffner [2020]. Asai and Muise [2020].
Silver et [2021] assume access to the underlying factored states of objects, such as object colors and other physical properties. Our work bridges the gap between two groups: we do not assume pre-factored state representations given as the input, learn to ground different factors of object states. More importantly, instead of relying off-the-shelff models for predicting pixels or for

---

The objective of model learning and planning is to obtain a cond-cond policy [Kaelbling, 1993b,a]. While many others have studied model-free [Dayan and Hinton, 1992, Schaul et al., 2015] or hybrid model-free and model-based approaches [Pong et al., 2018, Nasiriany et al., 2019], in this paper, we focus learning structured models from data and leveraging domain-independent planners for planning in latent representations. Such structured models support novel goal specifications via a first-order logic language, and fomain-independen heuristics to accelerate search. Another alternative approach towards learning and planning is to learning to predict subgoals that needs to be achieved before other goals [Xu et al., 2019]. However, their work assumes that all preconditions and effects can be represented in a predefined symbolic language, and there are controllers for achieving -based representations learning first-order rules in symbolic domains, our work considers how and sparsity structures can improve the model learning and planning efficiency.
locality
Our model learns an object-factored transition model, which generally falls into the category of learning object-oriented MDPs [OO-MDPs; Guestrin al., 2003, Diuk al., 2008]. OO-MDPs have been applied to neural network- -waoed representations for visual domains [Walsh, 2010, Kansky et al., 2017, Zhu et al., 2018, Xia et al., 2019, Veerapaneni et al.. 2020] and textual game domains [Liu et al., 2021]. Our model leverages object-centric representations to encode structures. Furthremore, we exploit fine-grate une-grained local and sparse structures in model learning.
individual subgoals. By contrast, PDSketch supports generic for predicates and action effects.
Our framework PDSketch combines model definition in structured language (e.g., first-order logic) and neural network learning. is closely related to the idea of neuro-symbolic programming and differentiable programming for relational reasoning [Manhaeve al., 2018, 1, Riegel et al., 2020, Huang et al., 2021] and policy learning [Sun et al., 2019]. Our language PDSketch borrows ideas from earlier work "soft" execution of logical formulas but works on planning domains.
5 Conclusions and Limitations
PDSketch supports flexible and effective specification of locality and sparsity structures of en- vironment transition models. Leveraging these structures enables more data-efficient learning, compositional generalization to novel goal specifications and environmental states, and also domain- independent heuristics that accelerates performance-ting planning. Limitations we hope address include the lack of hierarchy and the inability of the method to discover novel factorizations.
In terms of societal impact, PDSketch suggests hybrid method for building intelligent robots, by combining human programs to specify high-level structures with learning. to reduce the amount of data and computation required for learning complex and long-horizon behaviors. It also enables more modularized systems: users can specify new tasks using the predicates and actions that have already been defined more easily than in unstructured approaches. However, PDSketch may require more expertise (or, at least, a different type of expertise) from programmers than other approaches.
Acknowledgement. We thank all group members of the MIT Learning & Intelligent Systems Group for helpful comments on an early version of the project. This work is in part supported by ONR MURI N0001 1--1--11007 the Center for Brain, Minds, and Machines (CBMM, funded by NSF STC award ICF-1231216), NSF grant 2214177, AFOSR grant FA9550-22 1800022---0244 ONR grant N00014-18-1-2844 the MIT Quest for Intelligence. MIT-IBM Watson Lab. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of our sponsors.
References
Martin Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving. Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore. Derek G. Murray, Benoit Steiner, Paul Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: A System for Large-Scale Machine Learning, In OSDI, pages 265-283, 2016.
Masataro Asai and Christian Muise. Learning Neural-Symbolis Symbolic Descriptive Planning Models via Cube-Space Priors: The Voyage Home (To Strips). wcXi:2204.122500 2020. 9
10

---

Michael Bain and Claude Sammut. A Framework for Behavioural Cloning. Machine Intelligence, pages 103-129, 1995.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational Inductive Biases, Deep Learning, and Graph Networks. arXiv:/806.0/266 02261, 2018. 7
Blai Bonet and Héctor Geffner. Planning as Hcuristic Search. Artif. Intell., 129(1-2):5-35)) 2001.
Blai Bonct and Hector Geffner. Learning First-Order Symbolic Representations for Planning from the Structure of the State Space. ECAI, 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision Transformer: Reinforcement Learning via Sequence Modeling. In NewrIPS, 2021. 7
Maxime Chevalier- Boisvert, Dzmitry Bahdanau. Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. BabyAl: First Steps Towards Grounded Language Learning With Human In the Loop. In ICLR, 2019. 7
Silvia Chiappa. Sébastien Racaniere, Daan Wierstra, and Shakir Mohamed. Recurrent Environment Simulators.
ICLR, 2017. 9
Rohan Chitnis, Tom Silver, Joshua B Tenenbaum, Tomas Lozano-Perez, and Leslie Pack Kaelbling. Learning Neuro-Symbolic Relational Transition Models for Bilevel Planning. arXiv:2105. /4074, 2021.
Peter Dayan and Geoffrey E Hinton. Feudal Reinforcement Learning. In NewrlPS, 1992. 10
Carlos Diuk, Andre Cohen, Michael Littman. Object-Oriented Representation for Efficient Reinforce-
ment Learning. In /CML, 2008. 10
Richard E Fikes and Nils J Nilsson. STRIPS: New Approach the Application of Theorem Proving Problem Solving. Artif. Intell., 2(3-4):189-088, 1971. 4
Maria Fox and Derek Long. PDDL2.1: An Extension to PDDI. for Expressing Temporal Planning Domains. JAIR, 20:61-124, 2003.
Caclan Reed Garrett, Rohan Chitnis, Rachel Holladay, Bcomjoon Kim, Leslic Pack Kaclbling, and Tomás Lozano-Péree Integrated task and motion planning. Annual Review of Control, Robotics, & Autonomous Systems, 4:265-293, 2021.
Marco Gori, Gabriele Monfardini, and Franco Scarselli. A new model for learning in graph domains. In IJCNN,
2005. 7
Carlos Guestrin, Daphne Koller, Chris Gearhart, and Neal Kanodia. Generalizing Plans to New Environments in Relational MDPs. In IJCAI, 2003. 10
Malte Helmert. The Fast Downward Planning System. JAIR, 16:1911,46, 2006.
Jörg Hoffmann and Bernhard Nebel. The FF Planning System: Fast Plan Generation through Heuristic Search. JAIR, 14:253-302, 2001.
Jiani Huang, Ziyang Li. Binghong Chen, Karan Samel, Mayur Naik, Le Song, and Xujie Si. Scallop: From Probabilistic Deductive Datahases to Scalahle Differentiable Reasoning. In NeurlPS, 2021. 10
Nikolay Jetchev, Tobias Lang, and Mare Toussaint. Learning Grounded Relational Symbols from Continuous
Data for Abstract Reasoning, In ICRA Workshop, 2013. 9
Leslie Pack Kaelbling. Hierarchical Learning Stochastic Domains: Preliminary Results. JCML, 1993a. 10
Leslie Pack Kaelbling. Learning to Achieve Goals. In IJCAI, 1993h. 10
Ken Kansky. Tom Silver, David A Mély. Mohamed Eldawy, Miguel Lázaro-Gredilla Xinghua Lou, Nimrod Dorfman, Szymon Sidor, Scott Phoenix, and Dileep George. Schema Networks: Zero-Shot Transfer with Generative Causal Model of Intuitive Physics. In /CML, 2017. 10
George Konidaris, Leslic Pack Kaclbling, and Tomas Lozano- Perez. From skills to symbols: Learning symbolic
representations for abstract high-level planning. JAIR, 61:215-22 2-9,, 2018. 9
Lihong Li, Thomas J Walsh, and Michael Littman. Towards Unified Theory of State Abstraction MDPs. ISAIM, 4(5):9, 2006.
Guiliang Liu, Ashutosh Adhikari, Amir-massoud Farahmand, and Pascal Poupart. Learning Object-rriented Dynamics for Planning from Text. ICLR, 2021.
Robin Manhaeve, Sebastijan Dumancie, Angelika Kimmig, Thomas Demeester, and Lue De Raedt. Deep- ProbLog: Neural Probabilistic Logic Programming. In NewrlPS, 2018. 10
Soroush Nasiriany. Vitchyr Pong, Steven Lin, and Sergey Levine. Planning with Coal-Conditioned Policies. In NewrIPS, 2019. 10
11

---

Hanna M Pasula. Luke S Zettlemoyer, and Leslie Pack Kaelbling. Learning Symbolic Models of Stochastic Domains. JAIR, 29:309-352, 2007.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. PyTorch: Imperative Style, High-Performance Deep Learning Library. In NeurIPS, 2019. 2
Vitchyr Pong, Shixiang Gu, Murtaza Dalal, and Sergey Levine. Temporal Difference Models: Model-Free Deep RL for Model-Based Control. In ICLR, 2018. 10
Ryan Riegel, Alexander Gray, Francois Luus, Naweed Khan, Ndivhuwo Makondo, Ismail Yunus Akhalwaya, Haifeng Qian, Ronald Fagin, Francisco Barahona, Udit Sharma, et al. Logical Neural Networks. In NeurlPS, 2020. 10
Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal Value Function Approximators. In
ICML. 2015. 10
Julian Schrittwieeer, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model. Nat., 588(7839)6004 600-6699, 609, 2020. 9
Tom Silver, Rohan Chitnis, Joshua Tenenbaum, Leslie Pack Kaelbling, and Tomas Lozano-Perez. Learning Symbolic Operators for Task and Motion Planning. In IROS, 2021. 9
Armando Solar-Lezama. Program Synthesis by Sketching. PhD thesis, University of Califomia, Berkeley, 2008.
Shao-Hua Sun, Te-Lin Wu, and Joseph Lim. Program Guided Agent. In ICLR, 2019. 10
Aaron Van Den Oord, Oriol Vinyals, et al. Neural Discrete Representation Learning. In NeurIPS, 2017.
Rishi Veerapaneni, John D Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity Abstraction in Visual Modcl-Based Reinforcement Learning. In CoRL, 2020. 10
Thomas J Walsh. Efficient Learning of Relational Models for Sequential Decision Making. PhD thesis, Rutgers
The State University of New Jersey-New Brunswick, 2010. 10
Victoria Xia, Zi Wang, and Leslie Pack Kaelbling. Learning Sparse Relational Transition Models. ICLR,
2019. 10
Danfei Xu, Roberto Martin-Martin, De-An Huang, Yuke Zhu, Silvio Savarese, and Li Fei-Fei. Regression planning networks. NeurlPS, 2019.
Andy Zeng, Pete Florence, Jonathan Tompson, Stefan Welker, Jonathan Chien, Maria Attarian, Travis Armstrong. Ivan Krasin, Dan Duong. Vikas Sindhwani, al. Transporter Networks: Rearranging the Visual World for Robotic Manipulation. In CoRL. 2020.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal. and Sergey Levine. Learning invariant representa- tions for reinforeement lcarning without reconstruction. In ICLR, 2021. 9
Guangxiang Zhu, Zhiao Huang, and Chongjie Zhang. Object-Oriemt Dynamics Predictor. In NeurIPS, 2018.
4, 7, 10
12

---

Checklist
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 Section 5.
(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]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed reproduce the main ex- perimental results (either in the supplemental material or as a URL)? [Yes] see the supplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] see Section and the supplementary material.
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes] see the supplementary material.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] see the supplementary material.
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 eite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent obtained from people whose data you' 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "f9Y-cjNOfzN", "reviewer": "Reviewer_LBrw", "summary": "The paper presents a framework for learning environment dynamics models more efficiently by leveraging small amount of human guidance. Specifically, the paper points out that for common robotics tasks, transition models defined in terms of object attributes and relationships are inherently sparse and local (e.g., an action only affects a small subset of the state space), and exploiting such sparsity and locality can improve the data efficiency of model learning. The paper proposes a framework that allows human users to specify such structures with a PDDL-like planning domain language, and the framework converts such specification to a differentiable computational graph that can be optimized using environment interaction data. The method is evaluated on a BabyAI toy environment and a simulated robotic manipulation environment. The paper shows that the resulting sparse model tends to be more sample-efficient and generalizes better to unseen tasks compared to a “dense” object-level transition model.", "questions": "Please refer to my main comments (\"weaknesses\") for questions & comments. They mainly concern limitations of the method and baseline comparison.", "limitations": "Like mentioned in the main comment, the method makes assumptions about the primitive skills as well as the nature of environment dynamics. I hope the authors can elaborate on these assumptions and how they can be relaxed in the next revision of the paper.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall I enjoyed reading the paper — the idea of leveraging a small amount of human priors to build a sparse transition model for online planning is conceptually simple and is well-executed in the paper. I especially like the manner that human prior is introduced. Specifying the action effect sketch using a PDDL-like language and letting the neural network fill in the low-level details is both elegant and practical.\n\nHowever, I do have some comments that are mainly focused on the generality of the proposed method, especially when one wishes to apply it to more complex manipulation tasks.\nThe method assumes a finite number of atomic primitive skills. However, it is unclear to me how this can be applied to more practical settings where skills have parameters, e.g., grasping angles and placing targets. One might argue that the neural network can in addition take in the continuous parameters as input to model the effects of continuously-parameterized skills. But the fact that the skills are not atomic would likely make the locality & sparsity assumption invalid. For example, a grasp may fail, and the end effector may accidentally bump into other objects. In addition, it seems like the paper assumes fully observed & deterministic transition, which is not practical. How would the model learning and the online planning process handle stochastic transition? I hope the authors could address these limitations in the next revisions of the paper.\n\nMy second set of comments concern the experiment evaluation, especially the choice of baselines.\n- First of all, DT seems to be not the best baseline to compare against, as it tends not to perform well with sparse reward, and it’s a model-free method. How about comparing to more powerful & structured transition models such as PlaNet / Dreamer (Hafner et al., 2019)? They also assume observation-space input and sparse reward, which could make a fair comparison.\n- Second, how does PDSketch PDSketch compare to unsupervised methods that try to discover transition sparsity such as Xia et al., 2018? It would be interesting to see whether an unsupervised method can discover similar sparsity patterns as human supervision and how data efficiency of PDSketch scales with the number of training samples.\n- Thirdly, PDSketch also shares some similarities with Regression Neural Networks (Xu et al., 2019), which is not mentioned in the related works. RPN shares the same input & output as PDSketch (object-centric encoding input & first-order logic goal). RPN also assumes that humans provide prior guidance on predicate dependencies, which is similar to the idea of sketching action definitions. I would like to see conceptual and possibly empirical comparison between PDSketch and RPN.\n- Finally, can a network learn sparsity by itself? How does the effect of human-specified sparse transition rules scale with the amount of data? Table 4 only compares the sample efficiency between PDSketch and model-free method. This calls for an experiment on how PD-Base or similar model-based approaches scales with the amount of training data.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "SRlw_7zx8Q6", "reviewer": "Reviewer_CyxQ", "summary": "The paper proposes PDSketch to generate the object and its features as high-level constructs of the transition model. Compared to earlier tools like PDDL, PDSketch allows vector values in computation and supports learning of unspecified functions from data. PDSketch is a collection of neural networks to convert encode the state, evaluate the goal and build a transition model. The experiments show that PDSketch adapts to novel first order logic goal objectives for navigation in BabyAI and Painting Factory. ", "questions": "- What is #Obj. Gen. in Table 1?\n- PDS-Rob uses predefined transition model, so it seems that it is only learning object properties. Can you clarify how feasible it is obtain the transition model without simulator in real tasks?\n- What is the notion of timestep in change of properties? For example, the wetness property changes within the next timestep or based on location change?\n- How to compare the capability of PDSketch over existing planning languages in terms of modeling general-purpose TAMP domains?", "limitations": "The paper briefly mentions the limitations as not modeling object hierarchy and unable to detect novel objects and properties. Some discussion on potential societal impact can be included as how much modeling time or iterations are needed for human designer to specify the object attributes, or how a system built with PDSketch is easy/hard to update with task specification as compared to more unstructured learning approaches like BC or DT. ", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The objective of PDSketch is to efficiently model the “local” changes induced by an agent’s actions and “sparse” changes that occur in object properties for generalization in unseen scenarios. The paper describes how predicates, goal evaluation and transition models can be represented as computation graph and learned. Overall, the paper is well-motivated and detailed. \n\nThere are assumptions that on object properties that are task relevant, defining its predicates, transition models. This may increase significantly as the different category of objects, their properties and their interdependence increases. Although the authors motivate the use of factored representation of transition model and eval function, this implies that all the objects and their properties are known at test time. The tasks considered are quasi-static, and it is not clear if the proposed approach can be somehow extended to handle dynamics, noisy sensing and actuation of any physical robot. While the proposed approach is shown to perform better than BC or DT and its variants, some stronger baselines can include recent variants of PDDL like [1] and as cited in related work.\n\n[1] Learning symbolic operators for task and motion planning, Tom Silver*, Rohan Chitnis*, Josh Tenenbaum, Leslie Kaelbling, Tomas Lozano-Perez, IROS, 2021", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "3P2RRIHJRIk", "reviewer": "Reviewer_PhAv", "summary": "The paper proposes a symbolic language for adding structure to world models in robotics and for defining task-independent heuristics. The language describes predicates, which explain the state of the environment, and actions, which change the state using pre-conditions and effects. Based on the descriptions, neural networks are instantiated to ground the language.\n", "questions": "1. What training data is required by the system? How difficult is it to collect this data on a physical robot?\n2. I believe it is assumed that each symbolic action has a pre-defined controller. Could you comment on the difficulties of coming up with the symbolic actions and their controllers in your robotics domain?", "limitations": "\"Limitations we hope to address include the lack of hierarchy and the inability of the method to discover novel factorizations.\" – To me, the biggest limitations of this method concern gathering training data with annotated symbols and coming up with a PDSketch description of a domain. I am not saying that these limitations warrant a rejection, but they must be discussed in the paper.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I find it difficult to review this paper because many pages are spent describing the PDSketch language and important model and experiment details are delegated to the appendix. For example, the use of a VQ-VAE for discretization is mentioned in the third-to-last sentence of the methods section. Moreover, the robot experiment, which should be the highlight of this paper, is described in only one paragraph. The scientific contribution might be sufficient for NeurIPS, but the paper needs significant work.\n\nStrengths:\n* The paper proposes a general way of automatically adding structure to world models for robotics.\n* The proposed method has a fall-back mechanism onto a generic world model, so that the designer does not have to specify every detail of the environment.\n\nWeaknesses:\n* There is no discussion of the difficulties of getting training data for the proposed model. It is not clear to me if ground-truth predicates are used to train the predicate classifiers. \n* The results show that the model can fit training tasks from few demonstrations, but the evaluation of generalization to novel tasks is limited. Figure 9 shows examples of four goal states, but there are no quantitative results.\n\n# Detailed comments\n\n– Starting on line 34, an argument is made against “classical hand-engineered approaches”. There are no citations and the argument is not revisited in the rest of the paper. I think it would be useful to devote a paragraph to task and motion planning, and how the proposed approach overcomes its limitations.\n\n– [1] is a recent paper that appears to be quite related to this work (e.g. they also put emphasis on sparse transition models). While it is cited in this paper, it is not explicitly discussed. I would like to see a short discussion of the similarities and differences.\n\n– PDS-Abs and PDS-Rob are both described in one sentence. This makes the BabyAI experiment difficult to understand without going into the appendix.\n\n– In the program shown in Figure 5, I think it should be “wetness::assign ?o1 (??h (wetness ?o2))” because we are checking if o2 is a container with paint?\n\n– Line 58: space after comma.\n\n– Reference for [1] has the year 2010 instead of 2019.\n\n– The robotic environment seems to be derived from RAVENS (https://github.com/google-research/ravens) / Transporter Nets. This is not stated explicitly.\n\n– There are a few pre-2018 works that could be cited for graph neural networks. E.g. [2].\n\n# References\n\n[1] Victoria Xia, Zi Wang, and Leslie Pack Kaelbling. Learning Sparse Relational Transition Models. In ICLR, 2019.\n\n[2] M. Gori, G. Monfardini, and F. Scarselli. A new model for learning in graph domains. In Proceedings. 2005 IEEE International Joint Conference on Neural Networks, 2005., volume 2, pages 729–734 vol. 2, 2005.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept