1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom.py",
4  "start_line": 193,
5  "end_line": 397,
6  "content": "import sys as _sys\n\n\ndef _detect_env():\n    \"\"\"Parse --env from sys.argv.\"\"\"\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\n_ENV_NAME = _detect_env()\n\n# Per-environment hyperparameters (base from ReBRAC CORL configs)\n_HPARAMS = {\n    \"halfcheetah-medium-v2\": {\"actor_bc\": 0.001, \"critic_bc\": 0.01,  \"lr\": 1e-3},\n    \"walker2d-medium-v2\":    {\"actor_bc\": 0.05,  \"critic_bc\": 0.1,   \"lr\": 1e-3},\n    \"hopper-medium-v2\":      {\"actor_bc\": 0.01,  \"critic_bc\": 0.01,  \"lr\": 1e-3},\n    \"maze2d-large-v1\":       {\"actor_bc\": 0.003, \"critic_bc\": 0.001, \"lr\": 3e-4},\n    \"maze2d-medium-v1\":      {\"actor_bc\": 0.003, \"critic_bc\": 0.001, \"lr\": 3e-4},\n    \"maze2d-umaze-v1\":       {\"actor_bc\": 0.003, \"critic_bc\": 0.001, \"lr\": 3e-4},\n}\n_HP = _HPARAMS.get(_ENV_NAME, {\"actor_bc\": 0.01, \"critic_bc\": 0.01, \"lr\": 1e-3})\n\nCONFIG_OVERRIDES: Dict[str, Any] = {}\n\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy: 3x256 MLP -> tanh * max_action.\"\"\"\n\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        s = torch.tensor(state.reshape(1, -1), device=device, dtype=torch.float32)\n        return self(s).cpu().data.numpy().flatten()\n\n\nclass Actor(nn.Module):\n    \"\"\"Stochastic Tanh-Gaussian policy: 3x256 MLP (kept for compatibility).\"\"\"\n\n    def __init__(self, state_dim: int, action_dim: int, max_action: float,\n                 orthogonal_init: bool = False):\n        super().__init__()\n        self.max_action = max_action\n        self.action_dim = action_dim\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, 2 * action_dim),\n        )\n        init_module_weights(self.net, orthogonal_init)\n        self.log_std_min, self.log_std_max = -20.0, 2.0\n\n    def _get_dist(self, state):\n        out = self.net(state)\n        mean, log_std = torch.split(out, self.action_dim, dim=-1)\n        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)\n        return TransformedDistribution(\n            Normal(mean, torch.exp(log_std)), TanhTransform(cache_size=1)), mean\n\n    def forward(self, state, deterministic=False):\n        dist, mean = self._get_dist(state)\n        action = torch.tanh(mean) if deterministic else dist.rsample()\n        return self.max_action * action, dist.log_prob(action).sum(-1)\n\n    @torch.no_grad()\n    def act(self, state, device=\"cpu\"):\n        s = torch.tensor(state.reshape(1, -1), device=device, dtype=torch.float32)\n        a, _ = self(s, not self.training)\n        return a.cpu().data.numpy().flatten()\n\n\nclass Critic(nn.Module):\n    \"\"\"Q(s,a): 3x256 MLP with LayerNorm for stable offline training.\"\"\"\n\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, action):\n        return self.net(torch.cat([state, action], dim=-1)).squeeze(-1)\n\n\nclass ValueFunction(nn.Module):\n    \"\"\"V(s): 3x256 MLP (kept for compatibility).\"\"\"\n\n    def __init__(self, state_dim: int, orthogonal_init: bool = False):\n        super().__init__()\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, 1),\n        )\n        init_module_weights(self.net, orthogonal_init)\n\n    def forward(self, state):\n        return self.net(state).squeeze(-1)\n\n\nclass OfflineAlgorithm:\n    \"\"\"Annealed Softmin with Advantage-Conditioned BC (ASM-ACBC).\n\n    Novel contributions over standard ReBRAC:\n    1. Annealed Softmin (ASM): Critic target interpolates between\n       min(Q1,Q2) and mean(Q1,Q2) with eta increasing linearly during\n       training. Reduces underestimation bias as Q-estimates improve,\n       while ReBRAC's BC penalties still handle OOD conservatism.\n    2. Advantage-Conditioned BC (ACBC): Actor BC penalty weighted by\n       normalized advantage of dataset actions (after warmup). Focuses\n       regularization on high-quality demonstrations while allowing\n       deviation from low-quality ones. Uses tanh modulation for\n       bounded, mean-preserving weights.\n    3. Momentum Q-Norm: EMA-smoothed |Q| replaces per-batch |Q|.mean()\n       for the actor gradient normalization, reducing variance.\n    \"\"\"\n\n    def __init__(self, state_dim, action_dim, max_action, replay_buffer=None,\n                 discount=0.99, tau=5e-3, actor_lr=3e-4, critic_lr=3e-4,\n                 alpha_lr=3e-4, orthogonal_init=True, device=\"cuda\"):\n        self.device = device\n        self.discount = discount\n        self.tau = tau\n        self.max_action = max_action\n        self.total_it = 0\n        self.replay_buffer = replay_buffer\n        # Per-env hyperparameters\n        self.actor_bc_coef = _HP[\"actor_bc\"]\n        self.critic_bc_coef = _HP[\"critic_bc\"]\n        _lr = _HP[\"lr\"]\n        self.policy_noise = 0.2\n        self.noise_clip = 0.5\n        self.policy_freq = 2\n        # ASM: Annealed Softmin parameters\n        self.eta_base = 0.2           # Max interpolation toward mean\n        self.eta_ramp_steps = 8e5     # Steps to reach full eta\n        # ACBC: Advantage-Conditioned BC parameters\n        self.acbc_warmup = 50000      # Warmup before enabling\n        self.acbc_scale = 0.5         # Modulation strength [0.5, 1.5] range\n        # Momentum Q-normalization\n        self.ema_q_abs = 1.0\n        self.ema_rate = 0.005\n        # Networks\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        self.critic_1 = Critic(state_dim, action_dim).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        self.critic_2 = Critic(state_dim, action_dim).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            next_actions = (self.actor_target(next_states) + noise).clamp(\n                -self.max_action, self.max_action)\n            bc_penalty = ((next_actions - next_actions_data) ** 2).sum(-1)\n            tq1 = self.critic_1_target(next_states, next_actions)\n            tq2 = self.critic_2_target(next_states, next_actions)\n            # ASM: Annealed softmin — gradually reduce underestimation\n            q_min = torch.min(tq1, tq2)\n            q_mean = (tq1 + tq2) * 0.5\n            eta = self.eta_base * min(self.total_it / self.eta_ramp_steps, 1.0)\n            target_q = (1.0 - eta) * q_min + eta * q_mean\n            target_q = target_q - self.critic_bc_coef * bc_penalty\n            target_q = rewards_flat + not_done * self.discount * target_q\n\n        cq1 = self.critic_1(states, actions)\n        cq2 = self.critic_2(states, actions)\n        critic_loss = F.mse_loss(cq1, target_q) + F.mse_loss(cq2, target_q)\n        log_dict[\"critic_loss\"] = critic_loss.item()\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            q = self.critic_1(states, pi)\n            bc_mse = ((pi - actions) ** 2).sum(-1)\n            # ACBC: Advantage-conditioned BC weights (enabled after warmup)\n            if self.total_it > self.acbc_warmup:\n                with torch.no_grad():\n                    q_data = torch.min(\n                        self.critic_1(states, actions),\n                        self.critic_2(states, actions))\n                    adv = q_data - q_data.mean()\n                    adv_scale = adv.abs().mean() + 1e-8\n                    bc_w = 1.0 + self.acbc_scale * torch.tanh(adv / adv_scale)\n            else:\n                bc_w = 1.0\n            # Momentum Q-normalization\n            qabs = torch.abs(q).mean().detach().item()\n            self.ema_q_abs = (1 - self.ema_rate) * self.ema_q_abs + self.ema_rate * qabs\n            lmbda = 1.0 / (self.ema_q_abs + 1e-8)\n            actor_loss = (self.actor_bc_coef * bc_w * bc_mse - lmbda * q).mean()\n            log_dict[\"actor_loss\"] = actor_loss.item()\n            self.actor_optimizer.zero_grad()\n            actor_loss.backward()\n            self.actor_optimizer.step()\n            soft_updat\n... (truncated)"
7}

1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom.py",
4  "start_line": 193,
5  "end_line": 429,
6  "content": "import sys as _sys\n\n\ndef _detect_env():\n    \"\"\"Parse --env from sys.argv.\"\"\"\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\n_ENV_NAME = _detect_env()\n\n# Per-environment hyperparameters (from ReBRAC CORL reference configs)\n_HPARAMS = {\n    \"halfcheetah-medium-v2\": {\"actor_bc\": 0.001, \"critic_bc\": 0.01,  \"lr\": 1e-3},\n    \"walker2d-medium-v2\":    {\"actor_bc\": 0.05,  \"critic_bc\": 0.1,   \"lr\": 1e-3},\n    \"hopper-medium-v2\":      {\"actor_bc\": 0.01,  \"critic_bc\": 0.01,  \"lr\": 1e-3},\n    \"maze2d-large-v1\":       {\"actor_bc\": 0.003, \"critic_bc\": 0.001, \"lr\": 3e-4},\n    \"maze2d-medium-v1\":      {\"actor_bc\": 0.003, \"critic_bc\": 0.001, \"lr\": 3e-4},\n    \"maze2d-umaze-v1\":       {\"actor_bc\": 0.003, \"critic_bc\": 0.001, \"lr\": 3e-4},\n}\n_HP = _HPARAMS.get(_ENV_NAME, {\"actor_bc\": 0.01, \"critic_bc\": 0.01, \"lr\": 1e-3})\n\nCONFIG_OVERRIDES: Dict[str, Any] = {}\n\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy: 3x256 MLP -> tanh * max_action.\"\"\"\n\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        s = torch.tensor(state.reshape(1, -1), device=device, dtype=torch.float32)\n        return self(s).cpu().data.numpy().flatten()\n\n\nclass Actor(nn.Module):\n    \"\"\"Stochastic Tanh-Gaussian policy: 3x256 MLP (kept for compatibility).\"\"\"\n\n    def __init__(self, state_dim: int, action_dim: int, max_action: float,\n                 orthogonal_init: bool = False):\n        super().__init__()\n        self.max_action = max_action\n        self.action_dim = action_dim\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, 2 * action_dim),\n        )\n        init_module_weights(self.net, orthogonal_init)\n        self.log_std_min, self.log_std_max = -20.0, 2.0\n\n    def _get_dist(self, state):\n        out = self.net(state)\n        mean, log_std = torch.split(out, self.action_dim, dim=-1)\n        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)\n        return TransformedDistribution(\n            Normal(mean, torch.exp(log_std)), TanhTransform(cache_size=1)), mean\n\n    def forward(self, state, deterministic=False):\n        dist, mean = self._get_dist(state)\n        action = torch.tanh(mean) if deterministic else dist.rsample()\n        return self.max_action * action, dist.log_prob(action).sum(-1)\n\n    @torch.no_grad()\n    def act(self, state, device=\"cpu\"):\n        s = torch.tensor(state.reshape(1, -1), device=device, dtype=torch.float32)\n        a, _ = self(s, not self.training)\n        return a.cpu().data.numpy().flatten()\n\n\nclass Critic(nn.Module):\n    \"\"\"Q(s,a): 3x256 MLP with LayerNorm for stable offline training.\"\"\"\n\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, action):\n        return self.net(torch.cat([state, action], dim=-1)).squeeze(-1)\n\n\nclass ValueFunction(nn.Module):\n    \"\"\"V(s): 3x256 MLP (kept for compatibility).\"\"\"\n\n    def __init__(self, state_dim: int, orthogonal_init: bool = False):\n        super().__init__()\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, 1),\n        )\n        init_module_weights(self.net, orthogonal_init)\n\n    def forward(self, state):\n        return self.net(state).squeeze(-1)\n\n\nclass OfflineAlgorithm:\n    \"\"\"ReBRAC with Data-Anchored Bellman Target and Temporal BC (DABT-TBC).\n\n    Novel contributions over standard ReBRAC:\n\n    1. Data-Anchored Bellman Target (DABT): Blends the policy-based critic\n       target with a SARSA-style target evaluated at the dataset's actual\n       next action. Since data actions are always in-distribution, this\n       provides a stable anchor that naturally corrects the underestimation\n       caused by the BC penalty on out-of-distribution policy actions:\n         target = (1-alpha) * [min(Q1,Q2)(s',a_pol) - bc_pen]\n                  + alpha * min(Q1,Q2)(s', a_data')\n       When the policy deviates from data (high BC penalty), the data\n       anchor dominates, preventing excessive pessimism.\n\n    2. Temporally-Extended BC (TEBC): Distributes the actor's BC penalty\n       across current AND next state-action pairs from the dataset:\n         bc = (1-beta)*||pi(s)-a||^2 + beta*(1-done)*||pi(s')-a'||^2\n       This provides denser coverage of the data distribution and temporal\n       consistency without changing the total regularization strength.\n    \"\"\"\n\n    def __init__(self, state_dim, action_dim, max_action, replay_buffer=None,\n                 discount=0.99, tau=5e-3, actor_lr=3e-4, critic_lr=3e-4,\n                 alpha_lr=3e-4, orthogonal_init=True, device=\"cuda\"):\n        self.device = device\n        self.discount = discount\n        self.tau = tau\n        self.max_action = max_action\n        self.total_it = 0\n        self.replay_buffer = replay_buffer\n        # Per-env tuned ReBRAC hyperparameters\n        self.actor_bc_coef = _HP[\"actor_bc\"]\n        self.critic_bc_coef = _HP[\"critic_bc\"]\n        _lr = _HP[\"lr\"]\n        self.policy_noise = 0.2\n        self.noise_clip = 0.5\n        self.policy_freq = 2\n        self.normalize_q = True\n        # Novel: DABT — data-anchored Bellman target blend weight\n        self.dabt_alpha = 0.15\n        # Novel: TEBC — temporal BC weight for next-state term\n        self.tebc_beta = 0.2\n\n        # Networks (identical to ReBRAC)\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        self.critic_1 = Critic(state_dim, action_dim).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        self.critic_2 = Critic(state_dim, action_dim).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            next_actions_pol = (self.actor_target(next_states) + noise).clamp(\n                -self.max_action, self.max_action)\n            bc_penalty = ((next_actions_pol - next_actions_data) ** 2).sum(-1)\n\n            # Standard ReBRAC policy-based target\n            tq1_pol = self.critic_1_target(next_states, next_actions_pol)\n            tq2_pol = self.critic_2_target(next_states, next_actions_pol)\n            target_q_pol = torch.min(tq1_pol, tq2_pol) - self.critic_bc_coef * bc_penalty\n\n            # DABT: Data-anchored SARSA-style target (always in-distribution)\n            tq1_dat = self.critic_1_target(next_states, next_actions_data)\n            tq2_dat = self.critic_2_target(next_states, next_actions_data)\n            target_q_dat = torch.min(tq1_dat, tq2_dat)\n\n            # Blended target\n            alpha = self.dabt_alpha\n            target_q = (1.0 - alpha) * target_q_pol + alpha * target_q_dat\n            target_q = rewards_flat + not_done * self.discount * target_q\n\n        cq1 = self.critic_1(states, actions)\n        cq2 = self.critic_2(states, actions)\n        critic_loss = F.mse_loss(cq1, target_q) + F.mse_loss(cq2, 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            q = self.critic_1(states, pi)\n\n            # Current-state BC (standard)\n            bc_current = ((pi - actions) ** 2).sum(-1)\n\n            # TEBC: Next-state BC (temporal extension)\n            pi_next = self.actor(next_states)\n            bc_next = ((pi_next - next_actions_data) ** 2).sum(-1) * not_done\n\n            # Combined BC with maintained total weight\n            beta = self.tebc_beta\n            bc_total = (1.0 - beta) * bc_current + beta * bc_next\n\n            lmbda = 1.0\n            if self.normalize_q:\n                lmbda = 1.0 / (torch.abs(q).mean().detach() + 1e-8)\n\n            actor_loss = (self.actor_bc_coef * bc_total - 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 \n... (truncated)"
7}

1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom.py",
4  "start_line": 193,
5  "end_line": 440,
6  "content": "import sys as _sys\n\n\ndef _detect_env():\n    \"\"\"Parse --env from sys.argv.\"\"\"\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\n_ENV_NAME = _detect_env()\n\n# Per-environment hyperparameters (from ReBRAC CORL reference configs)\n_HPARAMS = {\n    \"halfcheetah-medium-v2\": {\"actor_bc\": 0.001, \"critic_bc\": 0.01,  \"lr\": 1e-3},\n    \"walker2d-medium-v2\":    {\"actor_bc\": 0.05,  \"critic_bc\": 0.1,   \"lr\": 1e-3},\n    \"hopper-medium-v2\":      {\"actor_bc\": 0.01,  \"critic_bc\": 0.01,  \"lr\": 1e-3},\n    \"maze2d-large-v1\":       {\"actor_bc\": 0.003, \"critic_bc\": 0.001, \"lr\": 3e-4},\n    \"maze2d-medium-v1\":      {\"actor_bc\": 0.003, \"critic_bc\": 0.001, \"lr\": 3e-4},\n    \"maze2d-umaze-v1\":       {\"actor_bc\": 0.003, \"critic_bc\": 0.001, \"lr\": 3e-4},\n}\n_HP = _HPARAMS.get(_ENV_NAME, {\"actor_bc\": 0.01, \"critic_bc\": 0.01, \"lr\": 1e-3})\n\nCONFIG_OVERRIDES: Dict[str, Any] = {}\n\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy: 3x256 MLP -> tanh * max_action.\"\"\"\n\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        s = torch.tensor(state.reshape(1, -1), device=device, dtype=torch.float32)\n        return self(s).cpu().data.numpy().flatten()\n\n\nclass Actor(nn.Module):\n    \"\"\"Stochastic Tanh-Gaussian policy: 3x256 MLP (kept for compatibility).\"\"\"\n\n    def __init__(self, state_dim: int, action_dim: int, max_action: float,\n                 orthogonal_init: bool = False):\n        super().__init__()\n        self.max_action = max_action\n        self.action_dim = action_dim\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, 2 * action_dim),\n        )\n        init_module_weights(self.net, orthogonal_init)\n        self.log_std_min, self.log_std_max = -20.0, 2.0\n\n    def _get_dist(self, state):\n        out = self.net(state)\n        mean, log_std = torch.split(out, self.action_dim, dim=-1)\n        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)\n        return TransformedDistribution(\n            Normal(mean, torch.exp(log_std)), TanhTransform(cache_size=1)), mean\n\n    def forward(self, state, deterministic=False):\n        dist, mean = self._get_dist(state)\n        action = torch.tanh(mean) if deterministic else dist.rsample()\n        return self.max_action * action, dist.log_prob(action).sum(-1)\n\n    @torch.no_grad()\n    def act(self, state, device=\"cpu\"):\n        s = torch.tensor(state.reshape(1, -1), device=device, dtype=torch.float32)\n        a, _ = self(s, not self.training)\n        return a.cpu().data.numpy().flatten()\n\n\nclass Critic(nn.Module):\n    \"\"\"Q(s,a): 3x256 MLP with LayerNorm for stable offline training.\"\"\"\n\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, action):\n        return self.net(torch.cat([state, action], dim=-1)).squeeze(-1)\n\n\nclass ValueFunction(nn.Module):\n    \"\"\"V(s): 3x256 MLP (kept for compatibility).\"\"\"\n\n    def __init__(self, state_dim: int, orthogonal_init: bool = False):\n        super().__init__()\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, 1),\n        )\n        init_module_weights(self.net, orthogonal_init)\n\n    def forward(self, state):\n        return self.net(state).squeeze(-1)\n\n\nclass OfflineAlgorithm:\n    \"\"\"Ensemble-Stabilized ReBRAC (ES-ReBRAC).\n\n    Novel contributions over standard ReBRAC:\n\n    1. Ensemble Policy Gradient (EPG): Instead of using a single critic\n       (Q1) for the actor gradient, uses the mean of both critics\n       q = (Q1(s,pi) + Q2(s,pi)) / 2. This reduces policy gradient\n       variance by averaging over the twin critics' independent\n       estimation errors, providing a more accurate gradient signal\n       for policy optimization. Unlike min(Q1,Q2) used in the target,\n       using mean for the actor avoids additional pessimism while still\n       benefiting from the ensemble.\n\n    2. Gradient Norm Stabilization (GNS): Applies adaptive gradient\n       clipping to actor and critic networks, preventing catastrophic\n       parameter updates from outlier mini-batches. Particularly\n       important for offline RL where distribution shift can cause\n       sudden large gradients that destabilize training.\n    \"\"\"\n\n    def __init__(self, state_dim, action_dim, max_action, replay_buffer=None,\n                 discount=0.99, tau=5e-3, actor_lr=3e-4, critic_lr=3e-4,\n                 alpha_lr=3e-4, orthogonal_init=True, device=\"cuda\"):\n        self.device = device\n        self.discount = discount\n        self.tau = tau\n        self.max_action = max_action\n        self.total_it = 0\n        self.replay_buffer = replay_buffer\n        # Per-env tuned ReBRAC hyperparameters\n        self.actor_bc_coef = _HP[\"actor_bc\"]\n        self.critic_bc_coef = _HP[\"critic_bc\"]\n        _lr = _HP[\"lr\"]\n        self.policy_noise = 0.2\n        self.noise_clip = 0.5\n        self.policy_freq = 2\n        self.normalize_q = True\n        # Novel: Gradient norm stabilization thresholds\n        self.actor_grad_clip = 2.0\n        self.critic_grad_clip = 10.0\n\n        # Networks (identical architecture to ReBRAC)\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        self.critic_1 = Critic(state_dim, action_dim).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        self.critic_2 = Critic(state_dim, action_dim).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 (identical to ReBRAC) ──\n        with torch.no_grad():\n            noise = (torch.randn_like(actions) * self.policy_noise).clamp(\n                -self.noise_clip, self.noise_clip)\n            next_actions = (self.actor_target(next_states) + noise).clamp(\n                -self.max_action, self.max_action)\n            bc_penalty = ((next_actions - next_actions_data) ** 2).sum(-1)\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)\n            target_q = target_q - 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        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        # GNS: Clip critic gradients to prevent destabilizing updates\n        torch.nn.utils.clip_grad_norm_(self.critic_1.parameters(), self.critic_grad_clip)\n        torch.nn.utils.clip_grad_norm_(self.critic_2.parameters(), self.critic_grad_clip)\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            # EPG: Ensemble policy gradient using both critics\n            q1 = self.critic_1(states, pi)\n            q2 = self.critic_2(states, pi)\n            q = (q1 + q2) * 0.5\n\n            bc_mse = ((pi - actions) ** 2).sum(-1)\n\n            lmbda = 1.0\n            if self.normalize_q:\n                lmbda = 1.0 / (torch.abs(q).mean().detach() + 1e-8)\n\n            actor_loss = (self.actor_bc_coef * bc_mse - lmbda * q).mean()\n            log_dict[\"actor_loss\"] = actor_loss.item()\n\n            self.actor_optimizer.zero_grad()\n            actor_loss.backward()\n            # GNS: Clip actor gradients for training stability\n            torch.nn.utils.clip_grad_norm_(self.actor.parameters(), self.actor_grad_clip)\n            self.actor_optimizer.step()\n\n            soft_update(self.critic_1_target, self.critic_1, self.tau)\n            soft_update(self.critic_2_target, self.critic_2, self.tau)\n            soft_update(self.actor_target, self.actor, self.tau)\n\n        return log_dict"
7}