add regenerative regularization and cite

README.md

@@ -151,3 +151,12 @@ actions = model.get_optimal_actions(video, instructions)
     year    = {2022}
 }
 ```
+
+```bibtex
+@inproceedings{Kumar2023MaintainingPI,
+    title   = {Maintaining Plasticity in Continual Learning via Regenerative Regularization},
+    author  = {Saurabh Kumar and Henrik Marklund and Benjamin Van Roy},
+    year    = {2023},
+    url     = {https://api.semanticscholar.org/CorpusID:261076021}
+}
+```

q_transformer/optimizer.py

@@ -1,4 +1,4 @@
-from torch.optim import AdamW, Adam
+from adam_atan2_pytorch import AdamAtan2
 
 def separate_weight_decayable_params(params):
     wd_params, no_wd_params = [], []
@@ -10,9 +10,9 @@ def separate_weight_decayable_params(params):
 def get_adam_optimizer(
     params,
     lr = 1e-4,
-    wd = 1e-2,
+    wd = 0,
     betas = (0.9, 0.99),
-    eps = 1e-8,
+    regen_reg_rate = 1e-2,
     filter_by_requires_grad = False,
     group_wd_params = True
 ):
@@ -29,7 +29,4 @@ def get_adam_optimizer(
             {'params': no_wd_params, 'weight_decay': 0},
         ]
 
-    if not has_wd:
-        return Adam(params, lr = lr, betas = betas, eps = eps)
-
-    return AdamW(params, lr = lr, weight_decay = wd, betas = betas, eps = eps)
+    return AdamAtan2(params, lr = lr, weight_decay = wd, betas = betas, regen_reg_rate = regen_reg_rate)

q_transformer/q_learner.py

@@ -93,6 +93,7 @@ def __init__(
         grad_accum_every: int = 1,
         monte_carlo_return: float | None = None,
         weight_decay: float = 0.,
+        regen_reg_rate: float = 1e-3,
         accelerator: Accelerator | None = None,
         accelerator_kwargs: dict = dict(),
         dataloader_kwargs: dict = dict(
@@ -143,6 +144,7 @@ def __init__(
             model.parameters(),
             lr = learning_rate,
             wd = weight_decay,
+            regen_reg_rate = regen_reg_rate,
             **optimizer_kwargs
         )
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'q-transformer',
   packages = find_packages(exclude=[]),
-  version = '0.1.17',
+  version = '0.2.0',
   license='MIT',
   description = 'Q-Transformer',
   author = 'Phil Wang',
@@ -19,6 +19,7 @@
   ],
   install_requires=[
     'accelerate',
+    'adam-atan2-pytorch>=0.0.12',
     'beartype',
     'classifier-free-guidance-pytorch>=0.6.10',
     'einops>=0.8.0',