1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom_adroit.py",
4  "start_line": 214,
5  "end_line": 416,
6  "content": "CONFIG_OVERRIDES: Dict[str, Any] = {}\n\n\nclass DeterministicActor(nn.Module):\n    \"\"\"Deterministic policy pi(s) = tanh(net(s)) * max_action.\n    Suitable for BC, TD3+BC style algorithms. Default: 2 x 256 MLP.\"\"\"\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    \"\"\"Tanh-Gaussian stochastic policy. Default: 3 x 256 MLP.\n    Suitable for CQL, IQL, AWAC style algorithms.\"\"\"\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 = -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        \"\"\"Log-probability of a dataset action under the current policy.\"\"\"\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). Default: 3 x 256 MLP.\"\"\"\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, 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 function V(s). Default: 3 x 256 MLP. Useful for IQL-style algorithms.\"\"\"\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: torch.Tensor) -> torch.Tensor:\n        return self.net(state).squeeze(-1)\n\n\nclass OfflineAlgorithm:\n    \"\"\"Expert-Guided Implicit Q-Learning (EGIQL) for narrow expert data.\n    \n    Combines IQL's expectile regression with advantage-weighted policy learning,\n    plus behavior cloning regularization and conservative Q regularization.\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        # Hyperparameters\n        self.iql_tau = 0.8          # expectile for V loss\n        self.beta = 3.0             # inverse temperature for advantage weighting\n        self.exp_adv_max = 100.0    # clamp advantage exponent\n        self.bc_coef = 0.1          # weight for BC MSE loss\n        self.cql_alpha = 0.1        # weight for Q conservative penalty\n        self.awac_lambda = 0.1      # temperature for AWAC-style weighting\n\n        # Actor (stochastic Tanh-Gaussian)\n        self.actor = Actor(state_dim, action_dim, max_action, orthogonal_init).to(device)\n        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)\n\n        # Twin Q-networks + targets\n        self.qf1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.qf2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.qf1_target = deepcopy(self.qf1)\n        self.qf2_target = deepcopy(self.qf2)\n        self.qf1_target.requires_grad_(False)\n        self.qf2_target.requires_grad_(False)\n        self.q_optimizer = torch.optim.Adam(\n            list(self.qf1.parameters()) + list(self.qf2.parameters()), lr=critic_lr\n        )\n\n        # Value function V(s)\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 _asymmetric_l2_loss(self, u: torch.Tensor, tau: float) -> torch.Tensor:\n        return torch.mean(torch.abs(tau - (u < 0).float()) * u ** 2)\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        dones = dones.squeeze(-1)\n        log_dict: Dict[str, float] = {}\n\n        # ---------- V update: expectile regression ----------\n        with torch.no_grad():\n            target_q1 = self.qf1_target(states, actions)\n            target_q2 = self.qf2_target(states, actions)\n            target_q = torch.min(target_q1, target_q2)\n        v = self.vf(states)\n        adv = target_q - v\n        v_loss = self._asymmetric_l2_loss(adv, self.iql_tau)\n        log_dict[\"value_loss\"] = v_loss.item()\n\n        self.v_optimizer.zero_grad()\n        v_loss.backward()\n        self.v_optimizer.step()\n\n        # ---------- Q update: Bellman with V(s') + conservative penalty ----------\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.qf1(states, actions)\n        q2 = self.qf2(states, actions)\n        q1_loss = F.mse_loss(q1, q_target)\n        q2_loss = F.mse_loss(q2, q_target)\n\n        # Conservative penalty: penalize Q for policy actions far from dataset actions\n        with torch.no_grad():\n            # Sample policy actions (no grad)\n            pi_actions, _ = self.actor(states, deterministic=False)\n        q1_pi = self.qf1(states, pi_actions.detach())\n        q2_pi = self.qf2(states, pi_actions.detach())\n        # Penalty based on action distance (MSE)\n        action_mse = F.mse_loss(pi_actions, actions, reduction='none').mean(dim=-1)\n        cql_penalty = (q1_pi + q2_pi) * action_mse.detach()\n        cql_loss = self.cql_alpha * cql_penalty.mean()\n\n        q_loss = (q1_loss + q2_loss) / 2.0 + cql_loss\n        log_dict[\"critic_loss\"] = q_loss.item()\n        log_dict[\"cql_loss\"] = cql_loss.item()\n\n        self.q_optimizer.zero_grad()\n        q_loss.backward()\n        self.q_optimizer.step()\n\n        # Target Q update\n        soft_update(self.qf1_target, self.qf1, self.tau)\n        soft_update(self.qf2_target, self.qf2, self.tau)\n\n        # ---------- Actor update: advantage-weighted regression + BC ----------\n        with torch.no_grad():\n            # Compute advantage using current Q and V\n            q1_val = self.qf1(states, actions)\n            q2_val = self.qf2(states, actions)\n            q_val = torch.min(q1_val, q2_val)\n            v_val = self.vf(states)\n            adv_detached = q_val - v_val\n            # AWAC-style weighting\n            weights = torch.exp(adv_detached / self.awac_lambda).clamp(max=self.exp_adv_max)\n\n        # Advantage-weighted log-likelihood\n        action_log_prob = self.actor.log_prob(states, actions)\n        awac_loss = -(action_log_prob * weights).mean()\n        log_dict[\"awac_loss\"] = awac_loss.item()\n\n        # BC MSE loss between deterministic policy action and dataset action\n        pi_actions_det, _ = self.actor(states, deterministic=True)\n        bc_loss = F.mse_loss(pi_actions_det, actions)\n        log_dict[\"bc_loss\"] = bc_loss.item()\n\n        actor_loss = awac_loss + self.bc_coef * bc_loss\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        return log_dict"
7}

1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom_adroit.py",
4  "start_line": 319,
5  "end_line": 459,
6  "content": "class OfflineAlgorithm:\n    \"\"\"IQL with Behavior Cloning regularization (IQL+BC) for narrow expert data.\n    \n    Combines IQL's expectile regression and advantage-weighted policy learning\n    with an additional BC loss to anchor to expert actions.\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        # Hyperparameters\n        self.iql_tau = 0.8          # expectile for V loss\n        self.beta = 3.0             # inverse temperature for advantage weighting\n        self.exp_adv_max = 100.0    # clamp advantage exponent\n        self.bc_coef = 0.1          # weight for BC MSE loss\n        self.grad_clip = 1.0        # gradient clipping norm\n\n        # Actor (stochastic Tanh-Gaussian)\n        self.actor = Actor(state_dim, action_dim, max_action, orthogonal_init).to(device)\n        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)\n        self.actor_lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n            self.actor_optimizer, T_max=int(1e6)\n        )\n\n        # Twin Q-networks + targets (shared optimizer as in IQL)\n        self.qf1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.qf2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.qf1_target = deepcopy(self.qf1)\n        self.qf2_target = deepcopy(self.qf2)\n        self.qf1_target.requires_grad_(False)\n        self.qf2_target.requires_grad_(False)\n        self.q_optimizer = torch.optim.Adam(\n            list(self.qf1.parameters()) + list(self.qf2.parameters()), lr=critic_lr\n        )\n\n        # Value function V(s)\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 _asymmetric_l2_loss(self, u: torch.Tensor, tau: float) -> torch.Tensor:\n        return torch.mean(torch.abs(tau - (u < 0).float()) * u ** 2)\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        dones = dones.squeeze(-1)\n        log_dict: Dict[str, float] = {}\n\n        # ---------- V update: expectile regression ----------\n        with torch.no_grad():\n            target_q1 = self.qf1_target(states, actions)\n            target_q2 = self.qf2_target(states, actions)\n            target_q = torch.min(target_q1, target_q2)\n        v = self.vf(states)\n        adv = target_q - v\n        v_loss = self._asymmetric_l2_loss(adv, self.iql_tau)\n        log_dict[\"value_loss\"] = v_loss.item()\n\n        self.v_optimizer.zero_grad()\n        v_loss.backward()\n        torch.nn.utils.clip_grad_norm_(self.vf.parameters(), self.grad_clip)\n        self.v_optimizer.step()\n\n        # ---------- Q update: Bellman with V(s') ----------\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.qf1(states, actions)\n        q2 = self.qf2(states, actions)\n        q1_loss = F.mse_loss(q1, q_target)\n        q2_loss = F.mse_loss(q2, q_target)\n        q_loss = (q1_loss + q2_loss) / 2.0\n        log_dict[\"critic_loss\"] = q_loss.item()\n\n        self.q_optimizer.zero_grad()\n        q_loss.backward()\n        torch.nn.utils.clip_grad_norm_(list(self.qf1.parameters()) + list(self.qf2.parameters()), self.grad_clip)\n        self.q_optimizer.step()\n\n        # Target Q update\n        soft_update(self.qf1_target, self.qf1, self.tau)\n        soft_update(self.qf2_target, self.qf2, self.tau)\n\n        # ---------- Actor update: advantage-weighted regression + BC ----------\n        with torch.no_grad():\n            # Compute advantage using target Q and V (as in IQL)\n            adv_detached = target_q - self.vf(states)\n            weights = torch.exp(self.beta * adv_detached).clamp(max=self.exp_adv_max)\n\n        # Advantage-weighted log-likelihood\n        action_log_prob = self.actor.log_prob(states, actions)\n        awac_loss = -(action_log_prob * weights).mean()\n        log_dict[\"awac_loss\"] = awac_loss.item()\n\n        # BC MSE loss between deterministic policy action and dataset action\n        pi_actions_det, _ = self.actor(states, deterministic=True)\n        bc_loss = F.mse_loss(pi_actions_det, actions)\n        log_dict[\"bc_loss\"] = bc_loss.item()\n\n        actor_loss = awac_loss + self.bc_coef * bc_loss\n        log_dict[\"actor_loss\"] = actor_loss.item()\n\n        self.actor_optimizer.zero_grad()\n        actor_loss.backward()\n        torch.nn.utils.clip_grad_norm_(self.actor.parameters(), self.grad_clip)\n        self.actor_optimizer.step()\n        self.actor_lr_scheduler.step()\n\n        return log_dict"
7}

1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom_adroit.py",
4  "start_line": 319,
5  "end_line": 447,
6  "content": "class OfflineAlgorithm:\n    \"\"\"IQL with Behavior Cloning regularization (IQL+BC) for narrow expert data.\n    \n    Combines IQL's expectile regression and advantage-weighted policy learning\n    with an additional BC loss to anchor to expert actions.\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        # Detect environment for per-env hyperparameters\n        env_name = os.environ.get(\"ENV\", \"\")\n        \n        # Default hyperparameters (pen-human-v1)\n        self.iql_tau = 0.8          # expectile for V loss\n        self.beta = 3.0             # inverse temperature for advantage weighting\n        self.exp_adv_max = 100.0    # clamp advantage exponent\n        self.bc_coef = 0.1          # weight for BC MSE loss\n        self.grad_clip = 1.0        # gradient clipping norm\n        \n        if \"hammer\" in env_name:\n            # Hammer requires less BC constraint to allow exploration of nailing\n            self.iql_tau = 0.7      # more optimistic expectile\n            self.beta = 5.0         # more selective advantage weighting\n            self.bc_coef = 0.01     # lighter BC\n        elif \"door\" in env_name:\n            # Door-cloned data may need stronger BC due to suboptimal demonstrations\n            self.iql_tau = 0.9      # more conservative expectile\n            self.beta = 2.0         # broader advantage weighting\n            self.bc_coef = 0.2      # stronger BC\n        # else pen defaults\n\n        # Actor (stochastic Tanh-Gaussian)\n        self.actor = Actor(state_dim, action_dim, max_action, orthogonal_init).to(device)\n        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)\n        self.actor_lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n            self.actor_optimizer, T_max=int(1e6)\n        )\n\n        # Twin Q-networks + targets (shared optimizer as in IQL)\n        self.qf1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.qf2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.qf1_target = deepcopy(self.qf1)\n        self.qf2_target = deepcopy(self.qf2)\n        self.qf1_target.requires_grad_(False)\n        self.qf2_target.requires_grad_(False)\n        self.q_optimizer = torch.optim.Adam(\n            list(self.qf1.parameters()) + list(self.qf2.parameters()), lr=critic_lr\n        )\n\n        # Value function V(s)\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 _asymmetric_l2_loss(self, u: torch.Tensor, tau: float) -> torch.Tensor:\n        return torch.mean(torch.abs(tau - (u < 0).float()) * u ** 2)\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        dones = dones.squeeze(-1)\n        log_dict: Dict[str, float] = {}\n\n        # ---------- V update: expectile regression ----------\n        with torch.no_grad():\n            target_q1 = self.qf1_target(states, actions)\n            target_q2 = self.qf2_target(states, actions)\n            target_q = torch.min(target_q1, target_q2)\n        v = self.vf(states)\n        adv = target_q - v\n        v_loss = self._asymmetric_l2_loss(adv, self.iql_tau)\n        log_dict[\"value_loss\"] = v_loss.item()\n\n        self.v_optimizer.zero_grad()\n        v_loss.backward()\n        torch.nn.utils.clip_grad_norm_(self.vf.parameters(), self.grad_clip)\n        self.v_optimizer.step()\n\n        # ---------- Q update: Bellman with V(s') ----------\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.qf1(states, actions)\n        q2 = self.qf2(states, actions)\n        q1_loss = F.mse_loss(q1, q_target)\n        q2_loss = F.mse_loss(q2, q_target)\n        q_loss = (q1_loss + q2_loss) / 2.0\n        log_dict[\"critic_loss\"] = q_loss.item()\n\n        self.q_optimizer.zero_grad()\n        q_loss.backward()\n        torch.nn.utils.clip_grad_norm_(list(self.qf1.parameters()) + list(self.qf2.parameters()), self.grad_clip)\n        self.q_optimizer.step()\n\n        # Target Q update\n        soft_update(self.qf1_target, self.qf1, self.tau)\n        soft_update(self.qf2_target, self.qf2, self.tau)\n\n        # ---------- Actor update: advantage-weighted regression + BC ----------\n        with torch.no_grad():\n            # Compute advantage using target Q and V (as in IQL)\n            adv_detached = target_q - self.vf(states)\n            weights = torch.exp(self.beta * adv_detached).clamp(max=self.exp_adv_max)\n\n        # Advantage-weighted log-likelihood\n        action_log_prob = self.actor.log_prob(states, actions)\n        awac_loss = -(action_log_prob * weights).mean()\n        log_dict[\"awac_loss\"] = awac_loss.item()\n\n        # BC MSE loss between deterministic policy action and dataset action\n        pi_actions_det, _ = self.actor(states, deterministic=True)\n        bc_loss = F.mse_loss(pi_actions_det, actions)\n        log_dict[\"bc_loss\"] = bc_loss.item()\n\n        actor_loss = awac_loss + self.bc_coef * bc_loss\n        log_dict[\"actor_loss\"] = actor_loss.item()\n\n        self.actor_optimizer.zero_grad()\n        actor_loss.backward()\n        torch.nn.utils.clip_grad_norm_(self.actor.parameters(), self.grad_clip)\n        self.actor_optimizer.step()\n        self.actor_lr_scheduler.step()\n\n        return log_dict"
7}

1{
2  "op": "replace",
3  "filename": "CORL/algorithms/offline/custom_adroit.py",
4  "start_line": 319,
5  "end_line": 462,
6  "content": "class OfflineAlgorithm:\n    \"\"\"IQL with Behavior Cloning regularization (IQL+BC) for narrow expert data.\n    \n    Combines IQL's expectile regression and advantage-weighted policy learning\n    with an additional BC loss to anchor to expert actions.\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        # Detect environment for per-env hyperparameters\n        env_name = os.environ.get(\"ENV\", \"\")\n        \n        # Default hyperparameters (pen-human-v1)\n        self.iql_tau = 0.8          # expectile for V loss\n        self.beta = 3.0             # inverse temperature for advantage weighting\n        self.exp_adv_max = 100.0    # clamp advantage exponent\n        self.bc_coef = 0.1          # weight for BC MSE loss\n        self.critic_bc_coef = 0.0   # weight for critic BC penalty (ReBRAC style)\n        self.grad_clip = 1.0        # gradient clipping norm\n        \n        if \"hammer\" in env_name:\n            # Hammer requires less BC constraint to allow exploration of nailing\n            self.iql_tau = 0.7      # more optimistic expectile\n            self.beta = 5.0         # more selective advantage weighting\n            self.bc_coef = 0.01     # lighter BC\n            self.critic_bc_coef = 0.5  # strong critic BC penalty as in ReBRAC\n        elif \"door\" in env_name:\n            # Door-cloned data may need stronger BC due to suboptimal demonstrations\n            self.iql_tau = 0.9      # more conservative expectile\n            self.beta = 2.0         # broader advantage weighting\n            self.bc_coef = 0.2      # stronger BC\n            self.critic_bc_coef = 0.1  # moderate critic BC penalty\n        else:  # pen\n            self.critic_bc_coef = 0.5  # as in ReBRAC\n\n        # Actor (stochastic Tanh-Gaussian)\n        self.actor = Actor(state_dim, action_dim, max_action, orthogonal_init).to(device)\n        self.actor_target = deepcopy(self.actor)\n        self.actor_target.requires_grad_(False)\n        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)\n        self.actor_lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n            self.actor_optimizer, T_max=int(1e6)\n        )\n\n        # Twin Q-networks + targets (shared optimizer as in IQL)\n        self.qf1 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.qf2 = Critic(state_dim, action_dim, orthogonal_init).to(device)\n        self.qf1_target = deepcopy(self.qf1)\n        self.qf2_target = deepcopy(self.qf2)\n        self.qf1_target.requires_grad_(False)\n        self.qf2_target.requires_grad_(False)\n        self.q_optimizer = torch.optim.Adam(\n            list(self.qf1.parameters()) + list(self.qf2.parameters()), lr=critic_lr\n        )\n\n        # Value function V(s)\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 _asymmetric_l2_loss(self, u: torch.Tensor, tau: float) -> torch.Tensor:\n        return torch.mean(torch.abs(tau - (u < 0).float()) * u ** 2)\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        dones = dones.squeeze(-1)\n        log_dict: Dict[str, float] = {}\n\n        # ---------- V update: expectile regression ----------\n        with torch.no_grad():\n            target_q1 = self.qf1_target(states, actions)\n            target_q2 = self.qf2_target(states, actions)\n            target_q = torch.min(target_q1, target_q2)\n        v = self.vf(states)\n        adv = target_q - v\n        v_loss = self._asymmetric_l2_loss(adv, self.iql_tau)\n        log_dict[\"value_loss\"] = v_loss.item()\n\n        self.v_optimizer.zero_grad()\n        v_loss.backward()\n        torch.nn.utils.clip_grad_norm_(self.vf.parameters(), self.grad_clip)\n        self.v_optimizer.step()\n\n        # ---------- Q update: Bellman with V(s') + critic BC penalty ----------\n        with torch.no_grad():\n            next_v = self.vf(next_states)\n            # Compute next actions from target actor (deterministic)\n            next_actions_policy, _ = self.actor_target(next_states, deterministic=True)\n            # Critic BC penalty (ReBRAC style)\n            bc_penalty = ((next_actions_policy - next_actions) ** 2).sum(dim=-1)\n            next_q = next_v - self.critic_bc_coef * bc_penalty\n            q_target = rewards + (1.0 - dones) * self.discount * next_q\n\n        q1 = self.qf1(states, actions)\n        q2 = self.qf2(states, actions)\n        q1_loss = F.mse_loss(q1, q_target)\n        q2_loss = F.mse_loss(q2, q_target)\n        q_loss = (q1_loss + q2_loss) / 2.0\n        log_dict[\"critic_loss\"] = q_loss.item()\n        log_dict[\"critic_bc_penalty\"] = bc_penalty.mean().item()\n\n        self.q_optimizer.zero_grad()\n        q_loss.backward()\n        torch.nn.utils.clip_grad_norm_(list(self.qf1.parameters()) + list(self.qf2.parameters()), self.grad_clip)\n        self.q_optimizer.step()\n\n        # Target Q update\n        soft_update(self.qf1_target, self.qf1, self.tau)\n        soft_update(self.qf2_target, self.qf2, self.tau)\n        # Target actor update\n        soft_update(self.actor_target, self.actor, self.tau)\n\n        # ---------- Actor update: advantage-weighted regression + BC ----------\n        with torch.no_grad():\n            # Compute advantage using target Q and V (as in IQL)\n            adv_detached = target_q - self.vf(states)\n            weights = torch.exp(self.beta * adv_detached).clamp(max=self.exp_adv_max)\n\n        # Advantage-weighted log-likelihood\n        action_log_prob = self.actor.log_prob(states, actions)\n        awac_loss = -(action_log_prob * weights).mean()\n        log_dict[\"awac_loss\"] = awac_loss.item()\n\n        # BC MSE loss between deterministic policy action and dataset action\n        pi_actions_det, _ = self.actor(states, deterministic=True)\n        bc_loss = F.mse_loss(pi_actions_det, actions)\n        log_dict[\"bc_loss\"] = bc_loss.item()\n\n        actor_loss = awac_loss + self.bc_coef * bc_loss\n        log_dict[\"actor_loss\"] = actor_loss.item()\n\n        self.actor_optimizer.zero_grad()\n        actor_loss.backward()\n        torch.nn.utils.clip_grad_norm_(self.actor.parameters(), self.grad_clip)\n        self.actor_optimizer.step()\n        self.actor_lr_scheduler.step()\n\n        return log_dict"
7}