[Draft]: add CTRL

examples/online/config/dmc/algo/td3_ctrl.yaml

@@ -0,0 +1,41 @@
+# @package _global_
+
+algo:
+  cls: ctrl_td3
+  discount: 0.99
+
+  tau: 0.005 # hard update
+  actor_update_freq: 1
+  target_update_freq: 1
+
+  # critic_hidden_dims: [512, 512, 512] # this is actually not used
+  actor_hidden_dims: [512, 512, 512]
+
+  critic_lr: 0.0003
+  actor_lr: 0.0003
+
+  target_policy_noise: 0.2
+  noise_clip: 0.3
+  exploration_noise: 0.2
+
+  # below are params specific to ctrl_td3
+  feature_dim: 512
+  feature_lr: 0.0001
+  feature_tau: 0.005
+  phi_hidden_dims: [512, 512]
+  mu_hidden_dims: [512, 512]
+  ctrl_coef: 1.0
+  reward_coef: 1.0
+  feature_update_ratio: 1
+  critic_hidden_dims: [512, ]
+  reward_hidden_dims: [512, ]
+  rff_dim: 1024
+  back_critic_grad: false
+  critic_coef: 1.0
+  aug_batch_size: 512
+  num_noises: 25
+  linear: false
+  beta: 1.0
+  ranking: true
+
+norm_obs: true

examples/online/main_dmc_offpolicy.py

@@ -23,6 +23,7 @@
     "td7": TD7Agent,
     "sdac": SDACAgent,
     "dpmd": DPMDAgent,
+    "td3_ctrl": Ctrl_TD3_Agent,
 }
 
 class OffPolicyTrainer():
@@ -116,7 +117,8 @@ def train(self):
                 if self.global_frame < cfg.warmup_frames:
                     train_metrics = {}
                 else:
-                    batch, indices = self.buffer.sample(batch_size=cfg.batch_size)
+                    batch_size = cfg.batch_size + getattr(cfg, "aug_batch_size", 0)
+                    batch, indices = self.buffer.sample(batch_size=batch_size)
                     train_metrics = self.agent.train_step(batch, step=self.global_frame)
                     if self.use_lap_buffer:
                         new_priorities = train_metrics.pop("priority")

flowrl/agent/online/__init__.py

@@ -5,6 +5,7 @@
 from .sdac import SDACAgent
 from .td3 import TD3Agent
 from .td7.td7 import TD7Agent
+from .ctrl.ctrl import Ctrl_TD3_Agent
 
 __all__ = [
     "BaseAgent",
@@ -14,4 +15,5 @@
     "SDACAgent",
     "DPMDAgent",
     "PPOAgent",
+    "Ctrl_TD3_Agent",
 ]