RiteeshGrandhi/sciclaimeval-shared-task huggingface dataset:
Notation
Throughout, we refer to “tasks" as Markov decision processes (MDP) defined by the following tuple \tau=(\mathcal{S},\mathcal{A},\mathcal{P},r,\rho_{0}) , where \mathcal{S} and \mathcal{A} are respectively the state and action sets, \mathcal{P}:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\mapsto[0,1] refers to the state transition distribution measure associating a probability to each tuple (state, action, new state), r:\mathcal{A}\times\mathcal{S}\mapsto\mathbb{R} is a bounded reward function and \rho_{0} is the initial state distribution. For simplicity, we consider finite-horizon MDP with T time-steps although our discussion can be extended to the infinite horizon case as well as partially observed MDP. We further specify notation when needed by subscripting symbols with the corresponding task \tau or time-step t .
Experiments
In this section, we explore the potential of our meta-learnt synaptic update rule with respect to the three properties of the meta-RL problem exposed in section 2 . Namely, 1) efficiency , 2) capacity and 3) generality of the produced learning algorithm. We compare it with three state-of-the-art meta-RL models based on different adaptation mechanisms: RL 2 [ 48 ] a memory-based algorithm based on training a GRU-cell to perform Reinforcement Learning within the hidden space of its recurrent unit, MAML [ 61 ] which performs online gradient descent on the weights of three-layer MLP and PEARL [ 78 ] , which performs probabilistic task inference for conditioning policy. Details for each experimental setting are further discussed in S.I. and the code can be found at https://github.com/mathieuchal/metods2022 .
Efficiency: One-shot reinforcement learning and rapid motor control
To first illustrate that learnt synaptic dynamics can support fast behavioral adaptation, we use a classic experiment from the neuroscience literature originally presented by Harlow [ 79 ] and recently reintroduced in artificial meta-RL in [ 54 ] as well as a heavily-benchmarked MuJoCo directional locomotion task (see Fig. 3 ). To behave optimally in both settings, the agent must quickly identify the task rule and implement a relevant policy : The Harlow task consists of five sequential presentations of two random variables placed on a one-dimensional line with random permutation of their positions that an agent must select by reaching the corresponding position. One value is associated with a positive reward and the other with a negative reward. The five trials are presented in alternance with periods of fixation where the agent should return to a neutral position between items. In the MuJoCo robotic Ant-dir experiment, a 4-legged agent must produce a locomotive policy given a random rewarded direction.

\diamond Harlow : Since values location are randomly permuted across presentations within one episode, the learner cannot develop a mechanistic strategy to reach high rewards based on initial position. Instead, to reach the maximal expected reward over the episode, the agent needs to perform one-shot learning of the task-contingent association rule during the first presentation. We found that even a very small network of N=20 neurons proved to be sufficient to solve the task perfectly. We investigated the synaptic mechanism encoding the agent policy. A principal component analysis reveals a policy differentiation with respect to the initial value choice outcome supported by synaptic specialization in only a few time-steps (see Figure 1 -c). We can further interpret this adaptation in terms of sharp modifications of the Hopfield energy of the dynamic weights (Figure 3 -e). Finally, the largest synaptic variations measured by the sum of absolute synaptic variations occur for states that carry a non-null reward signal (see S.I). These results suggest that the recursive hebbian update combined with reward feedback is sufficient to support one-shot reinforcement learning of the task association rule.

\diamond MuJoCo Ant-dir : We trained models for performing the task on an episode of 200 timesteps and found that MetODS can adapt in a few time-steps, similar to memory-based models such as RL^{2} , (Figure 3 -f) thanks to its continual adaptation mechanism. By design, MAML and PEARL do not present such a property, and they need multiple episodes before being able to perform adaptation correctly. We still report MAML performance after running its gradient adaptation at timestep t=100. We further note that our agent overall performance in a single episode is still superior to MAML performance reported in [ 61 , 49 ] when more episodes are accessible for training.
Capacity : Maze exploration task
We further tested the systematicity of MetODS learnt reinforcement program on a partially observable Markov decision process (POMDP): An agent must locate a target in a randomly generated maze while starting from random locations and observing a small portion of its environment (depicted in Figure 4 ). While visual navigation has been previously explored in meta-RL [ 48 , 49 , 31 ] , we here focus on the mnemonic component of navigation by complexifying the task in two ways, we reduce the agent’s visual field to a small size of 3x3 pix. and randomize the agent’s position after every reward encounter. The agent can take discrete actions in the set \{\text{up,down,left,right}\} which moves it accordingly by one coordinate. The agent’s reward signal is solely received by hitting the target location, thus receiving a reward of 10 . Each time the agent hits the target, its position is randomly reassigned on the map (orange) and the exploration resumes until 100 steps are accumulated during which the agent must collect as much reward as possible. Note that the reward is invisible to the agent, and thus the agent only knows it has hit the reward location because of the activation of the reward input. The reduced observability of the environment and the sparsity of the reward signal (most of the state transitions yield no reward) requires the agent to perform logical binding between distant temporal events to navigate the maze. Again, this setting rules out PEARL since its latent context encoding mechanism erases crucial temporal dependencies between state transitions for efficient exploration. Despite having no particular inductive bias for efficient spatial exploration or path memorization, a strong policy emerges spontaneously from training.

\diamond Ablation study and variations: We explored the contribution of the different features combined in MetODS update rule, showing that they all contribute to the final performance of our meta-learning model. First, we tested the importance of the element-wise tuning of plasticity in weight-based learning models and note that while it adversely affects MAML gradient update, it greatly improves MetODS performance, suggesting different forms of weights update. Second, we verified that augmenting recursivity depth S was beneficial to performance, consistently with Harlow results. Third, we noted that transforming the rightmost vector of the writing equation in 5 with a linear projection ( Linear in figure 4 , see S.I for full experimental details) yields major improvement while non-linear ( MLP ) does not improve performance. Finally, we additionally test the capability of the learnt navigation skills to generalize to a larger maze size of 10 \times 10 pix. unseen during training. We show that MetODS is able to retain its advantage (see figure 4 and table 5 for results).
Generality : Motor control
Finally, we test the generality of the reinforcement learning program learnt by our model for different continuous control tasks:

MetaWorld: First, we use the dexterous manipulation benchmark proposed in [ 81 ] using the benchmark suite [ 82 ] , in which a Sawyer robot is tasked with diverse operations. A full adaptation episode consists in N=10 rollouts of 500 timesteps of the same task across which dynamic weights are carried over. Observation consists in the robot’s joint angles and velocities, and the actions are its joint torques. We compare MetODS with baseline methods in terms of meta-training and meta-testing success rate for 3 settings, push, reach and ML-10 . We show in Fig. 6 the meta-training results for all the methods in the MetaWorld environments. Due to computational resource constraints, we restrict our experiment to a budget of 10M steps per run. While we note that our benchmark does not reflect the final performance of previous methods reported in [ 81 ] at 300M steps, we note that MetODS test performance outperforms these methods early in training and keeps improving at 10M steps, potentially leaving room for improvement with additional training. (see S.I for additional discussion). Finally, we note that all tested approaches performed modestly on ML10 for test tasks, which highlights the limitation of current methods. We conjecture that this might be due to the absence of inductive bias for sharing knowledge between tasks or fostering systematic exploration of the environment of the tested meta-learning algorithms.

Robot impairment: We also tested the robustness of MetODS learnt reinforcement programs by evaluating the agent ability to perform in a setting not seen during training: specifically, when partially impairing the agent motor capabilities. We adopt the same experimental setting as section 5.1 for the Ant and Cheetah robots and evaluate the performance when freezing one of the robots torque. We show that our model policy retains a better proportion of its performance compared to other approaches. These results suggest that fast synaptic dynamics are not only better suited to support fast adaptation of a motor policy in the continuous domain, but also implement a more robust reinforcement learning program when impairing the agent motor capabilities.
Discussion
In this work, we introduce a novel meta-RL system, MetODS, which leverages a self-referential weight update mechanism for rapid specialization at the episodic level. Our approach is generic and supports discrete and continuous domains, giving rise to a promising repertoire of skills such as one-shot adaptation, spatial navigation or motor coordination. MetODS demonstrates that locally tuned synaptic updates whose form depends directly on the network configuration can be meta-learnt for solving complex reinforcement learning tasks. We conjecture that further tuning the hyperparameters as well as combining MetODS with more sophisticated reinforcement learning techniques can boost its performance. Generally, the success of the approach provides evidence for the benefits of fast plasticity in artificial neural networks, and the exploration of self-referential networks.
Broader Impact
Our work explores the emergence of reinforcement-learning programs through fast synaptic plasticity. The proposed pieces of evidence that information memorization and manipulation as well as behavioral specialization can be supported by such computational principles 1) help question the functional role of such mechanisms observed in biology and 2) reaffirms that alternative paradigms to gradient descent might exists for efficient artificial neural network control. Additionally, the proposed method is of interest for interactive machine learning systems that operates in quickly changing environments and under uncertainty (bayesian optimization, active learning and control theory). For instance, the meta RL approach proposed in this work could be applied to brain-computer interfaces for tuning controllers to rapidly drifting neural signals. Improving medical applications and robot control promises positive impact, however deeper theoretical understanding and careful deployment monitoring are required to avoid misuse.
Acknowledgments and Disclosure of Funding
This work was supported by ANR-3IA Artificial and Natural Intelligence Toulouse Institute (ANR-19-PI3A-0004), OSCI-DEEP ANR
(ANR-19-NEUC-0004), ONR (N00014-19-1-2029), and NSF (IIS-1912280 and EAR-1925481). Additional support provided by the Carney Institute for Brain Science and the Center for Computation and Visualization (CCV) via NIH Office of the Director grant S10OD025181. The authors would like to thank the anonymous reviewers for their thorough comments and suggestions that led to an improved version of this work.
Optimization
Defining the weight parameters \bm{W} of MetODS as dynamic variables lifts the optimization problem ( 2 ) into a functional space of control functions parameterized by \bm{\theta} . Hence, meta-optimizing the control necessitates the estimation of gradients with respect to \bm{\theta} over the space \mathbb{T} and for any possible trajectory \pi_{t} in \Pi . Interestingly, previous meta RL approaches have performed policy gradient optimization by sampling a single policy trajectory \pi_{t}\sim\mathcal{M}_{\bm{\theta}}(\tau) over M multiple tasks, showing that it is sufficient to obtain correct gradient estimates on \bm{\theta} . We proceed in the same way, by estimating the gradient policy update integrated over the space of tasks as mini-batches over tasks.

\begin{split}\frac{\partial}{\partial\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{]}\hskip 14.45377pt\approx\sum\limits_{\tau_{1},\dots,\tau_{n}}\sum\limits_{t=0}^{T}\frac{\partial\log\pi_{t}(\bm{a}_{t}|\bm{W}_{t},\bm{\theta})}{\partial\bm{\theta}}r_{\tau_{i}}(\bm{a}_{t},\bm{s}_{t})\end{split} (7)

Additionally, the memory cost of storing synaptic weights trajectories instead of hidden activity in a network of N neurons is \mathcal{O}(N^{2}) instead of \mathcal{O}(N) . This might lead to prohibitively large memory requirements for training with BPTT [ 83 ] over long episodes. We present in S.I an alternative solution to train the model through the discrete adjoint sensitivity method, leveraging the work of [ 84 ] yielding a memory cost of \mathcal{O}(1) . The agent’s log-policy total derivative with respect to \bm{\theta} can be computed as the solution of an augmented adjoint problem [ 67 ] .
Gradient policy update
We define the evolution of the agent policy in task \tau up to step T as the stochastic policy process (\pi_{t})_{t\leq T} in the space of policy \Pi with measure \mu^{\bm{\theta},\tau,1\cdots T}_{\pi} and write (\bm{W}_{t})_{t\leq T} the trajectory of weights, such that:

\big{(}\pi_{t}\big{)}_{t\leq T}=\big{(}\pi_{t}(\cdot|\bm{W}_{t},\bm{s}_{t},\bm{\theta})\big{)}_{t\leq T}\sim\mu^{\bm{\theta},\tau,1\cdots T}_{\pi} (8)

We also recall the definition of \mathcal{R}(\tau,\pi) as the average accumulated reward under the realisation of (\pi_{t})_{t\leq T} for task \tau . Then, The policy gradient update used to train MetODS can be written as the average gradient of \bm{\theta} with respect to \mu_{\mathbb{T}} and \mu^{\bm{\theta},\tau,1\cdots T}_{\pi} :

\begin{split}\frac{\partial}{\partial\bm{\theta}}\;\mathbb{E}_{\tau\sim\mu_{\mathbb{T}}}\bigg{[}\mathbb{E}_{\pi\sim\mu^{\bm{\theta},\tau,1\cdots T}_{\pi}}\big{[}\mathcal{R}(\tau,\pi)]\bigg{]}=\mathbb{E}_{\tau\sim\mu_{\mathbb{T}}}\bigg{[}\sum\limits_{t=0}^{T}\frac{\partial}{\partial\bm{\theta}}\mathbb{E}_{\pi_{t},\mathcal{P}_{\tau}}\big{[}r_{\tau}(\bm{a}_{t},\bm{s}_{t})\big{]}\bigg{]}\end{split} (9)

Then for any t\leq T , we can rewrite the gradient of the average, using the log-policy trick:

\begin{split}&\hskip 14.45377pt\frac{\partial}{\partial\bm{\theta}}\mathbb{E}_{\pi_{t},\mathcal{P}_{\tau}}\big{[}r_{\tau}(\bm{a}_{t},\bm{s}_{t})\big{]}=\mathbb{E}_{\pi_{t},\mathcal{P}_{\tau}}\big{[}r_{\tau}(\bm{a}_{t},\bm{s}_{t})\frac{\partial\log\pi_{t}(\bm{a}_{t}^{\tau}|\bm{W}_{t},\bm{\theta})}{\partial\bm{\theta}}\big{]}\end{split} (10)

As specified in section 3.3, we do not sample over the policy distribution \pi_{t} and probability transition \mathcal{P}_{\tau} to estimate the inner expectation in ( 11 ). Instead, we rely on a single evaluation, which yield, combining ( 9 ) with ( 10 ):

\frac{\partial}{\partial\bm{\theta}}\;\mathbb{E}_{\tau\sim\mu_{\mathbb{T}}}\bigg{[}\mathbb{E}_{\pi\sim\mu^{\bm{\theta},\tau,1\cdots T}_{\pi}}\big{[}\mathcal{R}(\tau,\pi)]\bigg{]}\hskip 7.22743pt\approx\hskip 7.22743pt\sum\limits_{t=0}^{T}\big{[}r_{\tau}(\bm{a}_{t},\bm{s}_{t})\frac{\partial\log\pi_{t}(\bm{a}_{t}^{\tau}|\bm{W}_{t},\bm{\theta})}{\partial\bm{\theta}}\big{]} (11)

By sampling over tasks, this last equation allows us to write the following gradient estimator

\frac{\partial}{\partial\bm{\bm{\theta}}}\mathbb{E}_{\tau\sim\mu_{\mathbb{T}}}\bigg{[}\mathbb{E}_{\pi\sim\mu^{\bm{\theta},\tau,1\cdots T}_{\pi}}\big{[}\mathcal{R}(\tau,\pi)]\bigg{]}\hskip 7.22743pt\approx\hskip 7.22743pt\frac{1}{M}\sum\limits_{i=0}^{M}\sum\limits_{t=0}^{T}\frac{\partial\log\pi_{t}(\bm{a}^{\tau_{i}}_{t}|\bm{W}_{t},\bm{\theta})}{\partial\bm{\theta}}r(\bm{a}^{\tau_{i}}_{t},\bm{s}^{\tau_{i}}_{t}) (12)
Discrete adjoint system
With the previous notation, let us define the total gradient function \frac{\partial\mathcal{J}}{\partial\bm{\theta}} as the gradient of our objective function:

\frac{\partial\mathcal{J}}{\partial\bm{\theta}}(\bm{\theta})=\frac{1}{M}\sum\limits_{t=0}^{T}\sum\limits_{i=0}^{M}\frac{\partial\log\pi_{t}(\bm{a}^{\tau_{i}}_{t}|\bm{W}_{t},\bm{\theta})}{\partial\bm{\theta}}r(\bm{a}^{\tau_{i}}_{t},\bm{s}^{\tau_{i}}_{t}) (13)

To clarify, we shall introduce the intermediary cost notation:

c_{t}(\bm{W},\bm{\theta},t)=\sum\limits_{i=0}^{M}\log\pi_{t}(\bm{a}^{\tau_{i}}_{t}|\bm{W}_{t},\bm{\theta})r(\bm{a}^{\tau_{i}}_{t},\bm{s}^{\tau_{i}}_{t}) (14)

and we note the following update equation

\bm{W}_{t+1}=\delta(\bm{W}_{t},\bm{\theta}) (15)

This identify a discrete dynamical system with finite sum and differentiable cost, whose gradient can be computed mediating the introduction of an adjoint dynamical system presented in section 2 of [ 84 ] . Defining (\bm{\nu}_{t})_{t\leq T} the adjoint sequence, the general gradient equation can be computed as:

\frac{\partial\mathcal{J}}{\partial\bm{\theta}}(\bm{\theta})=\bigg{[}\frac{\partial c_{0}}{\partial\bm{W}_{0}}-\bm{ u}_{0}\cdot\frac{\partial\delta}{\partial\bm{W}_{0}}-\bm{ u}_{0}\bigg{]}^{\dagger}\frac{\partial\bm{W}_{0}}{\partial\bm{\theta}}+\sum\limits_{i=1}^{T}\bigg{[}\frac{\partial c_{t}}{\partial\bm{\theta}}-\bm{ u}_{t}^{\dagger}\cdot\frac{\partial\delta}{\partial\bm{\theta}}\bigg{]} (16)

Applying formula ( 16 ) to MetODS yields the following gradient formula:

\begin{split}\frac{\partial\mathcal{J}}{\partial\bm{\theta}}(\bm{\theta})=&\bigg{[}\frac{\partial c_{0}}{\partial\bm{W}_{0}}-\bm{ u}_{0}\bigg{]}+\sum\limits_{i=1}^{T}\bigg{[}\frac{\partial c_{t}}{\partial\bm{\theta}}-\bm{ u}_{t}^{\dagger}\cdot\frac{\partial\delta}{\partial\bm{\theta}}\bigg{]}\end{split} (17)

where (\bm{\nu}_{t})_{t\leq T} follows the following update rule backwards:

\begin{cases}\bm{ u}_{T}=0\ \bm{ u}_{t-1}=\bm{ u}_{t}-\frac{\partial c_{t}}{\partial\bm{W}_{t}}\end{cases} (18)
Harlow task
This experiment consists in a 1-dimensional simplification of the task presented in [ 54 ] and inspired from https://github.com/bkhmsi/Meta-RL-Harlow . The action space is the discrete set \{-1,1\} moving the agent respectively to the left/right on a discretized line. The state space consist in 17 positions while the agent receptive field is eight dimensional. Values are placed at 3 positions from the fixation target position. The fixation position yields a reward of 0.2 while the values are drawn uniformly from [\![0,10]\!] and randomly associated with a reward of -1 and 1 at the beginning of each episode. The maximal duration for an episode is 250 steps, although the model solves the 5 trials in \sim 35 steps on average.
Gym Mujoco directional robot control
Setting: We consider the directional rewards task proposed in [ 61 ] with the Ant and Cheetah robots, as a more complex test of rapid adaptation. We apply standard RL practice for continuous control by parameterizing the stochastic policy as a product of independent normal distributions with fixed standard deviation \sigma=0.1 and mean inferred by the agent network. A training meta-episode consists in a single rollout of 200 steps with random sampling of the reward direction for each episode.

Robot impairment: We partially impaired the agent motor capabilities by "freezing" one of the robot actuator. Namely we consistently passed a value of zero to the the right_back_leg in Ant (coordinate 8 in OpenAI Gym XML asset files) and ffoot in Cheetah (coordinate 6 in OpenAI Gym XML asset files).

Continual adaptation: Finally, we further investigated whether the adaptation mechanism was also adjustable after performing initial adaptation by testing trained agents over 5 consecutive episodes with randomly changing rewarded direction while retaining weight state across episodes. We show in Fig 10 the average reward profile of the three meta-RL agents over 1000 test episodes, where MetODS adapt faster and better with respect to reward variations.
Maze Navigation task
Maze generation: The maze environments are created following Prim’s algorithm [ 80 ] which randomly propagates walls on a board of N\times N cells. We additionally add walls to cell location were no walls has been created at any of the neighboring 8 cells to avoid null inputs to the agent. The reward and agent locations are selected at random at the beginning of an episode. The reward location does not change during the episode while the agent is restarting from a uniform random location after every encounter with reward.

Models capacity: In this experiment, since model memory capacity seems to directly impact final performance, we specifically control for the number of learnable parameters in order to compare models and fix a training budget of 10M env. steps. Since MAML does not have a continual adaptation mechanism, we perform gradient adaptation every 20 time-steps during an episode, in order to balance noisiness of the gradient and rapidity of exploration. We show in Table 2 the mean over 1000 episodes of the accumulated reward of agents trained on 8\times 8 mazes and tested on different maze sizes N\in(4,6,8,10,12) . Agents performances are highly varying within a size setting due to differences in maze configurations and across maze sizes due to increasing complexity. However, MetODS and RL 2 are able to generalize at least partially to these new settings (see table 2 ). MetODS outperforms RL 2 in every setting, and generalizes better, retaining an advantage of up to 25\% in accumulated reward in the biggest maze.

Maze size MAML RL 2 MetODS Rel imp. to RL 2 6 21.0\pm 18.4 149.3\pm 66.7 169.1 \pm 66.1 13 % 8 14.9\pm 4.5 72.1\pm 45.6 87.3 \pm 48.3 20\% 10 5.7\pm 7.9 28.1\pm 29.7 34.9 \pm 34.9 21\% 12 3.9\pm 6.9 11.1\pm 15.8 13.9 \pm 19.8 25\%

Additionally, we investigated the selectivity of neurons of MetODS plastic weight layer with respect to spatial location, to investigate whether there could be emergence of activation patterns resembling those of grid cells found in the mammalian entorhinal cortex. We measured the selectivity of each neuron for specific agent locations by measuring the average activation rate normalized by the number of agent passages in different maze configurations. Interestingly, we found that without any explicit regularization, some neurons displayed sparse activation and consistent remapping between maps.