1{
2  "op": "replace",
3  "filename": "d3rlpy/atari_offline/custom_atari.py",
4  "start_line": 220,
5  "end_line": 258,
6  "content": "class QNetwork(nn.Module):\n    \"\"\"Q-network with separate V, Q, and policy heads for Discrete Implicit Q-Learning.\"\"\"\n    def __init__(self, observation_shape, action_dim, feature_dim=512):\n        super().__init__()\n        self.encoder = NatureDQNEncoder(observation_shape[0], feature_dim)\n        self.v_head = nn.Linear(feature_dim, 1)\n        self.q_head = nn.Linear(feature_dim, action_dim)\n        self.policy_head = nn.Linear(feature_dim, action_dim)\n        self.action_dim = action_dim\n\n    def forward(self, obs):\n        \"\"\"Returns mean Q-values (B, action_dim) — compatible with FIXED evaluate.\"\"\"\n        return self.q_head(self.encoder(obs))\n\n    def forward_v(self, obs):\n        features = self.encoder(obs)\n        return self.v_head(features).squeeze(-1)  # (B,)\n\n    def forward_q(self, obs):\n        features = self.encoder(obs)\n        return self.q_head(features)  # (B, A)\n\n    def forward_policy(self, obs):\n        features = self.encoder(obs)\n        return self.policy_head(features)  # (B, A)\n\n\nclass OfflineAlgorithm:\n    \"\"\"Discrete Implicit Q-Learning (DIQL): expectile regression for V, TD for Q, advantage-weighted policy.\"\"\"\n    def __init__(self, observation_shape, action_dim, config, device, buffer):\n        from copy import deepcopy\n        self.device = device\n        self.config = config\n        self.buffer = buffer\n        self.action_dim = action_dim\n        \n        # DIQL hyperparameters (hardcoded as per constraint)\n        self.expectile_tau = 0.9  # for expectile regression on V\n        self.temperature = 10.0  # beta for advantage weighting\n        self.gamma = config.gamma\n        self.tau = config.tau  # soft update coefficient\n        \n        self.model = QNetwork(observation_shape, action_dim).to(device)\n        self.target_model = deepcopy(self.model)\n        # Freeze target network\n        for param in self.target_model.parameters():\n            param.requires_grad = False\n        \n        self._step = 0\n\n    def parameters(self):\n        return self.model.parameters()\n\n    def expectile_loss(self, diff, tau):\n        \"\"\"Expectile regression loss L = |τ - 1{diff < 0}| * diff^2.\"\"\"\n        weight = torch.where(diff < 0, 1 - tau, tau)\n        return (weight * (diff ** 2)).mean()\n\n    def train_step(self, obs, actions, rewards, next_obs, dones):\n        batch_size = obs.shape[0]\n        \n        # Compute V(s), Q(s,a), π(s) from current model\n        V = self.model.forward_v(obs)  # (B,)\n        Q = self.model.forward_q(obs)  # (B, A)\n        Q_selected = Q.gather(1, actions.unsqueeze(1)).squeeze(1)  # (B,)\n        logits = self.model.forward_policy(obs)  # (B, A)\n        log_probs = F.log_softmax(logits, dim=-1)\n        log_prob_actions = log_probs.gather(1, actions.unsqueeze(1)).squeeze(1)  # (B,)\n        \n        # Compute target V(s') from target network\n        with torch.no_grad():\n            next_V = self.target_model.forward_v(next_obs)  # (B,)\n            target_V = rewards + self.gamma * (1 - dones) * next_V  # (B,)\n        \n        # 1. Expectile regression loss for V\n        diff = V - target_V\n        v_loss = self.expectile_loss(diff, self.expectile_tau)\n        \n        # 2. TD loss for Q (using same target as V)\n        q_loss = F.mse_loss(Q_selected, target_V)\n        \n        # 3. Advantage-weighted policy loss\n        advantages = Q_selected - V  # (B,)\n        weights = torch.exp(self.temperature * advantages).detach()  # (B,)\n        policy_loss = -(weights * log_prob_actions).mean()\n        \n        # Total loss\n        loss = v_loss + q_loss + policy_loss\n        \n        # Logging\n        info = {\n            \"loss\": loss.item(),\n            \"v_loss\": v_loss.item(),\n            \"q_loss\": q_loss.item(),\n            \"policy_loss\": policy_loss.item(),\n            \"mean_V\": V.mean().item(),\n            \"mean_Q\": Q_selected.mean().item(),\n            \"mean_advantage\": advantages.mean().item(),\n        }\n        return loss, info\n\n    def after_gradient_step(self, optimizer):\n        \"\"\"Soft update target network.\"\"\"\n        self._step += 1\n        soft_update(self.target_model, self.model, self.tau)\n\n    def select_action(self, obs):\n        \"\"\"Select action greedily from policy head (softmax).\"\"\"\n        with torch.no_grad():\n            logits = self.model.forward_policy(obs)\n            return torch.argmax(logits, dim=-1).item()\n\n    def begin_episode(self):\n        pass\n\n    def observe(self, reward):\n        pass"
7}

1{
2  "op": "replace",
3  "filename": "d3rlpy/atari_offline/custom_atari.py",
4  "start_line": 220,
5  "end_line": 336,
6  "content": "class QNetwork(nn.Module):\n    \"\"\"Q-network with separate V, Q, and policy heads for Discrete Implicit Q-Learning.\"\"\"\n    def __init__(self, observation_shape, action_dim, feature_dim=512):\n        super().__init__()\n        self.encoder = NatureDQNEncoder(observation_shape[0], feature_dim)\n        self.v_head = nn.Linear(feature_dim, 1)\n        self.q_head = nn.Linear(feature_dim, action_dim)\n        self.policy_head = nn.Linear(feature_dim, action_dim)\n        self.action_dim = action_dim\n        \n        # Initialize heads with smaller weights for stability\n        for head in [self.v_head, self.q_head, self.policy_head]:\n            nn.init.orthogonal_(head.weight, gain=0.01)\n            nn.init.constant_(head.bias, 0.0)\n\n    def forward(self, obs):\n        \"\"\"Returns mean Q-values (B, action_dim) — compatible with FIXED evaluate.\"\"\"\n        return self.q_head(self.encoder(obs))\n\n    def forward_v(self, obs):\n        features = self.encoder(obs)\n        return self.v_head(features).squeeze(-1)  # (B,)\n\n    def forward_q(self, obs):\n        features = self.encoder(obs)\n        return self.q_head(features)  # (B, A)\n\n    def forward_policy(self, obs):\n        features = self.encoder(obs)\n        return self.policy_head(features)  # (B, A)\n\n\nclass OfflineAlgorithm:\n    \"\"\"Discrete Implicit Q-Learning (DIQL): expectile regression for V, TD for Q, advantage-weighted policy.\"\"\"\n    def __init__(self, observation_shape, action_dim, config, device, buffer):\n        from copy import deepcopy\n        self.device = device\n        self.config = config\n        self.buffer = buffer\n        self.action_dim = action_dim\n        \n        # DIQL hyperparameters (hardcoded as per constraint)\n        self.expectile_tau = 0.7  # for expectile regression on V\n        self.temperature = 1.0  # beta for advantage weighting\n        self.gamma = config.gamma\n        self.tau = config.tau  # soft update coefficient\n        self.clip_grad_norm = 1.0  # gradient clipping\n        \n        self.model = QNetwork(observation_shape, action_dim).to(device)\n        self.target_model = deepcopy(self.model)\n        # Freeze target network\n        for param in self.target_model.parameters():\n            param.requires_grad = False\n        \n        self._step = 0\n\n    def create_optimizer(self):\n        \"\"\"Create optimizer with weight decay.\"\"\"\n        return torch.optim.Adam(\n            self.model.parameters(),\n            lr=self.config.learning_rate,\n            eps=self.config.optim_eps,\n            weight_decay=1e-4\n        )\n\n    def parameters(self):\n        return self.model.parameters()\n\n    def expectile_loss(self, diff, tau):\n        \"\"\"Expectile regression loss L = |τ - 1{diff < 0}| * diff^2.\"\"\"\n        weight = torch.where(diff < 0, 1 - tau, tau)\n        return (weight * (diff ** 2)).mean()\n\n    def train_step(self, obs, actions, rewards, next_obs, dones):\n        batch_size = obs.shape[0]\n        \n        # Compute V(s), Q(s,a), π(s) from current model\n        V = self.model.forward_v(obs)  # (B,)\n        Q = self.model.forward_q(obs)  # (B, A)\n        Q_selected = Q.gather(1, actions.unsqueeze(1)).squeeze(1)  # (B,)\n        logits = self.model.forward_policy(obs)  # (B, A)\n        log_probs = F.log_softmax(logits, dim=-1)\n        log_prob_actions = log_probs.gather(1, actions.unsqueeze(1)).squeeze(1)  # (B,)\n        \n        # Compute target V(s') from target network\n        with torch.no_grad():\n            next_V = self.target_model.forward_v(next_obs)  # (B,)\n            target_V = rewards + self.gamma * (1 - dones) * next_V  # (B,)\n        \n        # 1. Expectile regression loss for V\n        diff = V - target_V\n        v_loss = self.expectile_loss(diff, self.expectile_tau)\n        \n        # 2. TD loss for Q (using same target as V)\n        q_loss = F.mse_loss(Q_selected, target_V)\n        \n        # 3. Advantage-weighted policy loss with clipping\n        advantages = Q_selected - V  # (B,)\n        # Clip advantages to prevent explosive weights\n        advantages_clipped = torch.clamp(advantages, min=-10.0, max=10.0)\n        weights = torch.exp(self.temperature * advantages_clipped).detach()  # (B,)\n        policy_loss = -(weights * log_prob_actions).mean()\n        \n        # Total loss\n        loss = v_loss + q_loss + policy_loss\n        \n        # Logging\n        info = {\n            \"loss\": loss.item(),\n            \"v_loss\": v_loss.item(),\n            \"q_loss\": q_loss.item(),\n            \"policy_loss\": policy_loss.item(),\n            \"mean_V\": V.mean().item(),\n            \"mean_Q\": Q_selected.mean().item(),\n            \"mean_advantage\": advantages.mean().item(),\n        }\n        return loss, info\n\n    def after_gradient_step(self, optimizer):\n        \"\"\"Soft update target network.\"\"\"\n        self._step += 1\n        soft_update(self.target_model, self.model, self.tau)\n\n    def select_action(self, obs):\n        \"\"\"Select action greedily from policy head (softmax).\"\"\"\n        with torch.no_grad():\n            logits = self.model.forward_policy(obs)\n            return torch.argmax(logits, dim=-1).item()\n\n    def begin_episode(self):\n        pass\n\n    def observe(self, reward):\n        pass"
7}

1{
2  "op": "replace",
3  "filename": "d3rlpy/atari_offline/custom_atari.py",
4  "start_line": 220,
5  "end_line": 353,
6  "content": "class QNetwork(nn.Module):\n    \"\"\"Q-network with separate V, Q, and policy heads for Discrete Conservative Implicit Q-Learning.\"\"\"\n    def __init__(self, observation_shape, action_dim, feature_dim=512):\n        super().__init__()\n        self.encoder = NatureDQNEncoder(observation_shape[0], feature_dim)\n        self.v_head = nn.Linear(feature_dim, 1)\n        self.q_head = nn.Linear(feature_dim, action_dim)\n        self.policy_head = nn.Linear(feature_dim, action_dim)\n        self.action_dim = action_dim\n        \n        # Initialize heads with smaller weights for stability\n        for head in [self.v_head, self.q_head, self.policy_head]:\n            nn.init.orthogonal_(head.weight, gain=0.01)\n            nn.init.constant_(head.bias, 0.0)\n\n    def forward(self, obs):\n        \"\"\"Returns mean Q-values (B, action_dim) — compatible with FIXED evaluate.\"\"\"\n        return self.q_head(self.encoder(obs))\n\n    def forward_v(self, obs):\n        features = self.encoder(obs)\n        return self.v_head(features).squeeze(-1)  # (B,)\n\n    def forward_q(self, obs):\n        features = self.encoder(obs)\n        return self.q_head(features)  # (B, A)\n\n    def forward_policy(self, obs):\n        features = self.encoder(obs)\n        return self.policy_head(features)  # (B, A)\n\n\nclass OfflineAlgorithm:\n    \"\"\"Discrete Conservative Implicit Q-Learning (DCIQL): expectile regression for V, TD for Q, CQL penalty, advantage-weighted policy.\"\"\"\n    def __init__(self, observation_shape, action_dim, config, device, buffer):\n        from copy import deepcopy\n        self.device = device\n        self.config = config\n        self.buffer = buffer\n        self.action_dim = action_dim\n        \n        # Hyperparameters (hardcoded as per constraint)\n        self.expectile_tau = 0.9  # for expectile regression on V (optimistic)\n        self.temperature = 5.0  # beta for advantage weighting\n        self.cql_alpha = 0.1  # weight for CQL penalty\n        self.gamma = config.gamma\n        self.tau = config.tau  # soft update coefficient\n        self.clip_grad_norm = 1.0  # gradient clipping\n        \n        self.model = QNetwork(observation_shape, action_dim).to(device)\n        self.target_model = deepcopy(self.model)\n        # Freeze target network\n        for param in self.target_model.parameters():\n            param.requires_grad = False\n        \n        self._step = 0\n\n    def create_optimizer(self):\n        \"\"\"Create optimizer with weight decay.\"\"\"\n        return torch.optim.Adam(\n            self.model.parameters(),\n            lr=self.config.learning_rate,\n            eps=self.config.optim_eps,\n            weight_decay=1e-4\n        )\n\n    def parameters(self):\n        return self.model.parameters()\n\n    def expectile_loss(self, diff, tau):\n        \"\"\"Expectile regression loss L = |τ - 1{diff < 0}| * diff^2.\"\"\"\n        weight = torch.where(diff < 0, 1 - tau, tau)\n        return (weight * (diff ** 2)).mean()\n\n    def train_step(self, obs, actions, rewards, next_obs, dones):\n        batch_size = obs.shape[0]\n        \n        # Compute V(s), Q(s,a), π(s) from current model\n        V = self.model.forward_v(obs)  # (B,)\n        Q = self.model.forward_q(obs)  # (B, A)\n        Q_selected = Q.gather(1, actions.unsqueeze(1)).squeeze(1)  # (B,)\n        logits = self.model.forward_policy(obs)  # (B, A)\n        log_probs = F.log_softmax(logits, dim=-1)\n        log_prob_actions = log_probs.gather(1, actions.unsqueeze(1)).squeeze(1)  # (B,)\n        \n        # Compute target V(s') from target network\n        with torch.no_grad():\n            next_V = self.target_model.forward_v(next_obs)  # (B,)\n            target_V = rewards + self.gamma * (1 - dones) * next_V  # (B,)\n        \n        # 1. Expectile regression loss for V\n        diff = V - target_V\n        v_loss = self.expectile_loss(diff, self.expectile_tau)\n        \n        # 2. TD loss for Q (using same target as V)\n        q_loss = F.mse_loss(Q_selected, target_V)\n        \n        # 3. CQL penalty: logsumexp(Q) - Q_selected\n        # logsumexp(Q) approximates max Q, encourages conservatism\n        cql_penalty = torch.logsumexp(Q, dim=-1) - Q_selected\n        cql_loss = self.cql_alpha * cql_penalty.mean()\n        \n        # 4. Advantage-weighted policy loss with clipping\n        advantages = Q_selected - V  # (B,)\n        # Clip advantages to prevent explosive weights\n        advantages_clipped = torch.clamp(advantages, min=-10.0, max=10.0)\n        weights = torch.exp(self.temperature * advantages_clipped).detach()  # (B,)\n        policy_loss = -(weights * log_prob_actions).mean()\n        \n        # Total loss\n        loss = v_loss + q_loss + cql_loss + policy_loss\n        \n        # Logging\n        info = {\n            \"loss\": loss.item(),\n            \"v_loss\": v_loss.item(),\n            \"q_loss\": q_loss.item(),\n            \"cql_loss\": cql_loss.item(),\n            \"policy_loss\": policy_loss.item(),\n            \"mean_V\": V.mean().item(),\n            \"mean_Q\": Q_selected.mean().item(),\n            \"mean_advantage\": advantages.mean().item(),\n        }\n        return loss, info\n\n    def after_gradient_step(self, optimizer):\n        \"\"\"Soft update target network.\"\"\"\n        self._step += 1\n        soft_update(self.target_model, self.model, self.tau)\n\n    def select_action(self, obs):\n        \"\"\"Select action greedily from Q-values (argmax Q).\"\"\"\n        with torch.no_grad():\n            Q = self.model.forward_q(obs)\n            return torch.argmax(Q, dim=-1).item()\n\n    def begin_episode(self):\n        pass\n\n    def observe(self, reward):\n        pass"
7}

1{
2  "n": 2
3}