A Tale of Three Implicit Planners
and the XLVIN AgentPetar Veličković
DeepMind/ELLIS CSML Seminar Series
University College London
12 March 2021

In this talk:Architectures and algorithms
for implicit planning1

What are wedoing here?

Reinforcement learning (RL) settingReinforcement learning (RL) setting (with planning)Reinforcement learning (RL) setting (variables)
action, a

policy, π(a | s)transitions, P(s’ | s, a)rewards, R(s, a)state, sreward, r ; state, s’Reinforcement learning (RL) setting
action, a

Want to optimise:

Discounted
cumulative rewardG = ∑t≥0 𝛾t rt

policy, π(a | s)transitions, P(s’ | s, a)rewards, R(s, a)state, sreward, r ; state, s’2

Do you havea plan?

What’s in a policy?
●Policies acting purely through adapting to observed rewards are often called reactive●

In many cases, they require large quantities of data and are slow to adapt●

Planning ameliorates such issues by maintaining an explicit model of the world○ State transition model: s’ ~ fT(s, a)○ Reward model: r ~ fR(s, a)○ Typically trained from observed trajectories●

Using these models, a planner can simulate the effects of actions before taking them!
○ Comes with many beneﬁts if done properly...

What are the benefits?
●

Gains in data efﬁciency: Good model implies fewer interactions are needed to learn to act●

Strong models allow quickly adapting to previously unexplored situations●

Being mindful of the consequences of acting enables better safety●

Allowing to explicitly account for external factors (e.g. human interactions)●

Impactful for game-playing (AlphaGo) and across the sciences (Segler et al., Nature’18)●

Encouraging theoretically: perfect models allow for planning perfect policies!

Intro to value iteration●Value Iteration: dynamic programming algorithm for perfectly solving an RL environment

where v(s) corresponds to the value of state s.
●

Guaranteed to converge to optimal solution (ﬁxed-point of Bellman optimality equation)!

Optimal policy takes actions that maximise expected value: argmaxa ∑s’ V*(s’) P(s’ | s, a)●

BUT requires full knowledge of underlying MDP (P / R)

3How do youeven plan?

Learning the model●In many cases of interest, P and R will be hidden from us!
○ Hence, we need to learn them through interacting with the environment●

Assume access to an interaction dataset of trajectories {(s, a, r, s’)}○ “I performed action a in state s, observed reward r and transitioned to state s’.”○ This can be built up as we learn to act---but need to tradeoff exploring vs exploiting.

●

We can then learn the transition function fT by supervising.
○ Make fT(s, a) be predictive of s’.
○ Proceed similarly for learning fR(s, a).

●There are other ways to incorporate assumptions about the env dynamics into fT/fR○ Here we assume a “tabula rasa” approach.

Latent-space transition models●Especially, we will ﬁnd models quite attractive if they operate in a latent state-space.
○ This particularly plays nicely with neural network approximators●Assume we have encoded our state (e.g. with a NN) into embeddings, z(s) ∈ ℝk
○ Our transition model is then of the form T : ℝk x A → ℝk
○ Optimised such that T(z(s), a) ≈ z(s’)●Many popular methods exist for learning T in the context of self-supervised learning●

Contrastive learning methods try to discriminate (s, a, s’) from negative pairs (s, a, s~)Implicit planning●Planning methods vary in how they exploit this information●

Policies can plan by explicitly rolling out the model, and then deciding how to act based on
the outcomes of the rollouts (leading to model-based RL)○ This is a discrete process, and doesn’t always play trivially with neuralnetwork-based optimisation●Here we will focus on learning to plan implicitly (leading to implicit planning).
○ Support planning computations using (architectural) inductive biases○ But still optimise using model-free RL losses e.g. DQN/A3C (~ “model-free planning”)●

Easy to compose and blend existing model-free approaches – “best of both worlds”●

Major challenge: need to build differentiable planning components

4What do you
mean, the plan
must bediﬀerentiable?

What we have covered so far

What we have covered so far

What we have covered so far(This skip connection canbe quite powerful!)5Is it okay tostart withsomethingsimpler?

Value iteration, revisited●

Now we will focus explicitly on the Value Iteration update rule●At each step, update each state’s values, by considering immediate successors●

Such a function is actually already differentiable w.r.t. its inputs!
○ But determining successors and aggregating over them may be tricky●

It makes sense to start with special RL environments where this determination is easier...

Value iteration in grid worlds●

Each state has known neighbours●

Actions are deterministic●In this setting, VI amounts to...

Value iteration in grid worlds●

Each state has known neighbours●

Actions are deterministic●

In this setting, VI amounts to…

●

Computing sums of neighbouring values!

●

Does this remind you of something?

Grid world VI ~ Convolution!

Generic MDP VI ~ Graph convolution!

If you’d like to know more
If GNNs are new(ish) to you, I recently gave a useful talk on theoretical GNN foundations:
https://www.youtube.com/watch?v=uF53xsT7mjc

Value Iteration Network●

Exactly this idea is leveraged by Value Iteration Networks (Tamar et al., NeurIPS’16)●

Assuming the underlying MDP is discrete, ﬁxed and known…

●

We can perform VI-style computation by stacking a shared convolutional layer⟹ We have our differentiable planning module!

●

Original VIN paper mainly dealt with grid worlds and hence used CNNs
○ Extended to generic MDPs and GNNs by GVINs (Niu et al., AAAI’18)

6

Relax these
constraints bytomorrow!

Moving beyond known worlds●

Assuming the MDP is ﬁxed and known was quite helpful○ We never needed to estimate transition models○ Didn’t have to deal with continuous state spaces●To relax these constraints, we learn a transition function T○ To operate directly over state embeddings z(s)○ Using it, we may construct a “local MDP” around a current state●We already discussed how to train T in prior sections

Using a transition model to expand

We can use the transition model
on every action, to be exhaustive(~breadth-ﬁrst search)
Doesn’t scale with large action
spaces / thinking times; O(|A|K)
Can ﬁnd more interesting rollout
policies, e.g. by distillingwell-performing model-free ones.

TreeQN / ATreeC●

Assume that we have reward/value models, giving us scalar values in every expanded node●We can now directly apply a VI-style update rule!

●Can then use the computed Q-values directly to decide the policy●

Exactly as leveraged by models like TreeQN / ATreeC (Farquhar et al., ICLR’18)○ Also related: Value Prediction Networks (Oh et al., NeurIPS’17)

TreeQN / ATreeC in action

7

Uh oh, there’s
bottlenecks inour plan

High-level view●It’s good to take a recap and realise what we have done so far○ We mapped our natural inputs (e.g. pixels) to the space of abstract inputs○ (local MDP + reward values in every node)○ This allowed us to execute VI-style algorithms directly on the abstract inputs●

The VI update is differentiable, and hence so is our entire implicit planner.

Algorithmic bottleneck●

Real-world data is often incredibly rich●

We still have to compress it down to scalar values●

The VI algorithmic solver:○ Commits to using this scalar○ Assumes it is perfect!

●If there are insufﬁcient training data to properly estimate the scalars…

●

We hit data efﬁciency issues again!
○ Algorithm will give a perfect solution, but in a suboptimal environment

Breaking the bottleneck●Neural networks derive great ﬂexibility from their latent representations○ They are inherently high-dimensional○ If any component is poorly predicted, others can step in and compensate!

●

To break the bottleneck, we replace the VI update with a neural network!

●

As before, we can use graph neural networks to perform VI-aligning computations.

Algorithmic reasoning●GNN over state representations aligns with VI, but may put pressure on the planner○ Same gradients used to construct correct graphs and make VI computations●

To alleviate this issue, we choose to pre-train the GNN to perform value iteration-stylecomputations (over many synthetic MDPs), then deploying it within our planner●This exploits the concept of algorithmic alignment (Xu et al., ICLR’20)

Algorithmic reasoning●

GNN over state representations aligns with VI, but may put pressure on the planner○ Same gradients used to construct correct graphs and make VI computations●To alleviate this issue, we choose to pre-train the GNN to perform value iteration-stylecomputations (over many synthetic MDPs), then deploying it within our planner●This exploits the concept of algorithmic alignment (Xu et al., ICLR’20)●

Relying on prescriptions on how to build effective extrapolating GNN reasoners, provided
by “Neural Execution of Graph Algorithms” (Veličković et al., ICLR’20)

Putting it all together!

XLVIN (Deac et al., NeurIPS’20 DeepRL)XLVIN Components●Encoder (z: S → ℝk) provides state representations●Transition (T: ℝk x A → ℝk) simulates effects of actions in latent space○ Pre-trained & Fine-tuned on the TransE loss (observed trajectories)●Executor (X: ℝk x ℝ|A| x k → ℝk) simulates a planning algorithm (Value Iteration) in latentspace○ Pre-trained to execute VI on synthetic MDPs of interest, then frozen●Policy / Value Head, computing action probabilities and state-values given embeddings○ Use PPO as the policy gradient method●The soft executor enables VIN-like models (Tamar et al., NIPS’16) on general MDPs

Further insightIf you would like to know more details about constructing good GNN executors:

https://www.youtube.com/watch?v=IPQ6CPoluok

https://drive.google.com/file/d/1_EQ9Yu7VEkvr
HaVHl_WbT5ABvxrSNY-s/view?usp=sharing

Algorithmic reasoning survey
Our 43-page survey on GNNs for CO!
https://arxiv.org/abs/2102.09544
Section 3.3. details algorithmic reasoning,with comprehensive references.

8So... did ourplan work?

Results on low-data AtariThank you!
Questions?
petarv@google.com | https://petar-v.com
With many thanks to Andreea Deac, Ognjen Milinković, Pierre-Luc Bacon,Jian Tang, Mladen Nikolić, Jess Hamrick, Feryal Behbahani and Charles Blundell