Meta-Reinforcement learning as an optimal transport problem:
Meta-Reinforcement learning considers the problem of building a program that generates a distribution \mu_{\pi} of policies \pi\in\Pi that are "adapted" for a distribution \mu_{\tau} of task \tau\in\mathbb{T} with respect to a given metric \mathcal{R} . For instance, \mathcal{R}(\tau,\pi) can be the expected cumulative reward for a task \tau and policy \pi , i.e where state transitions are governed by \bm{s}_{t+1}\sim\mathcal{P}_{\tau}(.|\bm{s}_{t},\bm{a}_{t}) , actions are sampled according to the policy \pi : \bm{a}_{t}\sim\pi and initial state \bm{s}_{0} follows the distribution \rho_{0,\tau} . Provided the existence of an optimal policy \pi^{*} for any task \tau\in\mathbb{T} , we can define the distribution measure \mu_{\pi^{*}} of these policies over \Pi . Arguably, an ideal system aims at associating to any task \tau its optimal policy \pi^{*} , i.e, finding the transport plan \bm{\gamma} in the space \Gamma(\mu_{\mathbb{T}},\mu_{\pi^{*}}) of coupling distributions with marginals \mu_{\mathbb{T}} and \mu_{\pi^{*}} that maximizes \mathcal{R} .

\max\limits_{\gamma\in\Gamma(\mu_{\mathbb{T}},\mu_{\pi^{*}})}\mathbb{E}_{(\tau,\pi)\sim\gamma}\bigg{[}\mathcal{R}(\tau,\pi)\bigg{]}\quadwhere\quad\mathcal{R}(\tau,\pi)=\mathbb{E}_{\pi,\mathcal{P}_{\tau}}\bigg{[}\sum\limits_{t=0}^{T}r_{\tau}(\bm{a}_{t},\bm{s}_{t})\bigg{]} (1)

Most generally, problem ( 1 ) is intractable, since \mu_{\pi^{*}} is unknow or has no explicit form. Instead, previous approaches optimize a surrogate problem, by defining a parametric specialization procedure which builds for any task \tau , a sequence (\pi_{t})_{t} of improving policies (see Fig. 1 ). Defining \bm{\theta} some meta-parameters governing the evolution of the sequences (\pi_{t})_{t} and \mu^{\bm{\theta},\tau,t}_{\pi} the distribution measure of the policy \pi_{t} after learning task \tau during some period t , the optimization problem amounts to finding the meta-parameters \bm{\theta} that best adapt \pi_{t}\sim\mu^{\bm{\theta},\tau,t}_{\pi} over the task distribution.

\max\limits_{\bm{\theta}}\;\;\mathbb{E}_{\tau\sim\mu_{\mathbb{T}}}\bigg{[}\mathbb{E}_{\pi\sim\mu^{\bm{\theta},\tau,t}_{\pi}}\big{[}\mathcal{R}(\tau,\pi)]\bigg{]} (2)

Equation ( 2 ) cast meta-Reinforcement learning as the problem of building automatic policy control in \Pi . We discuss below desirable properties of such control for which we show that our proposed meta-learnt synaptic rule has a good potential.

\diamond Efficiency: How "fast" the distribution \mu^{\bm{\theta},\tau,t}_{\pi} is transported towards a distribution of high-performing policies. Arguably, an efficient learning mechanism should require few interaction steps t with the environment to identify the task rule and adapt its policy accordingly. This can be seen as minimizing the agent’s cumulative regret during learning [ 37 ] . This notion is also connected to the policy exploration-exploitation trade-off [ 38 , 39 ] , where the agent learning program should foster rapid discovery of policy improvements while retaining performance gains acquired through exploration. We show that our learnt synaptic update rule is able to change the agent transfer function drastically in a few updates, and settle to a policy when task structure has been identified, thus supporting one-shot learning of a task-contingent association rule in the Harlow task, adapting a motor policy in a few steps or exploring original environments quickly in the Maze experiment.

\diamond Capacity: That property defines the learner sensitivity to particular task structures and its ability to convert them into precise states in the policy space, which determines the achievable level of performance for a distribution of tasks \mu_{\tau} . This notion is linked to the sensitivity of the learner, i.e how the agent captures and retains statistics and structures describing a task. Precedent work has shown that memory-based Meta-RL models operate a Bayesian task regression problem [ 40 ] , and further fostering task identification through diverse regularizations [ 41 , 42 , 39 ] benefited performance of the optimized learning strategy. Because our mechanism is continual, it allows for constant tracking of the environment information and policy update similar to memory-based models. We particularly test this property in the maze experiment in Section 5 , showing that tuned online synaptic updates obtain the best capacity under systematic variation of the environment.

\diamond Generality: We refer here to the overall ability of the meta-learnt policy flow to drive \mu^{\bm{\theta},\tau,t}_{\pi} towards high-performing policy regions for a diverse set of tasks ( generic trainability ) but also to how general is the resulting reinforcement learning program and how well it transfers to tasks unseen during training ( transferability ). In the former case. since the proposed synaptic mechanism is model-free, it allows for tackling diverse types of policy learning, from navigation to motor control. Arguably, to build reinforcing agents that learn in open situations, we should strive for generic and efficient computational mechanisms rather than learnt heuristics. Transferability corresponds to the ability of the meta-learned policy flow to generally yield improving updates even in unseen policy regions of the space \bm{\Pi} or conditioned by unseen task properties: new states, actions and transitions, new reward profile. etc. We show in a motor-control experiment using the Meta-World benchmark that meta-tuned synaptic updates are a potential candidate to produce a more systematic learner agnostic to environment setting and reward profile. The generality property remains the hardest for current meta-RL approaches, demonstrating the importance of building more stable and invariant control principles.