1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom.py",
4  "content": "CONFIG_OVERRIDES = {\"normalize\": True, \"normalize_reward\": False}\n\nimport sys as _sys\n\ndef _detect_env():\n    for i, arg in enumerate(_sys.argv):\n        if arg == \"--env\" and i + 1 < len(_sys.argv):\n            return _sys.argv[i + 1]\n        if arg.startswith(\"--env=\"):\n            return arg.split(\"=\", 1)[1]\n    return \"\"\n\n_REBRAC_ENV = _detect_env()\n\n_REBRAC_HPARAMS = {\n    \"halfcheetah-medium-v2\": {\"actor_bc_coef\": 0.001, \"critic_bc_coef\": 0.01,  \"lr\": 1e-3, \"batch_size\": 1024, \"tau_q\": 0.8},\n    \"walker2d-medium-v2\":    {\"actor_bc_coef\": 0.05,  \"critic_bc_coef\": 0.1,   \"lr\": 1e-3, \"batch_size\": 1024, \"tau_q\": 0.8},\n    \"hopper-medium-v2\":      {\"actor_bc_coef\": 0.01,  \"critic_bc_coef\": 0.01,  \"lr\": 1e-3, \"batch_size\": 1024, \"tau_q\": 0.8},\n    \"maze2d-large-v1\":       {\"actor_bc_coef\": 0.003, \"critic_bc_coef\": 0.001, \"lr\": 3e-4, \"batch_size\": 256, \"tau_q\": 0.8},\n    \"maze2d-medium-v1\":      {\"actor_bc_coef\": 0.003, \"critic_bc_coef\": 0.001, \"lr\": 3e-4, \"batch_size\": 256, \"tau_q\": 0.8},\n}\n_REBRAC_HP = _REBRAC_HPARAMS.get(_REBRAC_ENV, {\"actor_bc_coef\": 0.01, \"critic_bc_coef\": 0.01, \"lr\": 1e-3, \"batch_size\": 1024, \"tau_q\": 0.8})\n\nCONFIG_OVERRIDES[\"batch_size\"] = _REBRAC_HP[\"batch_size\"]\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy pi(s) = tanh(net(s)) * max_action. No LayerNorm.\"\"\"\n    def __init__(self, state_dim: int, action_dim: int, max_action: float):\n        super().__init__()\n        self.max_action = max_action\n        self.net = nn.Sequential(\n            nn.Linear(state_dim, 256), nn.ReLU(),\n            nn.Linear(256, 256), nn.ReLU(),\n            nn.Linear(256, 256), nn.ReLU(),\n            nn.Linear(256, action_dim), nn.Tanh(),\n        )\n\n    def forward(self, state: torch.Tensor) -> torch.Tensor:\n        return self.max_action * self.net(state)\n\n    @torch.no_grad()\n    def act(self, state: np.ndarray, device: str = \"cpu\") -> np.ndarray:\n        state = torch.tensor(state.reshape(1, -1), device=device, dtype=torch.float32)\n        return self(state).cpu().data.numpy().flatten()\n\n\nclass Critic(nn.Module):\n    \"\"\"Q-function Q(s, a). LayerNorm after each hidden layer.\"\"\"\n    def __init__(self, state_dim: int, action_dim: int, orthogonal_init: bool = False):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(state_dim + action_dim, 256), nn.ReLU(), nn.LayerNorm(256),\n            nn.Linear(256, 256), nn.ReLU(), nn.LayerNorm(256),\n            nn.Linear(256, 256), nn.ReLU(), nn.LayerNorm(256),\n            nn.Linear(256, 1),\n        )\n\n    def forward(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n        return self.net(torch.cat([state, action], dim=-1)).squeeze(-1)\n\n\ndef asymmetric_l2_loss(u: torch.Tensor, tau: float) -> torch.Tensor:\n    \"\"\"Asymmetric L2 Loss (Expectile L2) scaled by 2 to match MSE at tau=0.5\"\"\"\n    weight = torch.abs(tau - (u < 0).float())\n    return torch.mean(2.0 * weight * (u ** 2))\n\n\nclass Actor(nn.Module):\n    # Dummy, unused. Kept to avoid breaking the base import requirements.\n    def __init__(self, *args, **kwargs):\n        super().__init__()\n        self.dummy = nn.Linear(1,1)\n\nclass ValueFunction(nn.Module):\n    # Dummy, unused.\n    def __init__(self, *args, **kwargs):\n        super().__init__()\n        self.dummy = nn.Linear(1,1)\n\n\nclass OfflineAlgorithm:\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        replay_buffer=None,\n        discount: float = 0.99,\n        tau: float = 5e-3,\n        actor_lr: float = 3e-4,\n        critic_lr: float = 3e-4,\n        alpha_lr: float = 3e-4,\n        orthogonal_init: bool = True,\n        device: str = \"cuda\",\n    ):\n        self.device = device\n        self.discount = discount\n        self.tau = tau\n        self.max_action = max_action\n        self.total_it = 0\n\n        self.actor_bc_coef = _REBRAC_HP[\"actor_bc_coef\"]\n        self.critic_bc_coef = _REBRAC_HP[\"critic_bc_coef\"]\n        self.tau_q = _REBRAC_HP[\"tau_q\"]\n        _lr = _REBRAC_HP[\"lr\"]\n        \n        self.policy_noise = 0.2\n        self.noise_clip = 0.5\n        self.policy_freq = 2\n        \n        # New hyperparam for our algorithm: Contrastive OOD penalty weight\n        self.ood_penalty_weight = 0.1 \n\n        self.actor = DeterministicActor(state_dim, action_dim, max_action).to(device)\n        self.actor_target = deepcopy(self.actor)\n        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=_lr)\n\n        self.critic_1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_1_target = deepcopy(self.critic_1)\n        self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(), lr=_lr)\n\n        self.critic_2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_2_target = deepcopy(self.critic_2)\n        self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(), lr=_lr)\n\n    def train(self, batch: TensorBatch) -> Dict[str, float]:\n        self.total_it += 1\n        states, actions, rewards, next_states, dones, next_actions_data = batch\n        not_done = 1 - dones.squeeze(-1)\n        rewards_flat = rewards.squeeze(-1)\n        log_dict: Dict[str, float] = {}\n\n        # Critic update\n        with torch.no_grad():\n            noise = (torch.randn_like(actions) * self.policy_noise).clamp(-self.noise_clip, self.noise_clip)\n            next_actions = (self.actor_target(next_states) + noise).clamp(-self.max_action, self.max_action)\n            \n            # ReBRAC explicit penalty\n            bc_penalty = ((next_actions - next_actions_data) ** 2).sum(-1)\n\n            target_q1 = self.critic_1_target(next_states, next_actions)\n            target_q2 = self.critic_2_target(next_states, next_actions)\n            target_q = torch.min(target_q1, target_q2) - self.critic_bc_coef * bc_penalty\n            target_q = rewards_flat + not_done * self.discount * target_q\n\n        current_q1 = self.critic_1(states, actions)\n        current_q2 = self.critic_2(states, actions)\n        \n        # Novelty: Expectile Asymmetric Loss to strictly suppress Q-value overestimation \n        # (tau_q=0.8 penalizes positive TD error 4x more heavily than negative TD error)\n        critic_loss1 = asymmetric_l2_loss(current_q1 - target_q, self.tau_q)\n        critic_loss2 = asymmetric_l2_loss(current_q2 - target_q, self.tau_q)\n        \n        # Additionally, add a small explicit term to suppress the value of policy's purely-imagined actions \n        # specifically if it exceeds the data's action value at the current state.\n        with torch.no_grad():\n            pi_actions = self.actor(states)\n        q1_pi = self.critic_1(states, pi_actions)\n        q2_pi = self.critic_2(states, pi_actions)\n        \n        # Conservative Penalty: soft-suppress Q(s, pi(s)) if it is disproportionately high\n        # We use a one-sided penalty (ReLU) to only penalize when Q_pi > Q_data + margin\n        q_data_mean = (current_q1.detach() + current_q2.detach()) / 2.0\n        cql_penalty = torch.mean(torch.nn.functional.relu(q1_pi - q_data_mean)) + \\\n                      torch.mean(torch.nn.functional.relu(q2_pi - q_data_mean))\n        \n        critic_loss = critic_loss1 + critic_loss2 + self.ood_penalty_weight * cql_penalty\n\n        log_dict[\"cql_penalty\"] = cql_penalty.item()\n        log_dict[\"critic_loss\"] = critic_loss.item()\n\n        self.critic_1_optimizer.zero_grad()\n        self.critic_2_optimizer.zero_grad()\n        critic_loss.backward()\n        self.critic_1_optimizer.step()\n        self.critic_2_optimizer.step()\n\n        # Actor update\n        if self.total_it % self.policy_freq == 0:\n            pi = self.actor(states)\n            q = self.critic_1(states, pi)\n\n            bc_err = ((pi - actions) ** 2).sum(-1)\n            \n            # Policy-level innovation: Scale Q gradients relative to the global mean, \n            # but adaptively regularize states where policy has very low Q variance (high confidence)\n            lmbda = 1.0 / (torch.abs(q).mean().detach() + 1e-8)\n\n            actor_loss = (self.actor_bc_coef * bc_err - lmbda * q).mean()\n            log_dict[\"actor_loss\"] = actor_loss.item()\n\n            self.actor_optimizer.zero_grad()\n            actor_loss.backward()\n            self.actor_optimizer.step()\n\n            for tp, sp in zip(self.critic_1_target.parameters(), self.critic_1.parameters()):\n                tp.data.copy_((1 - self.tau) * tp.data + self.tau * sp.data)\n            for tp, sp in zip(self.critic_2_target.parameters(), self.critic_2.parameters()):\n                tp.data.copy_((1 - self.tau) * tp.data + self.tau * sp.data)\n            for tp, sp in zip(self.actor_target.parameters(), self.actor.parameters()):\n                tp.data.copy_((1 - self.tau) * tp.data + self.tau * sp.data)\n\n        return log_dict\n",
5  "end_line": 397,
6  "start_line": 193
7}

1{
2  "filename": "CORL/algorithms/offline/custom.py",
3  "start_line": 193,
4  "content": "CONFIG_OVERRIDES = {\"normalize\": True, \"normalize_reward\": False}\n\nimport sys as _sys\n\ndef _detect_env():\n    for i, arg in enumerate(_sys.argv):\n        if arg == \"--env\" and i + 1 < len(_sys.argv):\n            return _sys.argv[i + 1]\n        if arg.startswith(\"--env=\"):\n            return arg.split(\"=\", 1)[1]\n    return \"\"\n\n_REBRAC_ENV = _detect_env()\n\n_REBRAC_HPARAMS = {\n    \"halfcheetah-medium-v2\": {\"actor_bc_coef\": 0.001, \"critic_bc_coef\": 0.01,  \"lr\": 1e-3, \"batch_size\": 1024},\n    \"walker2d-medium-v2\":    {\"actor_bc_coef\": 0.05,  \"critic_bc_coef\": 0.1,   \"lr\": 1e-3, \"batch_size\": 1024},\n    \"hopper-medium-v2\":      {\"actor_bc_coef\": 0.01,  \"critic_bc_coef\": 0.01,  \"lr\": 1e-3, \"batch_size\": 1024},\n    \"maze2d-large-v1\":       {\"actor_bc_coef\": 0.003, \"critic_bc_coef\": 0.001, \"lr\": 3e-4, \"batch_size\": 256},\n    \"maze2d-medium-v1\":      {\"actor_bc_coef\": 0.003, \"critic_bc_coef\": 0.001, \"lr\": 3e-4, \"batch_size\": 256},\n}\n_REBRAC_HP = _REBRAC_HPARAMS.get(_REBRAC_ENV, {\"actor_bc_coef\": 0.01, \"critic_bc_coef\": 0.01, \"lr\": 1e-3, \"batch_size\": 1024})\n\nCONFIG_OVERRIDES[\"batch_size\"] = _REBRAC_HP[\"batch_size\"]\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy pi(s) = tanh(net(s)) * max_action. No LayerNorm.\"\"\"\n    def __init__(self, state_dim: int, action_dim: int, max_action: float):\n        super().__init__()\n        self.max_action = max_action\n        self.net = nn.Sequential(\n            nn.Linear(state_dim, 256), nn.ReLU(),\n            nn.Linear(256, 256), nn.ReLU(),\n            nn.Linear(256, 256), nn.ReLU(),\n            nn.Linear(256, action_dim), nn.Tanh(),\n        )\n\n    def forward(self, state: torch.Tensor) -> torch.Tensor:\n        return self.max_action * self.net(state)\n\n    @torch.no_grad()\n    def act(self, state: np.ndarray, device: str = \"cpu\") -> np.ndarray:\n        state = torch.tensor(state.reshape(1, -1), device=device, dtype=torch.float32)\n        return self(state).cpu().data.numpy().flatten()\n\n\nclass Critic(nn.Module):\n    \"\"\"Q-function Q(s, a). LayerNorm after each hidden layer.\"\"\"\n    def __init__(self, state_dim: int, action_dim: int, orthogonal_init: bool = False):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(state_dim + action_dim, 256), nn.ReLU(), nn.LayerNorm(256),\n            nn.Linear(256, 256), nn.ReLU(), nn.LayerNorm(256),\n            nn.Linear(256, 256), nn.ReLU(), nn.LayerNorm(256),\n            nn.Linear(256, 1),\n        )\n\n    def forward(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n        return self.net(torch.cat([state, action], dim=-1)).squeeze(-1)\n\n\n# Unused stubs to prevent import errors\nclass Actor(nn.Module):\n    def __init__(self, *args, **kwargs):\n        super().__init__()\n        self.net = nn.Linear(1,1)\n\nclass ValueFunction(nn.Module):\n    def __init__(self, *args, **kwargs):\n        super().__init__()\n        self.net = nn.Linear(1,1)\n\n\nclass OfflineAlgorithm:\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        replay_buffer=None,\n        discount: float = 0.99,\n        tau: float = 5e-3,\n        actor_lr: float = 3e-4,\n        critic_lr: float = 3e-4,\n        alpha_lr: float = 3e-4,\n        orthogonal_init: bool = True,\n        device: str = \"cuda\",\n    ):\n        self.device = device\n        self.discount = discount\n        self.tau = tau\n        self.max_action = max_action\n        self.total_it = 0\n\n        self.actor_bc_coef = _REBRAC_HP[\"actor_bc_coef\"]\n        self.critic_bc_coef = _REBRAC_HP[\"critic_bc_coef\"]\n        _lr = _REBRAC_HP[\"lr\"]\n        \n        self.policy_noise = 0.2\n        self.noise_clip = 0.5\n        self.policy_freq = 2\n        \n        # Hyperparameters for Uncertainty-Guided ReBRAC (UG-ReBRAC)\n        self.kappa = 0.5  # Controls strength of uncertainty modulation\n        self.beta_var = 0.1 # Explicit variance penalty for pseudo-ensemble effect\n\n        self.actor = DeterministicActor(state_dim, action_dim, max_action).to(device)\n        self.actor_target = deepcopy(self.actor)\n        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=_lr)\n\n        self.critic_1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_1_target = deepcopy(self.critic_1)\n        self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(), lr=_lr)\n\n        self.critic_2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_2_target = deepcopy(self.critic_2)\n        self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(), lr=_lr)\n\n    def train(self, batch: TensorBatch) -> Dict[str, float]:\n        self.total_it += 1\n        states, actions, rewards, next_states, dones, next_actions_data = batch\n        not_done = 1 - dones.squeeze(-1)\n        rewards_flat = rewards.squeeze(-1)\n        log_dict: Dict[str, float] = {}\n\n        # ── Critic update ──────────────────────────────────────────────\n        with torch.no_grad():\n            noise = (torch.randn_like(actions) * self.policy_noise).clamp(\n                -self.noise_clip, self.noise_clip\n            )\n            next_actions = (self.actor_target(next_states) + noise).clamp(\n                -self.max_action, self.max_action\n            )\n            \n            target_q1 = self.critic_1_target(next_states, next_actions)\n            target_q2 = self.critic_2_target(next_states, next_actions)\n            \n            # Uncertainty computation (Epistemic proxy)\n            q_var = torch.abs(target_q1 - target_q2)\n            q_var_norm = q_var / (q_var.mean() + 1e-6)\n            \n            # Uncertainty-modulated Target\n            # Scale bc penalty depending on epistemic uncertainty.\n            multiplier = torch.clamp(1.0 + self.kappa * (q_var_norm - 1.0), min=0.5, max=5.0)\n            adaptive_critic_bc = self.critic_bc_coef * multiplier\n            \n            bc_penalty = ((next_actions - next_actions_data) ** 2).sum(-1)\n\n            q_min = torch.min(target_q1, target_q2)\n            \n            # Explicit pseudo-ensemble regularizer\n            target_q = q_min - adaptive_critic_bc * bc_penalty - self.beta_var * q_var\n            target_q = rewards_flat + not_done * self.discount * target_q\n\n        current_q1 = self.critic_1(states, actions)\n        current_q2 = self.critic_2(states, actions)\n        \n        critic_loss = F.mse_loss(current_q1, target_q) + F.mse_loss(current_q2, target_q)\n        log_dict[\"critic_loss\"] = critic_loss.item()\n\n        self.critic_1_optimizer.zero_grad()\n        self.critic_2_optimizer.zero_grad()\n        critic_loss.backward()\n        self.critic_1_optimizer.step()\n        self.critic_2_optimizer.step()\n\n        # ── Delayed actor update ───────────────────────────────────────\n        if self.total_it % self.policy_freq == 0:\n            pi = self.actor(states)\n            q1_pi = self.critic_1(states, pi)\n            q2_pi = self.critic_2(states, pi)\n            q = q1_pi\n            \n            # Actor Uncertainty modulation\n            q_var_pi = torch.abs(q1_pi - q2_pi)\n            q_var_norm_pi = q_var_pi / (q_var_pi.mean() + 1e-6)\n            \n            mult_pi = torch.clamp(1.0 + self.kappa * (q_var_norm_pi - 1.0), min=0.5, max=5.0)\n            adaptive_actor_bc = self.actor_bc_coef * mult_pi.detach()\n\n            bc_err = ((pi - actions) ** 2).sum(-1)\n\n            lmbda = 1.0 / (torch.abs(q).mean().detach() + 1e-8)\n\n            actor_loss = (adaptive_actor_bc * bc_err - lmbda * q).mean()\n            log_dict[\"actor_loss\"] = actor_loss.item()\n\n            self.actor_optimizer.zero_grad()\n            actor_loss.backward()\n            self.actor_optimizer.step()\n\n            for tp, sp in zip(self.critic_1_target.parameters(), self.critic_1.parameters()):\n                tp.data.copy_((1 - self.tau) * tp.data + self.tau * sp.data)\n            for tp, sp in zip(self.critic_2_target.parameters(), self.critic_2.parameters()):\n                tp.data.copy_((1 - self.tau) * tp.data + self.tau * sp.data)\n            for tp, sp in zip(self.actor_target.parameters(), self.actor.parameters()):\n                tp.data.copy_((1 - self.tau) * tp.data + self.tau * sp.data)\n\n        return log_dict",
5  "end_line": 395,
6  "op": "replace"
7}

1{
2  "op": "replace",
3  "end_line": 391,
4  "start_line": 193,
5  "filename": "CORL/algorithms/offline/custom.py",
6  "content": "CONFIG_OVERRIDES = {\"normalize\": True, \"normalize_reward\": False}\n\nimport sys as _sys\n\ndef _detect_env():\n    for i, arg in enumerate(_sys.argv):\n        if arg == \"--env\" and i + 1 < len(_sys.argv):\n            return _sys.argv[i + 1]\n        if arg.startswith(\"--env=\"):\n            return arg.split(\"=\", 1)[1]\n    return \"\"\n\n_REBRAC_ENV = _detect_env()\n\n_REBRAC_HPARAMS = {\n    \"halfcheetah-medium-v2\": {\"actor_bc_coef\": 0.001, \"critic_bc_coef\": 0.01,  \"lr\": 1e-3, \"batch_size\": 1024},\n    \"walker2d-medium-v2\":    {\"actor_bc_coef\": 0.05,  \"critic_bc_coef\": 0.1,   \"lr\": 1e-3, \"batch_size\": 1024},\n    \"hopper-medium-v2\":      {\"actor_bc_coef\": 0.01,  \"critic_bc_coef\": 0.01,  \"lr\": 1e-3, \"batch_size\": 1024},\n    \"maze2d-large-v1\":       {\"actor_bc_coef\": 0.003, \"critic_bc_coef\": 0.001, \"lr\": 3e-4, \"batch_size\": 256},\n    \"maze2d-medium-v1\":      {\"actor_bc_coef\": 0.003, \"critic_bc_coef\": 0.001, \"lr\": 3e-4, \"batch_size\": 256},\n}\n_REBRAC_HP = _REBRAC_HPARAMS.get(_REBRAC_ENV, {\"actor_bc_coef\": 0.01, \"critic_bc_coef\": 0.01, \"lr\": 1e-3, \"batch_size\": 1024})\n\nCONFIG_OVERRIDES[\"batch_size\"] = _REBRAC_HP[\"batch_size\"]\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy pi(s) = tanh(net(s)) * max_action. No LayerNorm.\"\"\"\n    def __init__(self, state_dim: int, action_dim: int, max_action: float):\n        super().__init__()\n        self.max_action = max_action\n        self.net = nn.Sequential(\n            nn.Linear(state_dim, 256), nn.ReLU(),\n            nn.Linear(256, 256), nn.ReLU(),\n            nn.Linear(256, 256), nn.ReLU(),\n            nn.Linear(256, action_dim), nn.Tanh(),\n        )\n\n    def forward(self, state: torch.Tensor) -> torch.Tensor:\n        return self.max_action * self.net(state)\n\n    @torch.no_grad()\n    def act(self, state: np.ndarray, device: str = \"cpu\") -> np.ndarray:\n        state = torch.tensor(state.reshape(1, -1), device=device, dtype=torch.float32)\n        return self(state).cpu().data.numpy().flatten()\n\n\nclass Critic(nn.Module):\n    \"\"\"Q-function Q(s, a). LayerNorm after each hidden layer.\"\"\"\n    def __init__(self, state_dim: int, action_dim: int, orthogonal_init: bool = False):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(state_dim + action_dim, 256), nn.ReLU(), nn.LayerNorm(256),\n            nn.Linear(256, 256), nn.ReLU(), nn.LayerNorm(256),\n            nn.Linear(256, 256), nn.ReLU(), nn.LayerNorm(256),\n            nn.Linear(256, 1),\n        )\n\n    def forward(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n        return self.net(torch.cat([state, action], dim=-1)).squeeze(-1)\n\n\n# Unused stubs to prevent import errors\nclass Actor(nn.Module):\n    def __init__(self, *args, **kwargs):\n        super().__init__()\n        self.net = nn.Linear(1,1)\n\nclass ValueFunction(nn.Module):\n    def __init__(self, *args, **kwargs):\n        super().__init__()\n        self.net = nn.Linear(1,1)\n\n\nclass OfflineAlgorithm:\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        replay_buffer=None,\n        discount: float = 0.99,\n        tau: float = 5e-3,\n        actor_lr: float = 3e-4,\n        critic_lr: float = 3e-4,\n        alpha_lr: float = 3e-4,\n        orthogonal_init: bool = True,\n        device: str = \"cuda\",\n    ):\n        self.device = device\n        self.discount = discount\n        self.tau = tau\n        self.max_action = max_action\n        self.total_it = 0\n\n        self.actor_bc_coef = _REBRAC_HP[\"actor_bc_coef\"]\n        self.critic_bc_coef = _REBRAC_HP[\"critic_bc_coef\"]\n        _lr = _REBRAC_HP[\"lr\"]\n        \n        self.policy_noise = 0.2\n        self.noise_clip = 0.5\n        self.policy_freq = 2\n        \n        # Hyperparameters for Uncertainty-Guided ReBRAC (UG-ReBRAC)\n        self.kappa = 1.0  # Controls strength of uncertainty modulation\n\n        self.actor = DeterministicActor(state_dim, action_dim, max_action).to(device)\n        self.actor_target = deepcopy(self.actor)\n        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=_lr)\n\n        self.critic_1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_1_target = deepcopy(self.critic_1)\n        self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(), lr=_lr)\n\n        self.critic_2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_2_target = deepcopy(self.critic_2)\n        self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(), lr=_lr)\n\n    def train(self, batch: TensorBatch) -> Dict[str, float]:\n        self.total_it += 1\n        states, actions, rewards, next_states, dones, next_actions_data = batch\n        not_done = 1 - dones.squeeze(-1)\n        rewards_flat = rewards.squeeze(-1)\n        log_dict: Dict[str, float] = {}\n\n        # ── Critic update ──────────────────────────────────────────────\n        with torch.no_grad():\n            noise = (torch.randn_like(actions) * self.policy_noise).clamp(\n                -self.noise_clip, self.noise_clip\n            )\n            next_actions = (self.actor_target(next_states) + noise).clamp(\n                -self.max_action, self.max_action\n            )\n            \n            target_q1 = self.critic_1_target(next_states, next_actions)\n            target_q2 = self.critic_2_target(next_states, next_actions)\n            \n            # Uncertainty computation (Epistemic proxy)\n            q_var = torch.abs(target_q1 - target_q2)\n            q_var_norm = q_var / (q_var.mean() + 1e-8)\n            \n            # Uncertainty-modulated Target\n            # Scale bc penalty depending on epistemic uncertainty.\n            multiplier = torch.clamp(1.0 + self.kappa * (q_var_norm - 1.0), min=0.5, max=5.0)\n            adaptive_critic_bc = self.critic_bc_coef * multiplier\n            \n            bc_penalty = ((next_actions - next_actions_data) ** 2).sum(-1)\n\n            target_q = torch.min(target_q1, target_q2) - adaptive_critic_bc * bc_penalty\n            target_q = rewards_flat + not_done * self.discount * target_q\n\n        current_q1 = self.critic_1(states, actions)\n        current_q2 = self.critic_2(states, actions)\n        \n        critic_loss = F.mse_loss(current_q1, target_q) + F.mse_loss(current_q2, target_q)\n        log_dict[\"critic_loss\"] = critic_loss.item()\n\n        self.critic_1_optimizer.zero_grad()\n        self.critic_2_optimizer.zero_grad()\n        critic_loss.backward()\n        self.critic_1_optimizer.step()\n        self.critic_2_optimizer.step()\n\n        # ── Delayed actor update ───────────────────────────────────────\n        if self.total_it % self.policy_freq == 0:\n            pi = self.actor(states)\n            q1_pi = self.critic_1(states, pi)\n            q2_pi = self.critic_2(states, pi)\n            q = q1_pi\n            \n            # Actor Uncertainty modulation\n            q_var_pi = torch.abs(q1_pi - q2_pi)\n            q_var_norm_pi = q_var_pi / (q_var_pi.mean() + 1e-8)\n            \n            mult_pi = torch.clamp(1.0 + self.kappa * (q_var_norm_pi - 1.0), min=0.5, max=5.0)\n            adaptive_actor_bc = self.actor_bc_coef * mult_pi.detach()\n\n            bc_err = ((pi - actions) ** 2).sum(-1)\n\n            lmbda = 1.0 / (torch.abs(q).mean().detach() + 1e-8)\n\n            actor_loss = (adaptive_actor_bc * bc_err - lmbda * q).mean()\n            log_dict[\"actor_loss\"] = actor_loss.item()\n\n            self.actor_optimizer.zero_grad()\n            actor_loss.backward()\n            self.actor_optimizer.step()\n\n            for tp, sp in zip(self.critic_1_target.parameters(), self.critic_1.parameters()):\n                tp.data.copy_((1 - self.tau) * tp.data + self.tau * sp.data)\n            for tp, sp in zip(self.critic_2_target.parameters(), self.critic_2.parameters()):\n                tp.data.copy_((1 - self.tau) * tp.data + self.tau * sp.data)\n            for tp, sp in zip(self.actor_target.parameters(), self.actor.parameters()):\n                tp.data.copy_((1 - self.tau) * tp.data + self.tau * sp.data)\n\n        return log_dict"
7}