wehub-resource-sync
444 lines
16 KiB
Python
444 lines
16 KiB
Python
"""Example of how to use `TorchMetaLearner` and `DifferentiableLearner` for MAML.
|
||
|
||
Meta-learning, or “learning to learn,” trains models to quickly adapt to new tasks
|
||
using only a few examples. One prominent method is Model-Agnostic Meta-Learning
|
||
(MAML), which is compatible with any model trained via gradient descent. MAML has
|
||
been successfully applied across domains such as classification, regression, and
|
||
reinforcement learning.
|
||
|
||
In this MAML example, the goal is to train a model that can adapt to an infinite
|
||
number of tasks, where each task corresponds to a sinusoidal function with randomly
|
||
sampled amplitude and phase. Because each new task introduces a shift in data
|
||
distribution, traditional learning algorithms would fail to generalize — they’d
|
||
overfit to the training task and struggle on unseen ones. Meta-learning addresses
|
||
this by optimizing the model parameters such that they can be fine-tuned rapidly
|
||
for any new task.
|
||
|
||
During training, a DifferentiableLearner performs an inner-loop update using the
|
||
training error for each task. The outer-loop TorchMetaLearner then evaluates the
|
||
model’s performance on held-out data (the task's test set) and updates the meta-
|
||
parameters so that they lead to better generalization across all tasks. This bi-
|
||
level optimization ensures that gradients across tasks remain close, enabling
|
||
fast adaptation.
|
||
|
||
At inference time, the trained model can adapt to a new task using just a small
|
||
batch of examples — performing few-shot learning to adjust quickly and accurately.
|
||
|
||
This example shows:
|
||
- how to implement MAML with RLlib in just a few lines of code.
|
||
- how to define a `TorchDifferentiableLearner` to register a custom train loss
|
||
function.
|
||
- how to define a `TorchMetaLearner` class to implement a custom meta (test) train
|
||
loss function.
|
||
- how to configure both learners top be used with each others via the
|
||
`DifferentiableAlgorithmConfig` and `DifferentiableLearnerConfig`.
|
||
- how to update the `RLModule` in a meta-learning fashion.
|
||
- how to fine-tune an `RLModule` with gradient descent within a few iterations with
|
||
only using the meta (test) loss.
|
||
|
||
See :py:class:`~ray.rllib.examples.learners.classes.lr_meta_learner.LRTorchMetaLearner` # noqa
|
||
class for details on how to override the main `TorchMetaLearner`. And see
|
||
:py:class:`~ray.rllib.examples.learners.classes.lr_differentiable_learner.LRTorchDifferentiableLearner` # noqa
|
||
class for an example of how to override the main `TorchDifferentiableLearner`.
|
||
|
||
Note, the meta-learner needs a long-enough training (`default_iters`=~70,000) to learn
|
||
to adapt quickly to new tasks.
|
||
|
||
How to run this script
|
||
----------------------
|
||
`python [script file name].py --iters=70000 --meta-train-batch-size=5 --fine-tune-batch-size=5`
|
||
|
||
Use the `--meta-train-batch-size` to set the training/testing batch size in meta-learning and
|
||
the `--fine-tune-batch-size` to adjust the number of samples used in all updates during
|
||
few-shot learning.
|
||
|
||
To suppress plotting (plotting is the default) use `--no-plot` and for taking a longer
|
||
look at the plot increase the seconds for which plotting is paused at the end of the
|
||
script by `--pause-plot-secs`.
|
||
|
||
Results to expect
|
||
-----------------
|
||
You should expect to see sometimes alternating test losses ("Total Loss") due to new
|
||
(unseen) tasks during meta learning. In few-shot learning after the meta-learning the
|
||
(few shot) loss should decrease almost monotonically. In the plot you can expect to see
|
||
a decent adaption to the new task after fine-tuning updates of the `RLModule` weights.
|
||
|
||
With `--iters=70_000`, `--meta-train-batch-size=5`, `--fine-tune-batch-size=5`,
|
||
`--fine-tune-lr=0.01`, `--fine-tune-iters=10`, `--meta-lr=0.001`, `--noise-std=0.0`,
|
||
and no seed defined.
|
||
-------------------------
|
||
|
||
Iteration: 68000
|
||
Total loss: 0.013758559711277485
|
||
-------------------------
|
||
|
||
Iteration: 69000
|
||
Total loss: 0.7246640920639038
|
||
-------------------------
|
||
|
||
Iteration: 70000
|
||
Total loss: 3.091259002685547
|
||
|
||
Few shot loss: 2.754437208175659
|
||
Few shot loss: 2.7399725914001465
|
||
Few shot loss: 2.499554395675659
|
||
Few shot loss: 2.1763901710510254
|
||
Few shot loss: 1.793503999710083
|
||
Few shot loss: 1.4362313747406006
|
||
Few shot loss: 1.083552598953247
|
||
Few shot loss: 0.7845061421394348
|
||
Few shot loss: 0.5579453110694885
|
||
Few shot loss: 0.4087105393409729
|
||
"""
|
||
import gymnasium as gym
|
||
import matplotlib.pyplot as plt
|
||
import numpy as np
|
||
|
||
from ray.rllib.algorithms.algorithm_config import DifferentiableAlgorithmConfig
|
||
from ray.rllib.core import DEFAULT_MODULE_ID
|
||
from ray.rllib.core.columns import Columns
|
||
from ray.rllib.core.learner.differentiable_learner_config import (
|
||
DifferentiableLearnerConfig,
|
||
)
|
||
from ray.rllib.core.learner.training_data import TrainingData
|
||
from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
|
||
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
|
||
from ray.rllib.examples.algorithms.classes.maml_lr_differentiable_learner import (
|
||
MAMLTorchDifferentiableLearner,
|
||
)
|
||
from ray.rllib.examples.algorithms.classes.maml_lr_differentiable_rlm import (
|
||
DifferentiableTorchRLModule,
|
||
)
|
||
from ray.rllib.examples.algorithms.classes.maml_lr_meta_learner import (
|
||
MAMLTorchMetaLearner,
|
||
)
|
||
from ray.rllib.examples.utils import add_rllib_example_script_args
|
||
from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch
|
||
from ray.rllib.utils.framework import try_import_torch
|
||
|
||
# Import torch.
|
||
torch, _ = try_import_torch()
|
||
|
||
# Implement generation of data from sinusoid curves.
|
||
def generate_sinusoid_task(batch_size, noise_std=0.1, return_params=False):
|
||
"""Generate a sinusoid task with random amplitude and phase.
|
||
|
||
Args:
|
||
batch_size: The number of data points to be generated.
|
||
noise_std: An optional standard deviation to be used in the sinusoid
|
||
data generation. Defines a linear error term added to the sine
|
||
curve.
|
||
return_params: If the sampled amplitude and phase should be returned.
|
||
|
||
Returns:
|
||
Torch tensors with the support data and the labels of a sinusoid
|
||
curve.
|
||
"""
|
||
# Sample the amplitude and the phase for a task.
|
||
amplitude = np.random.uniform(0.1, 5.0)
|
||
phase = np.random.uniform(0.0, np.pi)
|
||
|
||
# Sample the support.
|
||
x = np.random.uniform(-5.0, 5.0, (batch_size, 1))
|
||
|
||
# Generate the labels.
|
||
y = amplitude * np.sin(x - phase)
|
||
|
||
# Add noise.
|
||
y += noise_std * np.random.random((batch_size, 1))
|
||
|
||
# If sampled parameters should be returned.
|
||
if return_params:
|
||
# Return torch tensors.
|
||
return (
|
||
torch.tensor(x, dtype=torch.float32),
|
||
torch.tensor(y, dtype=torch.float32),
|
||
amplitude,
|
||
phase,
|
||
)
|
||
# Otherwise, return only the sampled data.
|
||
else:
|
||
return (
|
||
torch.tensor(x, dtype=torch.float32),
|
||
torch.tensor(y, dtype=torch.float32),
|
||
)
|
||
|
||
|
||
def sample_task(batch_size=10, noise_std=0.1, training_data=False, return_params=False):
|
||
"""Samples training batches for meta learner and differentiable learner.
|
||
|
||
Args:
|
||
batch_size: The batch size for both meta learning and task learning.
|
||
noise_std: An optional standard deviation to be used in the sinusoid
|
||
data generation. Defines a linear error term added to the sine
|
||
curve.
|
||
training_data: Whether data should be returned as `TrainingData`.
|
||
Otherwise, a `MultiAgentBatch` is returned. Default is `False`.
|
||
return_params: If the sampled amplitude and phase should be returned.
|
||
|
||
Returns:
|
||
A tuple with training batches for the meta learner and the differentiable
|
||
learner. If `training_data` is `True`, the data is wrapped into
|
||
`TrainingData`, otherwise both batches are `MultiAgentBatch`es.
|
||
|
||
"""
|
||
# Generate training data for meta learner and differentiable learner.
|
||
train_batch = {}
|
||
generated_data = generate_sinusoid_task(
|
||
batch_size * 2, noise_std=noise_std, return_params=return_params
|
||
)
|
||
train_batch[Columns.OBS], train_batch["y"] = generated_data[:2]
|
||
|
||
# Convert to `MultiAgentBatch`.
|
||
meta_train_batch = MultiAgentBatch(
|
||
env_steps=batch_size,
|
||
policy_batches={
|
||
DEFAULT_MODULE_ID: SampleBatch(
|
||
{k: train_batch[k][:batch_size] for k in train_batch}
|
||
)
|
||
},
|
||
)
|
||
task_train_batch = MultiAgentBatch(
|
||
env_steps=batch_size,
|
||
policy_batches={
|
||
DEFAULT_MODULE_ID: SampleBatch(
|
||
{k: train_batch[k][batch_size:] for k in train_batch}
|
||
)
|
||
},
|
||
)
|
||
# If necessary convert to `TrainingData`.
|
||
if training_data:
|
||
meta_train_batch = TrainingData(
|
||
batch=meta_train_batch,
|
||
)
|
||
task_train_batch = TrainingData(
|
||
batch=task_train_batch,
|
||
)
|
||
|
||
# If amplitude and phase should be returned add them to the return tuple.
|
||
if return_params:
|
||
return meta_train_batch, task_train_batch, *generated_data[2:]
|
||
# Otherwise return solely train data.
|
||
else:
|
||
return meta_train_batch, task_train_batch
|
||
|
||
|
||
# Define arguments.
|
||
parser = add_rllib_example_script_args(default_iters=70_000)
|
||
|
||
parser.add_argument(
|
||
"--meta-train-batch-size",
|
||
type=int,
|
||
default=5,
|
||
help="The number of samples per train and test update (meta-learning).",
|
||
)
|
||
parser.add_argument(
|
||
"--meta-lr",
|
||
type=float,
|
||
default=0.001,
|
||
help="The learning rate to be used for meta learning (in the `MetaLearner`).",
|
||
)
|
||
parser.add_argument(
|
||
"--fine-tune-batch-size",
|
||
type=int,
|
||
default=10,
|
||
help="The number of samples for the fine-tuning updates.",
|
||
)
|
||
parser.add_argument(
|
||
"--noise-std",
|
||
type=float,
|
||
default=0.0,
|
||
help="The standard deviation for noise added to the single tasks.",
|
||
)
|
||
parser.add_argument(
|
||
"--seed",
|
||
type=int,
|
||
default=None,
|
||
help="An optional random seed. If not set, the experiment is not reproducable.",
|
||
)
|
||
parser.add_argument(
|
||
"--fine-tune-iters",
|
||
type=int,
|
||
default=10,
|
||
help="The number of updates in fine-tuning.",
|
||
)
|
||
parser.add_argument(
|
||
"--fine-tune-lr",
|
||
type=float,
|
||
default=0.01,
|
||
help="The learning rate to be used in fine-tuning the model in the test phase.",
|
||
)
|
||
parser.add_argument(
|
||
"--no-plot",
|
||
action="store_true",
|
||
help=(
|
||
"If plotting should suppressed. Otherwise user action is needed to close "
|
||
"the plot early."
|
||
),
|
||
)
|
||
parser.add_argument(
|
||
"--pause-plot-secs",
|
||
type=int,
|
||
default=1000,
|
||
help=(
|
||
"The number of seconds to keep the plot open. Note the plot can always be "
|
||
"closed by the user when open."
|
||
),
|
||
)
|
||
|
||
# Parse the arguments.
|
||
args = parser.parse_args()
|
||
|
||
# If a random seed is provided set it for torch and numpy.
|
||
if args.seed:
|
||
torch.random.manual_seed(args.seed)
|
||
np.random.seed(args.seed)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# Define the `RLModule`.
|
||
module_spec = RLModuleSpec(
|
||
module_class=DifferentiableTorchRLModule,
|
||
# Note, the spaces are needed by default but are not used.
|
||
observation_space=gym.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32),
|
||
action_space=gym.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32),
|
||
)
|
||
# `Learner`s work on `MultiRLModule`s.
|
||
multi_module_spec = MultiRLModuleSpec(
|
||
rl_module_specs={DEFAULT_MODULE_ID: module_spec}
|
||
)
|
||
|
||
# Build the `MultiRLModule`.
|
||
module = multi_module_spec.build()
|
||
|
||
# Configure the `DifferentiableLearner`.
|
||
diff_learner_config = DifferentiableLearnerConfig(
|
||
learner_class=MAMLTorchDifferentiableLearner,
|
||
minibatch_size=args.meta_train_batch_size,
|
||
lr=0.01,
|
||
)
|
||
|
||
# Configure the `TorchMetaLearner` via the `DifferentiableAlgorithmConfig`.
|
||
config = (
|
||
DifferentiableAlgorithmConfig()
|
||
.learners(
|
||
# Add the `DifferentiableLearnerConfig`s.
|
||
differentiable_learner_configs=[diff_learner_config],
|
||
num_gpus_per_learner=args.num_gpus_per_learner or 0,
|
||
)
|
||
.training(
|
||
lr=args.meta_lr,
|
||
train_batch_size=args.meta_train_batch_size,
|
||
# Use the full batch in a single update.
|
||
minibatch_size=args.meta_train_batch_size,
|
||
)
|
||
)
|
||
|
||
# Initialize the `TorchMetaLearner`.
|
||
meta_learner = MAMLTorchMetaLearner(config=config, module_spec=module_spec)
|
||
# Build the `TorchMetaLearner`.
|
||
meta_learner.build()
|
||
|
||
for i in range(args.stop_iters):
|
||
# Sample the training data.
|
||
meta_training_data, task_training_data = sample_task(
|
||
args.meta_train_batch_size, noise_std=args.noise_std, training_data=True
|
||
)
|
||
|
||
# Update the module.
|
||
outs = meta_learner.update(
|
||
training_data=meta_training_data,
|
||
num_epochs=1,
|
||
others_training_data=[task_training_data],
|
||
)
|
||
iter = i + 1
|
||
if iter % 1000 == 0:
|
||
total_loss = outs["default_policy"]["total_loss"].peek()
|
||
print("-------------------------\n")
|
||
print(f"Iteration: {iter}")
|
||
print(f"Total loss: {total_loss}")
|
||
|
||
# Generate test data.
|
||
test_batch, _, amplitude, phase = sample_task(
|
||
batch_size=args.fine_tune_batch_size,
|
||
noise_std=args.noise_std,
|
||
return_params=True,
|
||
)
|
||
|
||
if config.num_gpus_per_learner > 0:
|
||
test_batch = meta_learner._convert_batch_type(test_batch)
|
||
|
||
# Run inference and plot results.
|
||
with torch.no_grad():
|
||
# Generate a grid for the support.
|
||
x_grid = torch.tensor(
|
||
np.arange(-5.0, 5.0, 0.02), dtype=torch.float32, device=meta_learner._device
|
||
).view(-1, 1)
|
||
# Get label prediction from the model trained by MAML.
|
||
y_pred = meta_learner.module[DEFAULT_MODULE_ID]({Columns.OBS: x_grid})["y_pred"]
|
||
|
||
# Plot the results if requested.
|
||
if not args.no_plot:
|
||
# Sort the data by the support.
|
||
x_order = np.argsort(test_batch[DEFAULT_MODULE_ID][Columns.OBS].numpy()[:, 0])
|
||
x_sorted = test_batch[DEFAULT_MODULE_ID][Columns.OBS].numpy()[:, 0][x_order]
|
||
y_sorted = test_batch[DEFAULT_MODULE_ID]["y"][:, 0][x_order]
|
||
|
||
# Plot the data.
|
||
def sinusoid(t):
|
||
return amplitude * np.sin(t - phase)
|
||
|
||
plt.ion()
|
||
plt.figure(figsize=(5, 3))
|
||
# Plot the true sinusoid curve.
|
||
plt.plot(x_grid, sinusoid(x_grid), "r", label="Ground Truth")
|
||
# Add the sampled support values.
|
||
plt.plot(x_sorted, y_sorted, "^", color="purple")
|
||
# Add the prediction made by the model after MAML training.
|
||
plt.plot(x_grid, y_pred, ":", label="Prediction", color="#90EE90")
|
||
plt.title(f"MAML Results from {args.fine_tune_iters} fine-tuning steps.")
|
||
|
||
# Fine-tune with the meta loss for just a few steps.
|
||
optim = meta_learner.get_optimizers_for_module(DEFAULT_MODULE_ID)[0][1]
|
||
# Set the learning rate to a larger value.
|
||
for g in optim.param_groups:
|
||
g["lr"] = args.fine_tune_lr
|
||
# Now run the fine-tune iterations and update the model via the meta-learner loss.
|
||
for i in range(args.fine_tune_iters):
|
||
# Forward pass.
|
||
fwd_out = {
|
||
DEFAULT_MODULE_ID: meta_learner.module[DEFAULT_MODULE_ID](
|
||
test_batch[DEFAULT_MODULE_ID]
|
||
)
|
||
}
|
||
# Compute the MSE prediction loss.
|
||
loss_per_module = meta_learner.compute_losses(fwd_out=fwd_out, batch=test_batch)
|
||
# Optimize parameters.
|
||
optim.zero_grad(set_to_none=True)
|
||
loss_per_module[DEFAULT_MODULE_ID].backward()
|
||
optim.step()
|
||
# Show the loss for few-shot learning (fine-tuning).
|
||
print(f"Few shot loss: {loss_per_module[DEFAULT_MODULE_ID].item()}")
|
||
|
||
# Run the model again after fine-tuning.
|
||
with torch.no_grad():
|
||
y_pred_fine_tuned = meta_learner.module[DEFAULT_MODULE_ID](
|
||
{Columns.OBS: x_grid}
|
||
)["y_pred"]
|
||
|
||
if not args.no_plot:
|
||
# Plot the predictions of the fine-tuned model.
|
||
plt.plot(
|
||
x_grid,
|
||
y_pred_fine_tuned,
|
||
"-.",
|
||
label="Tuned Prediction",
|
||
color="green",
|
||
mfc="gray",
|
||
)
|
||
plt.legend()
|
||
plt.show()
|
||
|
||
# Pause the plot until the user closes it.
|
||
plt.pause(args.pause_plot_secs)
|