Add meta learning framework #183

jieyibi · 2024-05-27T15:57:43Z

Description

Add a new training framework based on meta learning. Details refer to Zhou et al. 2023.

Motivation and Context

To address the generalization issue.

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds core functionality)
Breaking change (fix or feature that would cause existing functionality to change)
Documentation (update in the documentation)
Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.

fedebotu

Great! 🚀

Left some comments here and there~

fedebotu · 2024-05-28T03:38:29Z

rl4co/utils/meta_trainer.py

+    def _alpha_scheduler(self):
+        self.alpha = max(self.alpha * self.alpha_decay, 0.0001)
+
+class RL4COMetaTrainer(Trainer):


Are there any differences compared to the RL4COTrainer? I could not find any at a first glance. I guess the only difference is that we pass the MetaModelCallback?

fedebotu · 2024-05-28T03:40:19Z

rl4co/utils/meta_trainer.py

+log = utils.get_pylogger(__name__)
+
+
+class MetaModelCallback(Callback):


I haven't thought of Meta-learning as Lightning Callbacks, but it looks neat! :D

Tagging @Junyoungpark for the good ol' Reptile

fedebotu · 2024-05-28T03:48:09Z

rl4co/utils/meta_trainer.py

+    def __init__(self, meta_params, print_log=True):
+        super().__init__()
+        self.meta_params = meta_params
+        assert meta_params["meta_method"] == 'reptile', NotImplementedError


Another general comment: it seems that this callback is only for Reptile, so we should consider calling it ReptileCallback. I think it would be cool to have these as say a meta_learning/ folder

Yes. ReptileCallback is better. Where should this new meta_learning/ folder be located? In utils/?

I''d say under models/rl (which includes LightningModules), since meta_learning is a way to optimize a policy

fedebotu · 2024-05-28T03:49:55Z

rl4co/utils/meta_trainer.py

+        self._sample_task()
+        self.selected_tasks[0] = (pl_module.env.generator.num_loc, )
+
+    def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:


[Minor] type hints should be classes, not strings. For example Trainer and LightningModule (Maybe even better: RL4COTrainer and RL4COLitModule since everything is inherited from that)

Hi. I am refactoring the ReptileCallback inherited from the REINFORCE and the RL4COLitModule. But in this case, maybe I need to add new meta_model.py inherited the model.py under the zoo/pomo/ folder or the zoo/am/ folder to call the REPTILE outside. It is a little bit redundant... Maybe a Lightning Callback is more generic. We could apply it to every model and every policy.

cbhua

Great job! 🚀 Really nice to have meta learning supporting 😁

cbhua · 2024-05-28T04:27:51Z

examples/2d-meta_train.py

@@ -0,0 +1,87 @@
+import pytz


[Minor] Nice tool but pytz is not included in RL4CO's dependence packages. Maybe better to have a package check here:

try: import pytz except ImportError: # raise a warning and use python default timeit pass

Yeah, maybe it could be removed.

cbhua · 2024-05-28T04:50:06Z

rl4co/utils/meta_trainer.py

+        super().__init__()
+        self.meta_params = meta_params
+        assert meta_params["meta_method"] == 'reptile', NotImplementedError
+        assert meta_params["data_type"] == 'size', NotImplementedError


It seems that this parameter data_type is not called anywhere? 🤔

fedebotu

Great!
Also:

Could you reproduce the learning curves of the original model?
Could you make a simple test as this for your model so that it can be automatically checked?

fedebotu · 2024-05-28T09:56:26Z

rl4co/utils/meta_trainer.py

+
+    # Meta training framework for addressing the generalization issue
+    # Based on Zhou et al. (2023): https://arxiv.org/abs/2305.19587
+    def __init__(self, meta_params, print_log=True):


[Minor] I think it would be more clear to list the Hyperparameters.

For example

def __init__(self, alpha=..., alpha_decay=... ... print_log=True):

which is generally easier to maintain and to document.

fedebotu · 2024-05-28T10:00:40Z

It seems there was a mistake here - I noticed that the "tasks" are a bit hardcoded, i.e. the "size" and "capacity" (I guess for TSP and CVRP, right?)

Speaking of testing: you can make sure things work on your device before committing by using pytest tests

cbhua · 2024-05-29T06:40:38Z

rl4co/utils/meta_trainer.py

+class ReptileCallback(Callback):
+
+    # Meta training framework for addressing the generalization issue
+    # Based on Zhou et al. (2023): https://arxiv.org/abs/2305.19587
+    def __init__(self,
+                 num_tasks,
+                 alpha,
+                 alpha_decay,
+                 min_size,
+                 max_size,
+                 sch_bar = 0.9,
+                 data_type = "size",
+                 print_log=True):
+        super().__init__()


[Documentation] It's recommended to have a doc with parameters, possibly including the data type, constraints, hints, etc. Better for us "non-experts" to understand 😆

class ReptileCallback(Callback): """Meta training framework for addressing the generalization issue Based on Zhou et al. (2023): https://arxiv.org/abs/2305.19587 Args: - num_tasks: number of task types, i.e. `B` in the original paper - alpha: ... - ... """ def __init__( self, num_tasks: int, alpha: float, alpha_decay: float, min_size: int, max_size: int, sch_bar: float = 0.9, data_type: str = "size", print_log: bool = True, ): super().__init__()

Hi Chuanbo. I have added the documentation as you recommended, along with the generation code for some distributions defined in the generalization-related works. Now the meta learning framework is supported for cross-distribution generalization.

fedebotu · 2024-05-29T16:24:49Z

rl4co/envs/common/distribution_utils.py

@@ -0,0 +1,184 @@
+import torch
+
+class Cluster():


I like this! @cbhua I think we should train some model with say TSP50 / CVRP50 with a mixed distribution and test its generalization performance

Minor comment: Shouldn't this be a subclass of torch.distributions.distribution.Distribution?

Yes actually this class is what we are missing in the distribution! This could be used by various environments' generator.

About the experiment, we want to test the distribution generalization ability right?

examples/2d-meta_train.py

fedebotu · 2024-06-08T13:26:52Z

Routine check, how is progress going? I think the multi-distribution generators in particular should be included, since they are part of MDPOMO, right?

jieyibi · 2024-06-11T08:23:34Z

The updated commit supports training on multiple mixed distributions by changing the argument loc_distribution to "mix_distribution" in the generator_params of the environment, e.g.

env = CVRPEnv(generator_params={'loc_distribution': "mix_distribution"})

The rest of the arguments remain the same.

Note that this mixed distribution setting follows the # setting in Bi et al., 2022.

fedebotu · 2024-06-11T10:53:17Z

Great! Should the code be merged?

Change some parameters for performance

jieyibi · 2024-06-19T13:23:41Z

Yes. I think it is ready to be merged. The newly reproduced performance (by changing some key parameters) seems similar to that in the literature.

fedebotu · 2024-06-19T16:02:57Z

Great!

How about the generalization experiments, i.e., MDPOMO? Do you have yaml configuration for them like this, or a training script?

jieyibi · 2024-06-19T17:29:56Z

Yeah, already added to the main branch:)

fedebotu · 2024-06-20T02:24:37Z

Awesome! Then we can go ahead and merge :)

jieyibi added 3 commits May 27, 2024 23:08

Update meta trainer

475e430

Update meta trainer

5cbe7b9

Add source

05e0870

fedebotu reviewed May 28, 2024

View reviewed changes

cbhua approved these changes May 28, 2024

View reviewed changes

Update Reptile callbacks

e8435c6

fedebotu reviewed May 28, 2024

View reviewed changes

Add test

bff31a9

fedebotu reviewed May 28, 2024

View reviewed changes

Update meta learning framework

0f4032c

cbhua reviewed May 29, 2024

View reviewed changes

jieyibi added 2 commits May 29, 2024 21:14

Add support for cross-distribution generalization

73d7a65

add document

5a029fa

fedebotu reviewed May 29, 2024

View reviewed changes

examples/2d-meta_train.py Outdated Show resolved Hide resolved

add support for training on multiple mixed distributions

d11788d

Update 2d-meta_train.py

60fa8c8

Change some parameters for performance

fedebotu added this to the 0.5.0 milestone Jun 19, 2024

fedebotu approved these changes Jun 19, 2024

View reviewed changes

fedebotu merged commit 911a754 into ai4co:main Jun 20, 2024
9 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add meta learning framework #183

Add meta learning framework #183

jieyibi commented May 27, 2024

fedebotu left a comment

fedebotu May 28, 2024

fedebotu May 28, 2024

fedebotu May 28, 2024

fedebotu May 28, 2024

jieyibi May 28, 2024

fedebotu May 28, 2024

fedebotu May 28, 2024

jieyibi May 28, 2024 •

edited

Loading

cbhua left a comment

cbhua May 28, 2024

jieyibi May 28, 2024

cbhua May 28, 2024

fedebotu left a comment •

edited

Loading

fedebotu May 28, 2024

fedebotu commented May 28, 2024

cbhua May 29, 2024

jieyibi May 29, 2024

fedebotu May 29, 2024

cbhua May 29, 2024

fedebotu commented Jun 8, 2024

jieyibi commented Jun 11, 2024

fedebotu commented Jun 11, 2024

jieyibi commented Jun 19, 2024 •

edited

Loading

fedebotu commented Jun 19, 2024

jieyibi commented Jun 19, 2024

fedebotu commented Jun 20, 2024

		log = utils.get_pylogger(__name__)


		class MetaModelCallback(Callback):

Add meta learning framework #183

Add meta learning framework #183

Conversation

jieyibi commented May 27, 2024

Description

Motivation and Context

Types of changes

Checklist

fedebotu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jieyibi May 28, 2024 • edited Loading

Choose a reason for hiding this comment

cbhua left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fedebotu left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fedebotu commented May 28, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fedebotu commented Jun 8, 2024

jieyibi commented Jun 11, 2024

fedebotu commented Jun 11, 2024

jieyibi commented Jun 19, 2024 • edited Loading

fedebotu commented Jun 19, 2024

jieyibi commented Jun 19, 2024

fedebotu commented Jun 20, 2024

jieyibi May 28, 2024 •

edited

Loading

fedebotu left a comment •

edited

Loading

jieyibi commented Jun 19, 2024 •

edited

Loading