Meta Reinforcement Learning for steering tasks

Use case: AWAKE beamline at CERN

Implementation example for the RL4AA'24 workshop

Simon Hirlaender, Jan Kaiser, Chenran Xu, Andrea Santamaria Garcia

Today!

In this tutorial notebook we will implement all the basic components of a Meta Reinforcement Learning (MLRL) algorithm to solve steering task in a linear accelerator.

Getting started
Part I: Quick introduction
Part II: Running PPO on our problem
Part III: Running MAML on our problem

Getting started

You will need Python 3.9 or higher to run this code ❗
You will require about 1 GB of free disk space ❗
Start by cloning locally the repository of the tutorial:

git clone https://github.com/RL4AA/rl4aa24-tutorial.git

The accelerator problem we want to solve

The goal is to minimize the distance $\Delta x_i$ of an initial beam trajectory to a target trajectory at different points $i$ across the accelerator(here marked as "position") in the least amount of steps.

No description has been provided for this image

Reference

"Ultra fast reinforcement learning demonstrated at CERN AWAKE" - IPAC (2023)

Formulating the RL problem

The problem is formulated in an episodic manner.

Actions

The actuators are the strengths of 10 corrector magnets that can steer the beam. They are normalized to [-1, 1]. In this tutorial, we apply the action by adding a delta change

$\Delta a$ to the current magnet strengths.

Formulating the RL problem

States/Observations

The observations are the readings of ten beam position monitors (BPMs), which read the position of the beam at a particular point in the beamline. The states are also normalized to [-1,1], corresponding to

$\pm$ 100 mm in the real accelerator. No description has been provided for this image

Formulating the RL problem

Reward

The reward is the negative RMS value of the distance to the target trajectory.

$r(x) = - \sqrt{ \frac{1}{10} \sum_{i=1}^{10} \Delta x_{i}^2} \,, \ \ \ \Delta x_{i} = x_{i} - x^{\text{target}}_{i}$

where $x^{\text{target}}=\vec{0}$ for a centered trajectory.

Formulating the RL problem

Environments/Tasks

The environment dynamics are determined by the response matrix, which in linear systems can encapsulate the dynamics of the problem.

More specifically: given the response matrix is $\mathbf R$ , the change in actions $\Delta a$ (corrector magnet strength), and the change in states $\Delta s$ (BPM readings), we have:

$\begin{align} \Delta s &= \mathbf{R}\Delta a\\ \end{align}$

Defining a benchmark policy

During this tutorial we want to compare the trained policies we obtain with different methods to a benchmark policy.

For this problem, our benchmark policy is just the inverse of the environment's response matrix.

$\implies$ Actions from benchmark policy: $\begin{align} \Delta a &= \mathbf{R}^{-1}\Delta s \end{align}$

$\implies$ Actions from deep RL policy: With the policy we get the actions:

So how do we compare policies in the evaluation stage? 🧐

At the beginning of each episode we reset the environment to suboptimal corrector strengths in a random way.
For each step within the episode we use the inverse of the response matrix (benchmark) or the trained policy to compute the next action (forward passes) until the episode ends (convergence or termination).
This will be performed for different evaluation tasks, just to assess how the policy performs in different lattices.

Side note:

The benchmark policy will not immediately find the settings for the target trajectory, because the actions are limited so that the maximum step is within $[-1,1]$ in the normalized space.
We can then compare the metrics of both policies.

PPO agent settings 🧐

n_env = 1
n_steps = 2048 (default params)
n_epochs = 10 (default params)
buffer_size = n_steps x n_env = 2048
backprops = (total_timesteps / buffer_size) * n_epochs

Questions 💻

Consider total_timesteps = 100. This parameter specifies the total number of timesteps (or steps) that the training process should run across all environments.

$\implies$ Considering the PPO agent settings: will we fill the buffer? what do you expect that happens?

Questions 💻

Run the code in the terminal with python ppo.py --train --steps 100 and observe the plot that pops-up.

$\implies$ What is the difference in episode length between the benchmark policy and PPO?

$\implies$ Look at the cumulative episode length, which policy takes longer?

$\implies$ Compare both cumulative rewards, which reward is higher and why?

$\implies$ Look at the final reward (-10*RMS(BPM readings)) and consider the sucessful (in red) and unsuccessful termination conditions mentioned before. What can you say about how the episode was ended?

Questions 💻

Set total_timesteps to 50,000 this time. Run it in the terminal with python ppo.py --train --steps 50000

$\implies$ What are the main differences between the untrained and trained PPO policies?

Train a bit longer setting total_timesteps to 100,000. Run it in the terminal with python ppo.py --train --steps 100000

$\implies$ In how many steps does it converge compared to the other training steps?

Optimization-based meta RL in this tutorial

In this tutorial we will adapt the parameters of our model (policy) through gradient descent with the MAML algorithm.

We have a meta policy $\phi(\theta)$ , where $\theta$ are the weights of a neural network. The meta policy starts untrained $\phi_0$ .

Step 1: outer loop

We randomly sample a number of tasks $i$ (in our case $i\in \{1,\dots,8\}$ different lattices, called meta-batch-size in the code) from a task distribution, each one with its particular initial task policy $\varphi_{0}^i=\phi_0$ .

Step 2: inner loop (adaptation)

For each task, we gather experience for several episodes, store the experience, and use it to perform gradient descent and update the weights of each task policy $\varphi_{0}^i \rightarrow \varphi_{1}^i$

This is repeated for $k$ gradient descent steps to generate $\varphi_{k}^i$ .

Optimization-based meta RL in this tutorial

Step 3: outer loop (meta training)

We generate episodes with the adapted task policies $\varphi_{k}^i$ . We sum the losses calculated for each task $\tau_{i}$ and perform gradient descent on the meta policy $\phi_0 \rightarrow \phi_1$

$\beta$ is the meta learning rate, $\alpha$ is the fast learning rate (for inner loop gradient updates)

Meta RL: summary

We start with a random meta policy, and we initialize the task policies with it: $\phi_0 = \varphi_{0}^i$

1 meta_step:  # Outer loop
   sample 8 tasks
   for task in tasks:
        for fast_step in num_steps:  # Inner loop
            for fast_batch in fast_batch_size:
                rollout 1 episode:
                    reset corrector_strength
                    while not stopped:
                        env.step()

We have gathered experience and trained 8 task policies: $\varphi_{0}^1 \rightarrow \varphi_{k}^1$ $\vdots$ $\varphi_{0}^8 \rightarrow \varphi_{k}^8$

The losses from the task policies are summed, and gradient descent is applied to update the meta policy $\phi_0 \rightarrow \phi_1$

Evaluation of a random task policy 💻

We will look at the inner loop only.
We consider only 1 task for now (task 0), from 5 fixed evaluation tasks.
The policy $\varphi_0^0$ starts as random and adapts for 500 steps (and show the progress every 50 steps).
This code does the training and evaluation.

$\implies$ Run the following code to train the task policy $\varphi_0^0$ for 500 steps:

python test.py --experiment-name tutorial --experiment-type adapt_from_scratch --num-batches 500 --plot-interval 50 --task-ids 0

Once it has run, you can look at the adaptation progress by running:

python read_out_train.py --experiment-name tutorial --experiment-type adapt_from_scratch

$\implies$ Run it now for all tasks:

python test.py --experiment-name tutorial --experiment-type adapt_from_scratch --num-batches 500 --plot-interval 50 --task-ids 0 1 2 3 4

$\implies$ Save the plot for comparison later

Evaluation of the pre-trained meta-policy 💻

We will now use a pre-trained policy located in awake/pretrained_policy.th and evalulate it against a certain number of fixed tasks.

$\implies$ Run the following code:

python test.py --experiment-name tutorial --experiment-type test_meta --use-meta-policy --policy awake/pretrained_policy.th --num-batches 500 --plot-interval 50 --task-ids 0 1 2 3 4

use --task-ids 0 1 2 3 4 to run evaluation against all 5 tasks, or e.g. --task-ids 0 to evaluate only for task 0.
here we set the flag --use-meta-policy so that it uses the pre-trained policy.

$\implies$ Afterwards, you can look at the adaptation progress by running:

python read_out_train.py --experiment-name tutorial --experiment-type test_meta

Evaluation of the trained meta-policy

$\implies$ What difference can you see compared to the untrained policy (previous plot saved)?