[update]: ctrl new architecture and clean the code (#14)

* add layer norm

* update: add metrics logging and regularizer

* fix minor bugs

* fix typo

* before using catgorical Q

* update: clean some code and finalize ctrl implementation

---------

Co-authored-by: Edward Chen <echen333us@gmail.com>

Author: typoverflow

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

@@ -2,49 +2,38 @@
 
 algo:
   name: ctrl_td3
-  cls: ctrl_td3
-  discount: 0.99
-
-  tau: 0.005 # hard update
   actor_update_freq: 1
   target_update_freq: 1
-
+  discount: 0.99
   ema: 0.005
-
-  # critic_hidden_dims: [512, 512, 512] # this is actually not used
   actor_hidden_dims: [512, 512, 512]
-
-  critic_lr: 0.0003
+  # critic_hidden_dims: [512, 512, 512] # not used
+  activation: elu  # not used
+  critic_ensemble_size: 2
+  layer_norm: true
   actor_lr: 0.0003
-
+  critic_lr: 0.0003
+  clip_grad_norm: null
   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
+  feature_ema: 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
+  ctrl_coef: 1.0
+  reward_coef: 1.0
   back_critic_grad: false
   critic_coef: 1.0
-  aug_batch_size: 512
 
   num_noises: 25
   linear: false
-  beta: 1.0 # not used
   ranking: true
-  activation: elu
-  layer_norm: true
-  critic_ensemble_size: 2
-  clip_grad_norm: null
 
 norm_obs: true
-batch_size: 1024

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

@@ -2,6 +2,8 @@
 
 algo:
   name: td3
+  actor_update_freq: 1
+  target_update_freq: 1
   discount: 0.99
   ema: 0.005
   actor_hidden_dims: [512, 512, 512]
@@ -12,8 +14,6 @@ algo:
   actor_lr: 0.0003
   critic_lr: 0.0003
   clip_grad_norm: null
-  actor_update_freq: 1
-  target_update_freq: 1
   target_policy_noise: 0.2
   noise_clip: 0.3
   exploration_noise: 0.2

examples/online/main_dmc_offpolicy.py

@@ -6,10 +6,10 @@
 import jax.numpy as jnp
 import numpy as np
 import omegaconf
-import wandb
 from omegaconf import OmegaConf
 from tqdm import tqdm, trange
 
+import wandb
 from flowrl.agent.online import *
 from flowrl.config.online.mujoco import Config
 from flowrl.dataset.buffer.state import ReplayBuffer
@@ -26,7 +26,7 @@
     "td7": TD7Agent,
     "sdac": SDACAgent,
     "dpmd": DPMDAgent,
-    "ctrl_td3": Ctrl_TD3_Agent,
+    "ctrl_td3": CtrlTD3Agent,
 }
 
 class OffPolicyTrainer():

flowrl/agent/online/__init__.py

@@ -1,5 +1,5 @@
 from ..base import BaseAgent
-from .ctrl.ctrl import Ctrl_TD3_Agent
+from .ctrl.ctrl import CtrlTD3Agent
 from .dpmd import DPMDAgent
 from .ppo import PPOAgent
 from .sac import SACAgent
@@ -15,5 +15,5 @@
     "SDACAgent",
     "DPMDAgent",
     "PPOAgent",
-    "Ctrl_TD3_Agent",
+    "CtrlTD3Agent",
 ]

flowrl/agent/online/ctrl/ctrl.py

@@ -1,5 +1,4 @@
 from functools import partial
-from operator import attrgetter
 from typing import Tuple
 
 import jax
@@ -8,7 +7,7 @@
 
 from flowrl.agent.online.ctrl.network import FactorizedNCE, update_factorized_nce
 from flowrl.agent.online.td3 import TD3Agent
-from flowrl.config.online.mujoco.algo.ctrl_td3 import CTRL_TD3_Config
+from flowrl.config.online.mujoco.algo.ctrl_td3 import CtrlTD3Config
 from flowrl.functional.ema import ema_update
 from flowrl.module.actor import SquashedDeterministicActor
 from flowrl.module.mlp import MLP
@@ -23,7 +22,6 @@ def update_critic(
     critic: Model,
     critic_target: Model,
     actor_target: Model,
-    nce: Model,
     nce_target: Model,
     batch: Batch,
     discount: float,
@@ -42,7 +40,6 @@ def update_critic(
 
     back_critic_grad = False
     if back_critic_grad:
-        # this part will use feature
         raise NotImplementedError("no back critic grad exists")
 
     feature = nce_target(batch.obs, batch.action, method="forward_phi")
@@ -53,7 +50,6 @@ def critic_loss_fn(critic_params: Param, dropout_rng: PRNGKey) -> Tuple[jnp.ndar
             feature,
             rngs={"dropout": dropout_rng},
         )
-        # q_pred (2, 512, 1), q_target (512, 1)
         critic_loss = critic_coef * ((q_pred - q_target[jnp.newaxis, :])**2).sum(0).mean()
         return critic_loss, {
             "loss/critic_loss": critic_loss,
@@ -69,7 +65,6 @@ def critic_loss_fn(critic_params: Param, dropout_rng: PRNGKey) -> Tuple[jnp.ndar
 def update_actor(
     rng: PRNGKey,
     actor: Model,
-    nce: Model,
     nce_target: Model,
     critic: Model,
     batch: Batch,
@@ -95,23 +90,21 @@ def actor_loss_fn(
     return rng, new_actor, metrics
 
 
-class Ctrl_TD3_Agent(TD3Agent):
+class CtrlTD3Agent(TD3Agent):
     """
-    CTRL Twin Delayed Deep Deterministic Policy Gradient (TD3) agent.
+    CTRL with Twin Delayed Deep Deterministic Policy Gradient (TD3) agent.
     """
 
-    name = "CTRLTD3Agent"
+    name = "CtrlTD3Agent"
     model_names = ["nce", "nce_target", "actor", "actor_target", "critic", "critic_target"]
 
-    def __init__(self, obs_dim: int, act_dim: int, cfg: CTRL_TD3_Config, seed: int):
+    def __init__(self, obs_dim: int, act_dim: int, cfg: CtrlTD3Config, seed: int):
         super().__init__(obs_dim, act_dim, cfg, seed)
         self.cfg = cfg
 
         self.ctrl_coef = cfg.ctrl_coef
         self.critic_coef = cfg.critic_coef
 
-        self.aug_batch_size = cfg.aug_batch_size
-        self.feature_tau = cfg.feature_tau
         self.linear = cfg.linear
         self.ranking = cfg.ranking
         self.feature_dim = cfg.feature_dim
@@ -120,10 +113,10 @@ def __init__(self, obs_dim: int, act_dim: int, cfg: CTRL_TD3_Config, seed: int):
         self.rff_dim = cfg.rff_dim
 
         # sanity checks for the hyper-parameters
-        assert not self.linear, "Removing linear version for now"
+        assert not self.linear, "linear mode is not supported yet"
 
         # networks
-        self.rng, nce_rng, actor_rng, critic_rng = jax.random.split(self.rng, 4)
+        self.rng, nce_rng, nce_init_rng, actor_rng, critic_rng = jax.random.split(self.rng, 5)
         nce_def = FactorizedNCE(
             self.obs_dim,
             self.act_dim,
@@ -139,6 +132,7 @@ def __init__(self, obs_dim: int, act_dim: int, cfg: CTRL_TD3_Config, seed: int):
             nce_def,
             nce_rng,
             inputs=(
+                nce_init_rng,
                 jnp.ones((1, self.obs_dim)),
                 jnp.ones((1, self.act_dim)),
                 jnp.ones((1, self.obs_dim)),
@@ -150,6 +144,7 @@ def __init__(self, obs_dim: int, act_dim: int, cfg: CTRL_TD3_Config, seed: int):
             nce_def,
             nce_rng,
             inputs=(
+                nce_init_rng,
                 jnp.ones((1, self.obs_dim)),
                 jnp.ones((1, self.act_dim)),
                 jnp.ones((1, self.obs_dim)),
@@ -201,28 +196,10 @@ def __init__(self, obs_dim: int, act_dim: int, cfg: CTRL_TD3_Config, seed: int):
     def train_step(self, batch: Batch, step: int) -> Metric:
         metrics = {}
 
-        split_index = batch.obs.shape[0] - self.aug_batch_size
-        obs, action, next_obs, reward, terminal = [
-            b[:split_index]
-            for b in attrgetter("obs", "action", "next_obs", "reward", "terminal")(
-                batch
-            )
-        ]
-        fobs, faction, fnext_obs, freward, fterminal = [
-            b[split_index:]
-            for b in attrgetter("obs", "action", "next_obs", "reward", "terminal")(
-                batch
-            )
-        ]
-        rl_batch = Batch(obs, action, reward, terminal, next_obs, None)
-
         self.rng, self.nce, nce_metrics = update_factorized_nce(
             self.rng,
             self.nce,
-            fobs,
-            faction,
-            fnext_obs,
-            freward,
+            batch,
             self.ranking,
             self.reward_coef,
         )
@@ -233,9 +210,8 @@ def train_step(self, batch: Batch, step: int) -> Metric:
             self.critic,
             self.critic_target,
             self.actor_target,
-            self.nce,
             self.nce_target,
-            rl_batch,
+            batch,
             discount=self.cfg.discount,
             target_policy_noise=self.target_policy_noise,
             noise_clip=self.noise_clip,
@@ -247,10 +223,9 @@ def train_step(self, batch: Batch, step: int) -> Metric:
             self.rng, self.actor, actor_metrics = update_actor(
                 self.rng,
                 self.actor,
-                self.nce,
                 self.nce_target,
                 self.critic,
-                rl_batch,
+                batch,
             )
             metrics.update(actor_metrics)
 
@@ -262,4 +237,4 @@ def train_step(self, batch: Batch, step: int) -> Metric:
 
     def sync_target(self):
         super().sync_target()
-        self.nce_target = ema_update(self.nce, self.nce_target, self.feature_tau)
+        self.nce_target = ema_update(self.nce, self.nce_target, self.cfg.feature_ema)