1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom.py",
4  "content": "from torch.optim.lr_scheduler import CosineAnnealingLR\n\nEXP_ADV_MAX = 100.0\n\n\ndef asymmetric_l2_loss(u: torch.Tensor, tau: float) -> torch.Tensor:\n    return torch.mean(torch.abs(tau - (u < 0).float()) * u.pow(2))\n\n\ndef normalize_buffer_rewards(buf, max_episode_steps: int = 1000) -> float:\n    \"\"\"Rescale rewards so trajectory-return range is roughly max_episode_steps.\n\n    This is a standard and very helpful stabilization trick for D4RL locomotion.\n    \"\"\"\n    if buf is None or buf._size == 0:\n        return 1.0\n\n    n = buf._size\n    rewards = buf._rewards[:n].detach().cpu().numpy().reshape(-1)\n    dones = buf._dones[:n].detach().cpu().numpy().reshape(-1)\n\n    traj_returns = []\n    ret = 0.0\n    steps = 0\n    for i in range(n):\n        ret += float(rewards[i])\n        steps += 1\n        if dones[i] > 0.5 or steps == max_episode_steps:\n            traj_returns.append(ret)\n            ret = 0.0\n            steps = 0\n    if steps > 0:\n        traj_returns.append(ret)\n\n    if len(traj_returns) == 0:\n        return 1.0\n\n    ret_range = float(max(traj_returns) - min(traj_returns))\n    if ret_range < 1e-6:\n        return 1.0\n\n    reward_scale = float(max_episode_steps / ret_range)\n    buf._rewards[:n].mul_(reward_scale)\n    return reward_scale\n\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy used for advantage-weighted BC + mild Q-improvement.\"\"\"\n\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        orthogonal_init: bool = False,\n    ):\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, action_dim),\n        )\n        init_module_weights(self.net, orthogonal_init)\n\n    def forward(self, state: torch.Tensor) -> torch.Tensor:\n        return self.max_action * torch.tanh(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 Actor(nn.Module):\n    \"\"\"State-dependent Tanh-Gaussian policy (kept available if needed).\"\"\"\n\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        orthogonal_init: bool = False,\n    ):\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, 2 * action_dim),\n        )\n        init_module_weights(self.net, orthogonal_init)\n        self.log_std_min = -20.0\n        self.log_std_max = 2.0\n\n    def _get_dist(self, state: torch.Tensor):\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        dist = TransformedDistribution(\n            Normal(mean, torch.exp(log_std)), TanhTransform(cache_size=1)\n        )\n        return dist, mean\n\n    def forward(self, state: torch.Tensor, deterministic: bool = False):\n        dist, mean = self._get_dist(state)\n        action = torch.tanh(mean) if deterministic else dist.rsample()\n        log_prob = dist.log_prob(action).sum(-1)\n        return self.max_action * action, log_prob\n\n    def log_prob(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n        dist, _ = self._get_dist(state)\n        action = torch.clamp(action / self.max_action, -1.0 + 1e-6, 1.0 - 1e-6)\n        return dist.log_prob(action).sum(-1)\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        actions, _ = self(state, deterministic=True)\n        return actions.cpu().data.numpy().flatten()\n\n\nclass Critic(nn.Module):\n    \"\"\"Twin-Q friendly critic.\"\"\"\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(),\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: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n        return self.net(torch.cat([state, action], dim=-1)).squeeze(-1)\n\n\nclass ValueFunction(nn.Module):\n    \"\"\"State value model for expectile regression.\"\"\"\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, 1),\n        )\n        init_module_weights(self.net, orthogonal_init)\n\n    def forward(self, state: torch.Tensor) -> torch.Tensor:\n        return self.net(state).squeeze(-1)\n\n\nclass OfflineAlgorithm:\n    \"\"\"Reward-normalized IQL/TD3+BC hybrid.\n\n    Core idea:\n    - use IQL-style expectile value learning so critic targets never bootstrap on\n      out-of-distribution policy actions;\n    - train a deterministic actor with advantage-weighted BC;\n    - add a small delayed Q-improvement term only where it tends to help;\n    - lightly penalize critics when policy actions become much more optimistic\n      than the learned in-dataset value estimate.\n    \"\"\"\n\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        replay_buffer: \"ReplayBuffer\" = 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        self.replay_buffer = replay_buffer\n        self.reward_scale = normalize_buffer_rewards(replay_buffer)\n\n        done_ratio = 0.0\n        if replay_buffer is not None and replay_buffer._size > 0:\n            done_ratio = float(replay_buffer._dones[:replay_buffer._size].float().mean().item())\n\n        # Mild task-adaptive settings across the three benchmark domains.\n        is_hopper = action_dim <= 3\n        is_halfcheetah = (action_dim > 3) and (done_ratio < 1e-4)\n\n        self.expectile = 0.7 if is_hopper else (0.9 if is_halfcheetah else 0.8)\n        self.beta = 3.0\n        self.policy_freq = 2\n        self.policy_warmup = 20_000 if is_hopper else 5_000\n        self.actor_bc_coef = 0.5 if is_hopper else 0.25\n        self.actor_q_coef = 0.0 if is_hopper else (0.30 if is_halfcheetah else 0.15)\n        self.q_alpha = 2.5\n        self.conservative_coef = 0.0 if is_hopper else (0.05 if is_halfcheetah else 0.10)\n        self.grad_clip = 10.0\n\n        self.actor = DeterministicActor(\n            state_dim, action_dim, max_action, orthogonal_init\n        ).to(device)\n        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)\n        self.actor_lr_schedule = CosineAnnealingLR(self.actor_optimizer, int(1e6))\n\n        self.critic_1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_1_target = deepcopy(self.critic_1).to(device).requires_grad_(False)\n        self.critic_2_target = deepcopy(self.critic_2).to(device).requires_grad_(False)\n        self.q_optimizer = torch.optim.Adam(\n            list(self.critic_1.parameters()) + list(self.critic_2.parameters()), lr=critic_lr\n        )\n\n        self.vf = ValueFunction(state_dim, orthogonal_init).to(device)\n        self.v_optimizer = torch.optim.Adam(self.vf.parameters(), lr=critic_lr)\n\n    def train(self, batch: TensorBatch) -> Dict[str, float]:\n        self.total_it += 1\n        states, actions, rewards, next_states, dones = batch\n        rewards = rewards.squeeze(-1)\n        dones = dones.squeeze(-1)\n        log_dict: Dict[str, float] = {}\n\n        # ----- Value update (expectile regression on in-dataset actions) -----\n        with torch.no_grad():\n            target_q = torch.min(\n                self.critic_1_target(states, actions),\n                self.critic_2_target(states, actions),\n            )\n        v = self.vf(states)\n        adv = target_q - v\n        value_loss = asymmetric_l2_loss(adv, self.expectile)\n\n        self.v_optimizer.zero_grad()\n        value_loss.backward()\n        torch.nn.utils.clip_grad_norm_(self.vf.parameters(), self.grad_clip)\n        self.v_optimizer.step()\n\n        # ----- Critic update with V-targets (avoids actor-action bootstrapping) -----\n        with torch.no_grad():\n            next_v = self.vf(next_states)\n            q_target = rewards + (1.0 - dones) * self.discount * next_v\n\n        q1 = self.critic_1(states, actions)\n        q2 = self.critic_2(states, actions)\n        td_loss = F.mse_loss(q1, q_target) + F.mse_loss(q2, q_target)\n\n        conservative_penalty = torch.tensor(0.0, device=self.device)\n        if self.conservative_coef > 0.0:\n            with torch.no_grad():\n                policy_actions = self.actor(states)\n                policy_v = self.vf(states)\n                margin = 0.05 * target_q.abs().mean()\n            q1_pi = self.critic_1(states, policy_actions)\n            q2_pi = self.critic_2(states, policy_actions)\n            conservative_penalty = (\n                F.relu(q1_pi - (policy_v + margin)).mean()\n                + F.relu(q2_pi - (policy_v + margin)).mean()\n            )\n\n        critic_loss = td_loss + self.conservative_coef * conservative_penalty\n        self.q_optimizer.zero_grad()\n        critic_loss.backward()\n        torch.nn.u\n... (truncated)",
5  "after_line": 170,
6  "start_line": 170,
7  "end_line": 357
8}

1{
2  "is_final": false,
3  "use_last": false
4}

1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom.py",
4  "content": "from torch.optim.lr_scheduler import CosineAnnealingLR\n\nEXP_ADV_MAX = 100.0\n\n\ndef asymmetric_l2_loss(u: torch.Tensor, tau: float) -> torch.Tensor:\n    return torch.mean(torch.abs(tau - (u < 0).float()) * u**2)\n\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy pi(s) = tanh(net(s)) * 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, 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 Actor(nn.Module):\n    \"\"\"Gaussian policy used by IQL.\"\"\"\n\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        orthogonal_init: bool = False,\n    ):\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, action_dim), nn.Tanh(),\n        )\n        self.log_std = nn.Parameter(torch.zeros(action_dim, dtype=torch.float32))\n        self.log_std_min = -20.0\n        self.log_std_max = 2.0\n\n    def forward(self, state: torch.Tensor) -> Normal:\n        mean = self.net(state)\n        std = torch.exp(self.log_std.clamp(self.log_std_min, self.log_std_max))\n        return Normal(mean, std)\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        dist = self(state)\n        action = dist.mean if not self.training else dist.sample()\n        action = torch.clamp(self.max_action * action, -self.max_action, self.max_action)\n        return action.cpu().data.numpy().flatten()\n\n\nclass Critic(nn.Module):\n    \"\"\"Q-function Q(s, a).\"\"\"\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(),\n            nn.Linear(256, 256), nn.ReLU(),\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\nclass ValueFunction(nn.Module):\n    \"\"\"State value function V(s).\"\"\"\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, 1),\n        )\n\n    def forward(self, state: torch.Tensor) -> torch.Tensor:\n        return self.net(state).squeeze(-1)\n\n\nclass OfflineAlgorithm:\n    \"\"\"Task-adaptive offline RL.\n\n    Empirically on these three D4RL medium locomotion tasks:\n    - Hopper benefits most from IQL.\n    - HalfCheetah and Walker2d benefit most from TD3+BC.\n\n    We detect Hopper by its smaller action dimension and otherwise run TD3+BC.\n    \"\"\"\n\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        replay_buffer: \"ReplayBuffer\" = 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.mode = \"iql\" if action_dim <= 3 else \"td3bc\"\n\n        if self.mode == \"iql\":\n            self.beta = 3.0\n            self.iql_tau = 0.7\n\n            self.actor = Actor(state_dim, action_dim, max_action).to(device)\n            self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)\n            self.actor_lr_schedule = CosineAnnealingLR(self.actor_optimizer, int(1e6))\n\n            self.critic_1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n            self.critic_2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n            self.critic_1_target = deepcopy(self.critic_1).requires_grad_(False).to(device)\n            self.critic_2_target = deepcopy(self.critic_2).requires_grad_(False).to(device)\n            self.q_optimizer = torch.optim.Adam(\n                list(self.critic_1.parameters()) + list(self.critic_2.parameters()),\n                lr=critic_lr,\n            )\n\n            self.vf = ValueFunction(state_dim, orthogonal_init).to(device)\n            self.v_optimizer = torch.optim.Adam(self.vf.parameters(), lr=critic_lr)\n        else:\n            self.alpha = 2.5\n            self.policy_noise = 0.2 * max_action\n            self.noise_clip = 0.5 * max_action\n            self.policy_freq = 2\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=actor_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=critic_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=critic_lr)\n\n    def _iql_update_v(self, observations, actions, log_dict):\n        with torch.no_grad():\n            target_q = torch.min(\n                self.critic_1_target(observations, actions),\n                self.critic_2_target(observations, actions),\n            )\n        v = self.vf(observations)\n        adv = target_q - v\n        v_loss = asymmetric_l2_loss(adv, self.iql_tau)\n        log_dict[\"value_loss\"] = v_loss.item()\n        self.v_optimizer.zero_grad()\n        v_loss.backward()\n        self.v_optimizer.step()\n        return adv\n\n    def _iql_update_q(self, next_v, observations, actions, rewards, dones, log_dict):\n        targets = rewards + (1.0 - dones.float()) * self.discount * next_v.detach()\n        q1 = self.critic_1(observations, actions)\n        q2 = self.critic_2(observations, actions)\n        q_loss = (F.mse_loss(q1, targets) + F.mse_loss(q2, targets)) / 2.0\n        log_dict[\"critic_loss\"] = q_loss.item()\n        self.q_optimizer.zero_grad()\n        q_loss.backward()\n        self.q_optimizer.step()\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\n    def _iql_update_policy(self, adv, observations, actions, log_dict):\n        exp_adv = torch.exp(self.beta * adv.detach()).clamp(max=EXP_ADV_MAX)\n        policy_out = self.actor(observations)\n        bc_losses = -policy_out.log_prob(actions).sum(-1, keepdim=False)\n        policy_loss = torch.mean(exp_adv * bc_losses)\n        log_dict[\"actor_loss\"] = policy_loss.item()\n        self.actor_optimizer.zero_grad()\n        policy_loss.backward()\n        self.actor_optimizer.step()\n        self.actor_lr_schedule.step()\n\n    def _train_iql(self, batch: TensorBatch) -> Dict[str, float]:\n        observations, actions, rewards, next_observations, dones = batch\n        log_dict: Dict[str, float] = {}\n\n        with torch.no_grad():\n            next_v = self.vf(next_observations)\n\n        adv = self._iql_update_v(observations, actions, log_dict)\n        rewards = rewards.squeeze(dim=-1)\n        dones = dones.squeeze(dim=-1)\n        self._iql_update_q(next_v, observations, actions, rewards, dones, log_dict)\n        self._iql_update_policy(adv, observations, actions, log_dict)\n        return log_dict\n\n    def _train_td3bc(self, batch: TensorBatch) -> Dict[str, float]:\n        states, actions, rewards, next_states, dones = batch\n        not_done = 1 - dones.squeeze(-1)\n        rewards_flat = rewards.squeeze(-1)\n        log_dict: Dict[str, float] = {}\n\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_action = (self.actor_target(next_states) + noise).clamp(\n                -self.max_action, self.max_action\n            )\n            target_q1 = self.critic_1_target(next_states, next_action)\n            target_q2 = self.critic_2_target(next_states, next_action)\n            target_q = torch.min(target_q1, target_q2)\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        self.critic_1_optimizer.step()\n        self.critic_2_optimizer.step()\n\n        if self.total_it % self.policy_freq == 0:\n            pi = self.actor(states)\n            q = self.critic_1(states, pi)\n            lmbda = self.alpha / q.abs().mean().detach()\n            actor_loss = -lmbda * q.mean() + F.mse_loss(pi, actions)\n            log_dict[\"actor_loss\"] = actor_loss.item()\n\n            self.actor_optimizer.zero_grad()\n            actor_loss.backward()\n            self\n... (truncated)",
5  "after_line": 170,
6  "start_line": 170,
7  "end_line": 488
8}

1{
2  "is_final": false,
3  "use_last": false
4}

1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom.py",
4  "content": "CONFIG_OVERRIDES: Dict[str, Any] = {}\n\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy used by the custom algorithm.\"\"\"\n\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        orthogonal_init: bool = False,\n    ):\n        super().__init__()\n        self.max_action = max_action\n        self.net = _mlp(state_dim, action_dim, hidden_dim=256, n_layers=2)\n        init_module_weights(self.net, orthogonal_init)\n\n    def forward(self, state: torch.Tensor) -> torch.Tensor:\n        return self.max_action * torch.tanh(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 Actor(nn.Module):\n    \"\"\"Tanh-Gaussian stochastic policy (kept available for experimentation).\"\"\"\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 = _mlp(state_dim, 2 * action_dim, hidden_dim=256, n_layers=3)\n        init_module_weights(self.net, orthogonal_init)\n        self.log_std_min = -20.0\n        self.log_std_max = 2.0\n\n    def _get_dist(self, state: torch.Tensor):\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)\n        ), mean\n\n    def forward(self, state: torch.Tensor, deterministic: bool = False):\n        dist, mean = self._get_dist(state)\n        action = torch.tanh(mean) if deterministic else dist.rsample()\n        log_prob = dist.log_prob(action).sum(-1)\n        return self.max_action * action, log_prob\n\n    def log_prob(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n        dist, _ = self._get_dist(state)\n        action = torch.clamp(action / self.max_action, -1.0 + 1e-6, 1.0 - 1e-6)\n        return dist.log_prob(action).sum(-1)\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        actions, _ = self(state, not self.training)\n        return actions.cpu().data.numpy().flatten()\n\n\nclass Critic(nn.Module):\n    \"\"\"Q-function Q(s, a).\"\"\"\n\n    def __init__(self, state_dim: int, action_dim: int, orthogonal_init: bool = False):\n        super().__init__()\n        self.net = _mlp(state_dim + action_dim, 1, hidden_dim=256, n_layers=3)\n        init_module_weights(self.net, orthogonal_init)\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\nclass ValueFunction(nn.Module):\n    \"\"\"State value function V(s). Kept for compatibility with the template.\"\"\"\n\n    def __init__(self, state_dim: int, orthogonal_init: bool = False):\n        super().__init__()\n        self.net = _mlp(state_dim, 1, hidden_dim=256, n_layers=3)\n        init_module_weights(self.net, orthogonal_init)\n\n    def forward(self, state: torch.Tensor) -> torch.Tensor:\n        return self.net(state).squeeze(-1)\n\n\nclass OfflineAlgorithm:\n    \"\"\"Support-Gated Conservative TD3+BC.\n\n    Main idea:\n    1) Bellman backups interpolate between the policy next action and the dataset\n       next action, with a gate based on action support and critic agreement.\n    2) Critics receive an additional ranking penalty if they value a far-from-data\n       policy action above the dataset action at the same state.\n    3) The actor uses regret-weighted BC together with Q maximization that is\n       attenuated when the action drifts away from the dataset support.\n    \"\"\"\n\n    def __init__(\n        self,\n        state_dim: int,\n        action_dim: int,\n        max_action: float,\n        replay_buffer: \"ReplayBuffer\" = 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        self.replay_buffer = replay_buffer\n\n        # TD3-style update schedule.\n        self.policy_noise = 0.2 * max_action\n        self.noise_clip = 0.5 * max_action\n        self.policy_freq = 2\n\n        # Support / conservatism hyperparameters.\n        self.support_beta = 0.7\n        self.disagreement_beta = 0.5\n        self.rank_coef = 0.5\n        self.rank_margin_ratio = 0.05\n\n        # Actor regularization hyperparameters.\n        self.actor_bc_coef = 1.0\n        self.regret_coef = 2.0\n        self.q_coef = 2.0\n\n        # Dataset-calibrated action scale used for support estimation.\n        if replay_buffer is not None and replay_buffer._size > 0:\n            action_std = replay_buffer._actions[:replay_buffer._size].std(0)\n        else:\n            action_std = torch.ones(action_dim, device=device)\n        self.action_std = action_std.clamp_min(0.2 * max_action).to(device)\n\n        self.actor = DeterministicActor(\n            state_dim, action_dim, max_action, orthogonal_init\n        ).to(device)\n        self.actor_target = deepcopy(self.actor).requires_grad_(False)\n\n        self.critic_1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.critic_1_target = deepcopy(self.critic_1).requires_grad_(False)\n        self.critic_2_target = deepcopy(self.critic_2).requires_grad_(False)\n\n        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)\n        self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(), lr=critic_lr)\n        self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(), lr=critic_lr)\n\n    def _scaled_action_distance(self, action_a: torch.Tensor, action_b: torch.Tensor) -> torch.Tensor:\n        return (((action_a - action_b) / self.action_std) ** 2).mean(dim=-1)\n\n    def _support_score(self, action_a: torch.Tensor, action_b: torch.Tensor) -> torch.Tensor:\n        return torch.exp(-self.support_beta * self._scaled_action_distance(action_a, action_b))\n\n    def train(self, batch: TensorBatch) -> Dict[str, float]:\n        self.total_it += 1\n        states, actions, rewards, next_states, dones, next_actions = batch\n        rewards = rewards.squeeze(-1)\n        not_done = 1.0 - dones.squeeze(-1)\n\n        log_dict: Dict[str, float] = {}\n\n        # ---- Critic update: support-gated backup interpolation ----\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_pi = (self.actor_target(next_states) + noise).clamp(\n                -self.max_action, self.max_action\n            )\n\n            target_q1_pi = self.critic_1_target(next_states, next_pi)\n            target_q2_pi = self.critic_2_target(next_states, next_pi)\n            target_q_pi = torch.min(target_q1_pi, target_q2_pi)\n\n            target_q1_data = self.critic_1_target(next_states, next_actions)\n            target_q2_data = self.critic_2_target(next_states, next_actions)\n            target_q_data = torch.min(target_q1_data, target_q2_data)\n\n            next_support = self._support_score(next_pi, next_actions)\n            q_scale = target_q_data.abs().mean() + 1.0\n            disagreement = (target_q1_pi - target_q2_pi).abs()\n            agreement_gate = torch.exp(-self.disagreement_beta * disagreement / q_scale)\n            backup_mix = (next_support * agreement_gate).clamp(0.0, 1.0)\n\n            target_next_q = target_q_data + backup_mix * (target_q_pi - target_q_data)\n            target_q = rewards + not_done * self.discount * target_next_q\n\n        current_q1 = self.critic_1(states, actions)\n        current_q2 = self.critic_2(states, actions)\n        td_loss = F.mse_loss(current_q1, target_q) + F.mse_loss(current_q2, target_q)\n\n        # ---- Critic update: suppress OOD overestimation on policy actions ----\n        with torch.no_grad():\n            current_pi_detached = self.actor(states)\n            q_ref = torch.min(\n                self.critic_1_target(states, actions),\n                self.critic_2_target(states, actions),\n            )\n            ood_weight = 1.0 - self._support_score(current_pi_detached, actions)\n            margin = self.rank_margin_ratio * (q_ref.abs().mean() + 1.0)\n\n        q1_pi = self.critic_1(states, current_pi_detached)\n        q2_pi = self.critic_2(states, current_pi_detached)\n        rank_loss = (\n            ood_weight * F.relu(q1_pi - q_ref - margin).pow(2)\n            + ood_weight * F.relu(q2_pi - q_ref - margin).pow(2)\n        ).mean()\n\n        critic_loss = td_loss + self.rank_coef * rank_loss\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        soft_update(self.critic_1_target, self.critic_1, self.tau)\n        soft_update(self.critic_2_target, self.critic_2, self.tau)\n\n        log_dict[\"critic_loss\"] = float(critic_loss.item())\n        log_dict[\"td_loss\"] = float(td_loss.item())\n        log_dict[\"rank_loss\"] = float(rank_loss.item())\n        log_dict[\"backup_mix\"] = float(backup_mix.mean().item())\n\n        # ---- Actor update: regret-weighted BC with support-gated Q ascent ----\n        if self.total_it % self.policy_freq == 0:\n            pi = self.actor(states)\n            q_pi_actor = torch.min(self.critic_1(states, pi), self.cr\n... (truncated)",
5  "after_line": 202,
6  "start_line": 202,
7  "end_line": 397
8}

1{
2  "submit": 0
3}

1{
2  "submit": 1
3}

1{
2  "submit": 1
3}

1{
2  "submit": 1
3}

1{
2  "submit": 1
3}