1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_value_atari.py",
4  "start_line": 186,
5  "end_line": 249,
6  "content": "class QNetwork(nn.Module):\n    \"\"\"Implicit Quantile Network (IQN) for distributional RL.\n    \n    Encodes quantile samples via cosine embedding and combines with state features.\n    \"\"\"\n    def __init__(self, envs, quantile_embed_dim=64):\n        super().__init__()\n        n_actions = envs.single_action_space.n\n        self.n_actions = n_actions\n        self.quantile_embed_dim = quantile_embed_dim\n        self.encoder = NatureDQNEncoder()\n        \n        # Cosine embedding for quantile samples\n        self.cos_embed = nn.Linear(quantile_embed_dim, 512)\n        # Main head: after elementwise product of features and quantile embedding\n        self.head = nn.Linear(512, n_actions)\n        \n    def forward(self, x, tau_samples=None):\n        \"\"\"Return Q-values (mean over quantiles) or quantile values if tau_samples provided.\n        \n        Args:\n            x: state tensor [batch, 4, 84, 84]\n            tau_samples: optional [batch, n_samples] quantile samples in (0,1)\n        \n        Returns:\n            If tau_samples is None: Q-values [batch, n_actions]\n            Else: quantile values [batch, n_samples, n_actions]\n        \"\"\"\n        features = self.encoder(x)  # [batch, 512]\n        \n        if tau_samples is None:\n            # For action selection: sample a single tau per state? Use expectation over tau.\n            # We'll sample 32 tau per state and average.\n            batch_size = x.shape[0]\n            tau = torch.rand(batch_size, 32, device=x.device)  # [batch, 32]\n            quantile_vals = self._quantile_values(features, tau)  # [batch, 32, n_actions]\n            q_values = quantile_vals.mean(dim=1)  # [batch, n_actions]\n            return q_values\n        else:\n            # tau_samples: [batch, n_samples]\n            quantile_vals = self._quantile_values(features, tau_samples)  # [batch, n_samples, n_actions]\n            return quantile_vals\n    \n    def _quantile_values(self, features, tau):\n        \"\"\"Compute quantile values given features and tau samples.\n        \n        Args:\n            features: [batch, 512]\n            tau: [batch, n_samples]\n        \n        Returns:\n            quantile values: [batch, n_samples, n_actions]\n        \"\"\"\n        batch_size, n_samples = tau.shape\n        # Compute cosine embedding of tau\n        # tau in (0,1), expand to [batch, n_samples, quantile_embed_dim]\n        pi = torch.tensor(torch.pi, device=tau.device)\n        # Create cosine basis: cos(pi * i * tau) for i=1..quantile_embed_dim\n        i_vals = torch.arange(1, self.quantile_embed_dim + 1, device=tau.device).float()\n        # i_vals: [quantile_embed_dim]\n        # tau: [batch, n_samples, 1]\n        tau = tau.unsqueeze(-1)  # [batch, n_samples, 1]\n        # cos(pi * i * tau): [batch, n_samples, quantile_embed_dim]\n        cos_vals = torch.cos(pi * i_vals * tau)\n        # Linear transform\n        cos_embed = self.cos_embed(cos_vals)  # [batch, n_samples, 512]\n        # Expand features to match\n        features_expanded = features.unsqueeze(1).expand(-1, n_samples, -1)  # [batch, n_samples, 512]\n        # Elementwise product\n        combined = features_expanded * cos_embed  # [batch, n_samples, 512]\n        # ReLU activation (as in IQN paper)\n        combined = F.relu(combined)\n        # Linear head\n        quantile_vals = self.head(combined)  # [batch, n_samples, n_actions]\n        return quantile_vals\n\n\nclass ValueAlgorithm:\n    \"\"\"Implicit Quantile Networks (IQN) with latent temporal consistency regularization.\"\"\"\n    \n    def __init__(self, envs, device, args):\n        self.device = device\n        self.gamma = args.gamma\n        self.tau = args.tau\n        self.target_network_frequency = args.target_network_frequency\n        self.n_quantile_samples = 8  # number of tau samples for loss\n        self.n_target_samples = 8\n        self.kappa = 1.0  # Huber loss threshold\n        self.lambda_latent = 0.1  # weight for latent consistency loss\n        \n        self.q_network = QNetwork(envs).to(device)\n        self.target_network = QNetwork(envs).to(device)\n        self.target_network.load_state_dict(self.q_network.state_dict())\n        self.optimizer = optim.Adam(self.q_network.parameters(), lr=args.learning_rate)\n        \n        # Small MLP for predicting next latent features\n        self.latent_predictor = nn.Sequential(\n            nn.Linear(ENCODER_FEATURE_DIM, 256),\n            nn.ReLU(),\n            nn.Linear(256, ENCODER_FEATURE_DIM)\n        ).to(device)\n        self.latent_optimizer = optim.Adam(self.latent_predictor.parameters(), lr=args.learning_rate)\n        \n    def select_action(self, obs, epsilon):\n        \"\"\"Epsilon-greedy action selection using Q-values.\"\"\"\n        if random.random() < epsilon:\n            return np.array([self.q_network.n_actions])\n        q_values = self.q_network(torch.Tensor(obs).to(self.device))\n        return torch.argmax(q_values, dim=1).cpu().numpy()\n    \n    def update(self, batch, global_step):\n        \"\"\"IQN update with latent temporal consistency regularization.\"\"\"\n        batch_size = batch.observations.shape[0]\n        \n        # Sample quantiles for current and target networks\n        tau = torch.rand(batch_size, self.n_quantile_samples, device=self.device)  # [batch, N]\n        tau_prime = torch.rand(batch_size, self.n_target_samples, device=self.device)  # [batch, N']\n        \n        with torch.no_grad():\n            # Compute target quantile values\n            next_quantiles = self.target_network._quantile_values(\n                self.target_network.encoder(batch.next_observations),\n                tau_prime\n            )  # [batch, N', n_actions]\n            next_q = next_quantiles.mean(dim=1)  # [batch, n_actions]\n            next_actions = next_q.argmax(dim=1)  # [batch]\n            # Select quantiles for best actions\n            next_quantiles_best = next_quantiles[torch.arange(batch_size), :, next_actions]  # [batch, N']\n            # Compute target quantile values\n            target_quantiles = batch.rewards + self.gamma * next_quantiles_best * (1 - batch.dones)\n            target_quantiles = target_quantiles.unsqueeze(1)  # [batch, 1, N'] for broadcast\n        \n        # Current quantile values for taken actions\n        current_quantiles_all = self.q_network._quantile_values(\n            self.q_network.encoder(batch.observations),\n            tau\n        )  # [batch, N, n_actions]\n        current_quantiles = current_quantiles_all[torch.arange(batch_size), :, batch.actions.flatten()]  # [batch, N]\n        \n        # Quantile Huber loss\n        # Pairwise TD errors: [batch, N, N']\n        td_errors = target_quantiles - current_quantiles.unsqueeze(2)\n        huber_loss = F.huber_loss(td_errors, torch.zeros_like(td_errors), reduction='none', delta=self.kappa)\n        quantile_loss = (huber_loss * (tau.unsqueeze(2) - (td_errors < 0).float()).abs()).mean()\n        \n        # Latent temporal consistency loss\n        features_current = self.q_network.encoder(batch.observations)\n        features_next = self.q_network.encoder(batch.next_observations)\n        predicted_next = self.latent_predictor(features_current)\n        latent_loss = F.mse_loss(predicted_next, features_next.detach())\n        \n        # Total loss\n        total_loss = quantile_loss + self.lambda_latent * latent_loss\n        \n        self.optimizer.zero_grad()\n        self.latent_optimizer.zero_grad()\n        total_loss.backward()\n        self.optimizer.step()\n        self.latent_optimizer.step()\n        \n        # Soft target network update\n        if global_step % self.target_network_frequency == 0:\n            for target_param, q_param in zip(self.target_network.parameters(), self.q_network.parameters()):\n                target_param.data.copy_(\n                    self.tau * q_param.data + (1.0 - self.tau) * target_param.data\n                )\n        \n        return {\n            \"td_loss\": quantile_loss.item(),\n            \"latent_loss\": latent_loss.item(),\n            \"total_loss\": total_loss.item(),\n            \"q_values\": current_quantiles.mean().item()\n        }"
7}

1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_value_atari.py",
4  "start_line": 289,
5  "end_line": 294,
6  "content": "    def select_action(self, obs, epsilon):\n        \"\"\"Epsilon-greedy action selection using Q-values.\"\"\"\n        if random.random() < epsilon:\n            # Return random action (single environment)\n            return np.array([random.randrange(self.q_network.n_actions)])\n        q_values = self.q_network(torch.Tensor(obs).to(self.device))\n        return torch.argmax(q_values, dim=1).cpu().numpy()"
7}

1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_value_atari.py",
4  "start_line": 263,
5  "end_line": 359,
6  "content": "class ValueAlgorithm:\n    \"\"\"Implicit Quantile Networks (IQN) with temporal feature smoothing regularization.\"\"\"\n    \n    def __init__(self, envs, device, args):\n        self.device = device\n        self.gamma = args.gamma\n        self.tau = args.tau\n        self.target_network_frequency = args.target_network_frequency\n        self.n_quantile_samples = 8  # number of tau samples for loss\n        self.n_target_samples = 8\n        self.kappa = 1.0  # Huber loss threshold\n        self.lambda_latent = 0.01  # weight for temporal feature smoothing\n        \n        self.q_network = QNetwork(envs).to(device)\n        self.target_network = QNetwork(envs).to(device)\n        self.target_network.load_state_dict(self.q_network.state_dict())\n        self.optimizer = optim.Adam(self.q_network.parameters(), lr=args.learning_rate)\n        \n    def select_action(self, obs, epsilon):\n        \"\"\"Epsilon-greedy action selection using Q-values.\"\"\"\n        if random.random() < epsilon:\n            # Return random action (single environment)\n            return np.array([random.randrange(self.q_network.n_actions)])\n        q_values = self.q_network(torch.Tensor(obs).to(self.device))\n        return torch.argmax(q_values, dim=1).cpu().numpy()\n    \n    def update(self, batch, global_step):\n        \"\"\"IQN update with temporal feature smoothing regularization.\"\"\"\n        batch_size = batch.observations.shape[0]\n        \n        # Sample quantiles for current and target networks\n        tau = torch.rand(batch_size, self.n_quantile_samples, device=self.device)  # [batch, N]\n        tau_prime = torch.rand(batch_size, self.n_target_samples, device=self.device)  # [batch, N']\n        \n        with torch.no_grad():\n            # Compute target quantile values\n            next_quantiles = self.target_network._quantile_values(\n                self.target_network.encoder(batch.next_observations),\n                tau_prime\n            )  # [batch, N', n_actions]\n            next_q = next_quantiles.mean(dim=1)  # [batch, n_actions]\n            next_actions = next_q.argmax(dim=1)  # [batch]\n            # Select quantiles for best actions\n            next_quantiles_best = next_quantiles[torch.arange(batch_size), :, next_actions]  # [batch, N']\n            # Compute target quantile values\n            target_quantiles = batch.rewards + self.gamma * next_quantiles_best * (1 - batch.dones)\n            target_quantiles = target_quantiles.unsqueeze(1)  # [batch, 1, N'] for broadcast\n        \n        # Current quantile values for taken actions\n        current_quantiles_all = self.q_network._quantile_values(\n            self.q_network.encoder(batch.observations),\n            tau\n        )  # [batch, N, n_actions]\n        current_quantiles = current_quantiles_all[torch.arange(batch_size), :, batch.actions.flatten()]  # [batch, N]\n        \n        # Quantile Huber loss\n        # Pairwise TD errors: [batch, N, N']\n        td_errors = target_quantiles - current_quantiles.unsqueeze(2)\n        huber_loss = F.huber_loss(td_errors, torch.zeros_like(td_errors), reduction='none', delta=self.kappa)\n        quantile_loss = (huber_loss * (tau.unsqueeze(2) - (td_errors < 0).float()).abs()).mean()\n        \n        # Temporal feature smoothing: encourage consecutive features to be similar\n        features_current = self.q_network.encoder(batch.observations)\n        features_next = self.q_network.encoder(batch.next_observations)\n        latent_loss = F.mse_loss(features_current, features_next.detach())\n        \n        # Total loss\n        total_loss = quantile_loss + self.lambda_latent * latent_loss\n        \n        self.optimizer.zero_grad()\n        total_loss.backward()\n        self.optimizer.step()\n        \n        # Soft target network update\n        if global_step % self.target_network_frequency == 0:\n            for target_param, q_param in zip(self.target_network.parameters(), self.q_network.parameters()):\n                target_param.data.copy_(\n                    self.tau * q_param.data + (1.0 - self.tau) * target_param.data\n                )\n        \n        return {\n            \"td_loss\": quantile_loss.item(),\n            \"latent_loss\": latent_loss.item(),\n            \"total_loss\": total_loss.item(),\n            \"q_values\": current_quantiles.mean().item()\n        }"
7}

1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_value_atari.py",
4  "start_line": 186,
5  "end_line": 260,
6  "content": "class QNetwork(nn.Module):\n    \"\"\"Implicit Quantile Network (IQN) for distributional RL.\n    \n    Encodes quantile samples via cosine embedding and combines with state features.\n    \"\"\"\n    def __init__(self, envs, quantile_embed_dim=32):\n        super().__init__()\n        n_actions = envs.single_action_space.n\n        self.n_actions = n_actions\n        self.quantile_embed_dim = quantile_embed_dim\n        self.encoder = NatureDQNEncoder()\n        \n        # Cosine embedding for quantile samples\n        self.cos_embed = nn.Linear(quantile_embed_dim, 512)\n        # Main head: after elementwise product of features and quantile embedding\n        self.head = nn.Linear(512, n_actions)\n        \n    def forward(self, x, tau_samples=None):\n        \"\"\"Return Q-values (mean over quantiles) or quantile values if tau_samples provided.\n        \n        Args:\n            x: state tensor [batch, 4, 84, 84]\n            tau_samples: optional [batch, n_samples] quantile samples in (0,1)\n        \n        Returns:\n            If tau_samples is None: Q-values [batch, n_actions]\n            Else: quantile values [batch, n_samples, n_actions]\n        \"\"\"\n        features = self.encoder(x)  # [batch, 512]\n        \n        if tau_samples is None:\n            # For action selection: sample 8 tau per state and average.\n            batch_size = x.shape[0]\n            tau = torch.rand(batch_size, 8, device=x.device)  # [batch, 8]\n            quantile_vals = self._quantile_values(features, tau)  # [batch, 8, n_actions]\n            q_values = quantile_vals.mean(dim=1)  # [batch, n_actions]\n            return q_values\n        else:\n            # tau_samples: [batch, n_samples]\n            quantile_vals = self._quantile_values(features, tau_samples)  # [batch, n_samples, n_actions]\n            return quantile_vals\n    \n    def _quantile_values(self, features, tau):\n        \"\"\"Compute quantile values given features and tau samples.\n        \n        Args:\n            features: [batch, 512]\n            tau: [batch, n_samples]\n        \n        Returns:\n            quantile values: [batch, n_samples, n_actions]\n        \"\"\"\n        batch_size, n_samples = tau.shape\n        # Compute cosine embedding of tau\n        # tau in (0,1), expand to [batch, n_samples, quantile_embed_dim]\n        pi = torch.tensor(torch.pi, device=tau.device)\n        # Create cosine basis: cos(pi * i * tau) for i=1..quantile_embed_dim\n        i_vals = torch.arange(1, self.quantile_embed_dim + 1, device=tau.device).float()\n        # i_vals: [quantile_embed_dim]\n        # tau: [batch, n_samples, 1]\n        tau = tau.unsqueeze(-1)  # [batch, n_samples, 1]\n        # cos(pi * i * tau): [batch, n_samples, quantile_embed_dim]\n        cos_vals = torch.cos(pi * i_vals * tau)\n        # Linear transform\n        cos_embed = self.cos_embed(cos_vals)  # [batch, n_samples, 512]\n        # Expand features to match\n        features_expanded = features.unsqueeze(1).expand(-1, n_samples, -1)  # [batch, n_samples, 512]\n        # Elementwise product\n        combined = features_expanded * cos_embed  # [batch, n_samples, 512]\n        # ReLU activation (as in IQN paper)\n        combined = F.relu(combined)\n        # Linear head\n        quantile_vals = self.head(combined)  # [batch, n_samples, n_actions]\n        return quantile_vals"
7}

1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_value_atari.py",
4  "start_line": 262,
5  "end_line": 347,
6  "content": "class ValueAlgorithm:\n    \"\"\"Implicit Quantile Networks (IQN) with TD-weighted temporal feature smoothing.\"\"\"\n    \n    def __init__(self, envs, device, args):\n        self.device = device\n        self.gamma = args.gamma\n        self.tau = args.tau\n        self.target_network_frequency = args.target_network_frequency\n        self.n_quantile_samples = 16  # number of tau samples for loss\n        self.n_target_samples = 16\n        self.kappa = 1.0  # Huber loss threshold\n        self.lambda_latent = 0.001  # weight for TD-weighted feature smoothing\n        \n        self.q_network = QNetwork(envs).to(device)\n        self.target_network = QNetwork(envs).to(device)\n        self.target_network.load_state_dict(self.q_network.state_dict())\n        self.optimizer = optim.Adam(self.q_network.parameters(), lr=args.learning_rate)\n        \n    def select_action(self, obs, epsilon):\n        \"\"\"Epsilon-greedy action selection using Q-values.\"\"\"\n        if random.random() < epsilon:\n            # Return random action (single environment)\n            return np.array([random.randrange(self.q_network.n_actions)])\n        q_values = self.q_network(torch.Tensor(obs).to(self.device))\n        return torch.argmax(q_values, dim=1).cpu().numpy()\n    \n    def update(self, batch, global_step):\n        \"\"\"IQN update with TD-weighted temporal feature smoothing.\"\"\"\n        batch_size = batch.observations.shape[0]\n        \n        # Sample quantiles for current and target networks\n        tau = torch.rand(batch_size, self.n_quantile_samples, device=self.device)  # [batch, N]\n        tau_prime = torch.rand(batch_size, self.n_target_samples, device=self.device)  # [batch, N']\n        \n        with torch.no_grad():\n            # Compute target quantile values\n            next_quantiles = self.target_network._quantile_values(\n                self.target_network.encoder(batch.next_observations),\n                tau_prime\n            )  # [batch, N', n_actions]\n            next_q = next_quantiles.mean(dim=1)  # [batch, n_actions]\n            next_actions = next_q.argmax(dim=1)  # [batch]\n            # Select quantiles for best actions\n            next_quantiles_best = next_quantiles[torch.arange(batch_size), :, next_actions]  # [batch, N']\n            # Compute target quantile values\n            target_quantiles = batch.rewards + self.gamma * next_quantiles_best * (1 - batch.dones)\n            target_quantiles = target_quantiles.unsqueeze(1)  # [batch, 1, N'] for broadcast\n        \n        # Current quantile values for taken actions\n        current_quantiles_all = self.q_network._quantile_values(\n            self.q_network.encoder(batch.observations),\n            tau\n        )  # [batch, N, n_actions]\n        current_quantiles = current_quantiles_all[torch.arange(batch_size), :, batch.actions.flatten()]  # [batch, N]\n        \n        # Quantile Huber loss\n        # Pairwise TD errors: [batch, N, N']\n        td_errors = target_quantiles - current_quantiles.unsqueeze(2)\n        huber_loss = F.huber_loss(td_errors, torch.zeros_like(td_errors), reduction='none', delta=self.kappa)\n        quantile_loss = (huber_loss * (tau.unsqueeze(2) - (td_errors < 0).float()).abs()).mean()\n        \n        # TD-weighted temporal feature smoothing\n        features_current = self.q_network.encoder(batch.observations)\n        features_next = self.q_network.encoder(batch.next_observations)\n        # Compute MSE per sample\n        mse_per_sample = F.mse_loss(features_current, features_next.detach(), reduction='none').mean(dim=1)  # [batch]\n        # Compute TD error magnitude per sample\n        td_mag = td_errors.abs().mean(dim=(1,2))  # [batch]\n        # Weight: higher weight where TD error is small (confidence)\n        weight = torch.exp(-td_mag.detach())  # [batch]\n        latent_loss = (weight * mse_per_sample).mean()\n        \n        # Total loss\n        total_loss = quantile_loss + self.lambda_latent * latent_loss\n        \n        self.optimizer.zero_grad()\n        total_loss.backward()\n        self.optimizer.step()\n        \n        # Soft target network update\n        if global_step % self.target_network_frequency == 0:\n            for target_param, q_param in zip(self.target_network.parameters(), self.q_network.parameters()):\n                target_param.data.copy_(\n                    self.tau * q_param.data + (1.0 - self.tau) * target_param.data\n                )\n        \n        return {\n            \"td_loss\": quantile_loss.item(),\n            \"latent_loss\": latent_loss.item(),\n            \"total_loss\": total_loss.item(),\n            \"q_values\": current_quantiles.mean().item()\n        }"
7}

1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_value_atari.py",
4  "start_line": 216,
5  "end_line": 222,
6  "content": "        if tau_samples is None:\n            # For action selection: sample 16 tau per state and average.\n            batch_size = x.shape[0]\n            tau = torch.rand(batch_size, 16, device=x.device)  # [batch, 16]\n            quantile_vals = self._quantile_values(features, tau)  # [batch, 16, n_actions]\n            q_values = quantile_vals.mean(dim=1)  # [batch, n_actions]\n            return q_values"
7}