Upload 10 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+figures/performance.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,176 @@
+
+<div align="center">
+  <h1>UD7</h1>
+  <h3>Provable Generalization of Clipped Double Q-Learning for Variance Reduction and Sample Efficiency</h3>
+  
+  <a href="https://www.python.org/">
+    <img src="https://img.shields.io/badge/Python-3.7+-blue?logo=python&style=flat-square" alt="Python Badge"/>
+  </a>
+  &nbsp;&nbsp;
+  <a href="https://pytorch.org/">
+    <img src="https://img.shields.io/badge/PyTorch-1.8+-EE4C2C?logo=pytorch&style=flat-square" alt="PyTorch Badge"/>
+  </a>
+  &nbsp;&nbsp;
+  <a href="https://www.sciencedirect.com/journal/neurocomputing">
+    <img src="https://img.shields.io/badge/Neurocomputing-Published-success?style=flat-square" alt="Neurocomputing Badge"/>
+  </a>
+  &nbsp;&nbsp;
+  <a href="https://www.elsevier.com/">
+    <img src="https://img.shields.io/badge/Elsevier-Journal-orange?style=flat-square" alt="Elsevier Badge"/>
+  </a>
+</div>
+
+---
+
+## Neurocomputing — PyTorch Implementation
+
+This repository contains a PyTorch implementation of **UD7** of the paper:
+
+> **Provable Generalization of Clipped Double Q-Learning for Variance Reduction and Sample Efficiency**  
+> Jangwon Kim, Jiseok Jeong, Soohee Han  
+> *Neurocomputing*, Volume 673, 7 April 2026, 132772
+
+### Paper Link
+https://www.sciencedirect.com/science/article/abs/pii/S0925231226001694
+
+---
+
+**UD7** is an off-policy actor–critic algorithm that builds on a TD7-style training pipeline, while replacing the critic target **formulation** with **UBOC**.
+
+---
+
+## 1) Background: Clipped Double Q-Learning (CDQ)
+
+Clipped double Q-learning is a widely-used bias correction in actor-critic methods (e.g., TD3). It maintains **two critics** and uses the **minimum** of the two as the TD target:
+
+$$
+y_{\text{CDQ}}(s_t,a_t)=r_t+\gamma \min_{i\in\{1,2\}} \bar Q_i(s_{t+1}, a_{t+1})
+$$
+
+### Strengths (why CDQ is popular)
+- **Effective overestimation control:** taking a minimum is conservative, often preventing exploding Q-values.
+- **Robust baseline behavior:** works well across many continuous-control tasks.
+
+### Limitations (what the paper highlights)
+- **High variance:** when critics are poorly learned early on, the min operator can yield high-variance TD targets, destabilizing TD learning and reducing sample efficiency.
+
+**UBOC is motivated by a concrete question:**
+> Can we obtain **the same expected target value as CDQ**, but with **smaller variance**?
+
+---
+
+## 2) UBOC: Uncertainty-Based Overestimation Correction (Detailed)
+
+UBOC views the critic outputs as a **distribution of Q estimates** (because function approximation is noisy).  
+Instead of using `min(Q1, Q2)`, UBOC uses **N critics** to estimate:
+- a **mean** \(m\),
+- an **(unbiased) standard deviation** \(\hat s\),
+and then forms a corrected value:
+
+$$
+Q_{\text{corrected}} = m - x\cdot \hat s
+$$
+
+where \(x>0\) controls conservativeness.
+
+### 2.1 Expectation equivalence to clipped double-Q
+
+Under the assumption that critic estimates behave like i.i.d. samples from a normal distribution, we can derive:
+
+$$
+\mathbb{E}\left[\min(Q_A, Q_B)\right]=\mathbb{E}\left[m - \frac{\hat s}{\sqrt{\pi}}\right]
+$$
+
+This is the key insight:
+- choosing $$x=1/\sqrt{\pi}$$ makes the corrected estimate **match CDQ in expectation**.
+
+### 2.2 Variance reduction (provable)
+
+We can further prove that the estimator
+
+$$
+m - \frac{\hat s}{\sqrt{\pi}}
+$$
+
+has **strictly smaller variance** than the CDQ minimum-based target, and the **variance gap is strictly positive for all $$N\ge 2\$$**.
+
+As $$N\to\infty$$, the maximum achievable variance reduction is upper-bounded by:
+
+$$
+\sigma^2\left(1-\frac{1}{\pi}\right)
+$$
+
+**It means that**
+- UBOC does not only “bias-correct”; it **reduces noise** in TD targets.
+- This is especially important early in training, where noisy targets can derail learning.
+
+### 2.3 UBOC TD target (what you implement)
+
+Using N target critics $$Q_1,\dots, Q_N$$, compute:
+
+
+**Mean**
+
+$$
+m(s,a) = \frac{1}{N}\sum_{i=1}^N  Q_i(s,a)
+$$
+
+**Unbiased variance (Approximation)**
+
+$$
+\hat v(s,a)=\frac{1}{N-1}\sum_{i=1}^N \left( Q_i(s,a)-m(s,a)\right)^2
+$$
+
+Then the **UBOC target** is:
+
+$$
+y_{\text{UBOC}}(s_t,a_t)=r_t + \gamma\left(m(s_{t+1},a_{t+1}) - \sqrt{\frac{\hat v(s_{t+1},a_{t+1})}{\pi}}\right)
+$$
+
+where $$a_{t+1}$$ can be computed with target policy smoothing.
+
+This gives a *dynamic* bias correction driven by critic uncertainty.
+
+---
+
+## 3) UD7: TD7 + UBOC Targets
+
+**UD7** integrates UBOC into a TD7-style pipeline and emphasizes strong sample efficiency.
+
+- UD7 uses the TD7 background for practical stability/efficiency.
+- **The main difference from TD7 is the critic training target:** UD7 uses **UBOC targets** and a multi-critic ensemble (commonly **N=5**).
+
+> If you already have a TD7 baseline, UD7 is best viewed as:  
+> **“swap the target rule + use N critics, then keep the rest of the training recipe.”**
+
+---
+
+## 4) Performance
+
+<div align="center">
+  <img src="figures/performance.png" alt="Fig. 1 — Performance comparison on MuJoCo benchmarks" width="800"/>
+</div>
+
+---
+
+## 5) Computational Overhead 
+
+Runtime figure (tested on RTX 3090 Ti + Intel i7-12700):
+
+<div align="center">
+  <img src="figures/runtime.png" alt="Fig. 2 — Runtime comparison" width="300"/>
+</div>
+
+---
+
+## Citation
+```
+@article{kim2026provable,
+  title={Provable generalization of clipped double Q-learning for variance reduction and sample efficiency},
+  author={Kim, Jangwon and Jeong, Jiseok and Han, Soohee},
+  journal={Neurocomputing},
+  pages={132772},
+  year={2026},
+  publisher={Elsevier}
+}
+```

figures/tmp.md

@@ -0,0 +1 @@
+

main.py

@@ -0,0 +1,78 @@
+import argparse
+import torch
+from ud7 import UD7
+from trainer import Trainer
+from utils import set_seed, make_env
+
+
+def get_parameters():
+    parser = argparse.ArgumentParser()
+
+    # Environment Setting
+    parser.add_argument('--env-name', default='Humanoid-v4')
+    parser.add_argument('--random-seed', default=-1, type=int)
+
+    # UBOC
+    parser.add_argument('--num_critics', default=5, type=int)
+
+    # Checkpointing
+    parser.add_argument('--use_checkpoints', default=True, type=bool)
+    parser.add_argument('--max-eps-when-checkpointing', default=20, type=int)
+    parser.add_argument('--steps-before-checkpointing', default=75e4, type=int)
+    parser.add_argument('--reset-weight', default=0.9, type=float)
+
+    # LAP
+    parser.add_argument('--alpha', default=0.4, type=float)
+    parser.add_argument('--min_priority', default=1, type=float)
+
+    # Generic
+    parser.add_argument('--target-update-rate', default=250, type=int)
+    parser.add_argument('--start-steps', default=25e3, type=int)
+    parser.add_argument('--max-steps', default=5000000, type=int)
+    parser.add_argument('--zs-dim', default=256, type=int)
+    parser.add_argument('--critic-hidden-dims', default=(256, 256))
+    parser.add_argument('--policy-hidden-dims', default=(256, 256))
+    parser.add_argument('--encoder-hidden-dims', default=(256, 256))
+    parser.add_argument('--hidden-dims', default=(256, 256))
+    parser.add_argument('--batch-size', default=256, type=int)
+    parser.add_argument('--buffer-size', default=1000000, type=int)
+    parser.add_argument('--policy-update-delay', default=2)
+    parser.add_argument('--gamma', default=0.99, type=float)
+    parser.add_argument('--actor-lr', default=0.0003, type=float)
+    parser.add_argument('--critic-lr', default=0.0003, type=float)
+    parser.add_argument('--encoder-lr', default=0.0003, type=float)
+
+    # TD3
+    parser.add_argument('--act-noise-scale', default=0.1, type=float)
+    parser.add_argument('--target-noise-scale', default=0.2, type=float)
+    parser.add_argument('--target-noise-clip', default=0.5, type=float)
+
+    # Log & Evaluation
+    parser.add_argument('--show-loss', default=False, type=bool)
+    parser.add_argument('--eval_flag', default=True, type=bool)
+    parser.add_argument('--eval-freq', default=5000, type=int)
+    parser.add_argument('--eval-episode', default=10, type=int)
+
+    param = parser.parse_args()
+
+    return param
+
+
+def main(args):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    random_seed = set_seed(args.random_seed)
+    env, eval_env = make_env(args.env_name, random_seed)
+
+    state_dim = env.observation_space.shape[0]
+    action_dim = env.action_space.shape[0]
+    action_bound = [env.action_space.low[0], env.action_space.high[0]]
+
+    agent = UD7(state_dim, action_dim, action_bound, device, args)
+
+    trainer = Trainer(env, eval_env, agent, args)
+    trainer.run()
+
+
+if __name__ == '__main__':
+    args = get_parameters()
+    main(args)

network.py

@@ -0,0 +1,145 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from utils import weight_init, AvgL1Norm
+
+
+class EnsembleQNet(nn.Module):
+    def __init__(self, num_critics, state_dim, action_dim, device, zs_dim=256, hidden_dims=(256, 256), activation_fc=F.elu):
+        super(EnsembleQNet, self).__init__()
+        self.device = device
+        self.activation_fc = activation_fc
+
+        self.num_critics = num_critics
+
+        self.q_nets = nn.ModuleList()
+        for _ in range(self.num_critics):
+            q_net = self._build_q_net(state_dim, action_dim, zs_dim, hidden_dims)
+            self.q_nets.append(q_net)
+
+        self.apply(weight_init)
+
+    def _build_q_net(self, state_dim, action_dim, zs_dim, hidden_dims):
+        q_net = nn.ModuleDict({
+            's_input_layer': nn.Linear(state_dim + action_dim, hidden_dims[0]),
+            'emb_input_layer': nn.Linear(2 * zs_dim + hidden_dims[0], hidden_dims[0]),
+            'emb_hidden_layers': nn.ModuleList([
+                nn.Linear(hidden_dims[i], hidden_dims[i + 1]) for i in range(len(hidden_dims) - 1)
+            ]),
+            'output_layer': nn.Linear(hidden_dims[-1], 1)
+        })
+        return q_net
+
+    def _format(self, state, action):
+        x, u = state, action
+        if not isinstance(x, torch.Tensor):
+            x = torch.tensor(x, device=self.device, dtype=torch.float32)
+            x = x.unsqueeze(0)
+
+        if not isinstance(u, torch.Tensor):
+            u = torch.tensor(u, device=self.device, dtype=torch.float32)
+            u = u.unsqueeze(0)
+
+        return x, u
+
+    def forward(self, state, action, zsa, zs):
+        s, a = self._format(state, action)
+        sa = torch.cat([s, a], dim=1)
+        embeddings = torch.cat([zsa, zs], dim=1)
+
+        q_values = []
+        for q_net in self.q_nets:
+            q = AvgL1Norm(q_net['s_input_layer'](sa))
+            q = torch.cat([q, embeddings], dim=1)
+            q = self.activation_fc(q_net['emb_input_layer'](q))
+            for hidden_layer in q_net['emb_hidden_layers']:
+                q = self.activation_fc(hidden_layer(q))
+            q = q_net['output_layer'](q)
+            q_values.append(q)
+
+        return torch.cat(q_values, dim=1)
+
+
+class Policy(nn.Module):
+    def __init__(self, state_dim, action_dim, device, zs_dim=256, hidden_dims=(256, 256), activation_fc=F.relu):
+        super(Policy, self).__init__()
+        self.device = device
+        self.apply(weight_init)
+
+        self.activation_fc = activation_fc
+
+        self.s_input_layer = nn.Linear(state_dim, hidden_dims[0])
+        self.zss_input_layer = nn.Linear(zs_dim + hidden_dims[0], hidden_dims[0])
+        self.zss_hidden_layers = nn.ModuleList()
+        for i in range(len(hidden_dims)-1):
+            hidden_layer = nn.Linear(hidden_dims[i], hidden_dims[i+1])
+            self.zss_hidden_layers.append(hidden_layer)
+        self.zss_output_layer = nn.Linear(hidden_dims[-1], action_dim)
+
+    def _format(self, state):
+        x = state
+        if not isinstance(x, torch.Tensor):
+            x = torch.tensor(x, device=self.device, dtype=torch.float32)
+            x = x.unsqueeze(0)
+
+        return x
+
+    def forward(self, state, zs):
+        state = self._format(state)
+
+        state = AvgL1Norm(self.s_input_layer(state))
+        zss = torch.cat([state, zs], 1)
+
+        zss = self.activation_fc(self.zss_input_layer(zss))
+        for i, hidden_layer in enumerate(self.zss_hidden_layers):
+            zss = self.activation_fc(hidden_layer(zss))
+        zss = self.zss_output_layer(zss)
+        action = torch.tanh(zss)
+        return action
+
+
+class Encoder(nn.Module):
+    def __init__(self, state_dim, action_dim, device, zs_dim=256, hidden_dims=(256, 256), activation_fc=F.elu):
+        super(Encoder, self).__init__()
+        self.device = device
+        self.activation_fc = activation_fc
+
+        self.s_encoder_input_layer = nn.Linear(state_dim, hidden_dims[0])
+        self.s_encoder_hidden_layers = nn.ModuleList()
+        for i in range(len(hidden_dims)-1):
+            hidden_layer = nn.Linear(hidden_dims[i], hidden_dims[i+1])
+            self.s_encoder_hidden_layers.append(hidden_layer)
+        self.s_encoder_output_layer = nn.Linear(hidden_dims[-1], zs_dim)
+
+        self.zsa_encoder_input_layer = nn.Linear(zs_dim + action_dim, hidden_dims[0])
+        self.zsa_encoder_hidden_layers = nn.ModuleList()
+        for i in range(len(hidden_dims)-1):
+            hidden_layer = nn.Linear(hidden_dims[i], hidden_dims[i+1])
+            self.zsa_encoder_hidden_layers.append(hidden_layer)
+        self.zsa_encoder_output_layer = nn.Linear(hidden_dims[-1], zs_dim)
+
+    def _format(self, state):
+        x = state
+        if not isinstance(x, torch.Tensor):
+            x = torch.tensor(x, device=self.device, dtype=torch.float32)
+            x = x.unsqueeze(0)
+        return x
+
+    def zs(self, state):
+        state = self._format(state)
+
+        zs = self.activation_fc(self.s_encoder_input_layer(state))
+        for i, hidden_layer in enumerate(self.s_encoder_hidden_layers):
+            zs = self.activation_fc(hidden_layer(zs))
+        zs = AvgL1Norm(self.s_encoder_output_layer(zs))
+        return zs
+
+    def zsa(self, zs, action):
+        action = self._format(action)
+        zsa = torch.cat([zs, action], 1)
+
+        zsa = self.activation_fc(self.zsa_encoder_input_layer(zsa))
+        for i, hidden_layer in enumerate(self.zsa_encoder_hidden_layers):
+            zsa = self.activation_fc(hidden_layer(zsa))
+        zsa = self.zsa_encoder_output_layer(zsa)
+        return zsa

replay_memory.py

@@ -0,0 +1,67 @@
+import torch
+import numpy as np
+
+
+class LAP:
+    def __init__(self, state_dim, action_dim, device, capacity=1e6, normalize_action=True, max_action=1, prioritized=True):
+        # Set the device
+        self.device = device
+
+        # Set the replay buffer capacity
+        self.capacity = int(capacity)
+        self.size = 0
+        self.position = 0
+
+        # Set the action normalization factor
+        self.do_normalize_action = normalize_action
+        self.normalize_action = max_action if normalize_action else 1
+        self.max_action = max_action
+
+        # Set the prioritized flag
+        self.prioritized = prioritized
+        if prioritized:
+            self.priority = torch.zeros(self.capacity, device=device)
+            self.max_priority = 1
+
+        # Initialize the replay buffer
+        self.state_buffer = np.empty(shape=(self.capacity, state_dim), dtype=np.float32)
+        self.action_buffer = np.empty(shape=(self.capacity, action_dim), dtype=np.float32)
+        self.reward_buffer = np.empty(shape=(self.capacity, 1), dtype=np.float32)
+        self.next_state_buffer = np.empty(shape=(self.capacity, state_dim), dtype=np.float32)
+        self.done_buffer = np.empty(shape=(self.capacity, 1), dtype=np.float32)
+
+    def push(self, state, action, reward, next_state, done):
+        self.state_buffer[self.position] = state
+        self.action_buffer[self.position] = action / self.normalize_action
+        self.reward_buffer[self.position] = reward
+        self.next_state_buffer[self.position] = next_state
+        self.done_buffer[self.position] = done
+
+        if self.prioritized:
+            self.priority[self.position] = self.max_priority
+
+        self.position = (self.position + 1) % self.capacity
+        self.size = min(self.size + 1, self.capacity)
+
+    def sample(self, batch_size):
+        if self.prioritized:
+            csum = torch.cumsum(self.priority[:self.size], 0)
+            val = torch.rand(size=(batch_size,), device=self.device) * csum[-1]
+            self.ind = torch.searchsorted(csum, val).cpu().data.numpy()
+        else:
+            self.ind = np.random.randint(0, self.size, size=batch_size)
+
+        states = torch.FloatTensor(self.state_buffer[self.ind]).to(self.device)
+        actions = torch.FloatTensor(self.action_buffer[self.ind]).to(self.device)
+        rewards = torch.FloatTensor(self.reward_buffer[self.ind]).to(self.device)
+        next_states = torch.FloatTensor(self.next_state_buffer[self.ind]).to(self.device)
+        dones = torch.FloatTensor(self.done_buffer[self.ind]).to(self.device)
+
+        return states, actions, rewards, next_states, dones
+
+    def update_priority(self, priority):
+        self.priority[self.ind] = priority.reshape(-1).detach()
+        self.max_priority = max(float(priority.max()), self.max_priority)
+
+    def reset_max_priority(self):
+        self.max_priority = float(self.priority[:self.size].max())

trainer.py

@@ -0,0 +1,112 @@
+import numpy as np
+
+class Trainer:
+    def __init__(self, env, eval_env, agent, args):
+        self.args = args
+
+        self.agent = agent
+        self.env_name = args.env_name
+        self.env = env
+        self.eval_env = eval_env
+
+        self.start_steps = args.start_steps
+        self.max_steps = args.max_steps
+        self.batch_size = args.batch_size
+        self.target_noise_scale = args.target_noise_scale
+
+        self.eval_flag = args.eval_flag
+        self.eval_episode = args.eval_episode
+        self.eval_freq = args.eval_freq
+
+        self.episode = 0
+        self.episode_reward = 0
+        self.total_steps = 0
+        self.eval_num = 0
+        self.finish_flag = False
+
+        self.target_noise_scale = args.target_noise_scale
+        self.policy_update_delay = args.policy_update_delay
+
+    def evaluate(self):
+        # Evaluate process
+        self.eval_num += 1
+        reward_list = []
+
+        for epi in range(self.eval_episode):
+            epi_reward = 0
+            state, _ = self.eval_env.reset()
+
+            done = False
+
+            while not done:
+                action = self.agent.get_action(state, use_checkpoint=self.args.use_checkpoints, add_noise=False)
+                next_state, reward, terminated, truncated, _ = self.eval_env.step(action)
+                done = terminated or truncated
+                epi_reward += reward
+                state = next_state
+            reward_list.append(epi_reward)
+
+        print("Eval  |  total_step {}  |  episode {}  |  Average Reward {:.2f}  |  Max reward: {:.2f}  |  "
+              "Min reward: {:.2f}".format(self.total_steps, self.episode, sum(reward_list)/len(reward_list),
+                                               max(reward_list), min(reward_list), np.std(reward_list)))
+
+    def run(self):
+        # Train-process start.
+        allow_train = False
+
+        while not self.finish_flag:
+            self.episode += 1
+            done = False
+            ep_total_reward, ep_timesteps = 0, 0
+
+            state, _ = self.env.reset()
+            # Episode start.
+            while not done:
+                self.total_steps += 1
+                ep_timesteps += 1
+
+                if allow_train:
+                    action = self.agent.get_action(state, use_checkpoint=False, add_noise=True)
+                else:
+                    action = self.env.action_space.sample()
+                next_state, reward, terminated, truncated, _ = self.env.step(action)
+                done = terminated or truncated
+
+                ep_total_reward += reward
+
+                done_mask = 0.0 if ep_timesteps == self.env._max_episode_steps else float(done)
+                self.agent.buffer.push(state, action, reward, next_state, done_mask)
+
+                state = next_state
+
+                if allow_train and not self.args.use_checkpoints:
+                    actor_loss, critic_loss, encoder_loss = self.agent.train()
+                    # Print loss.
+                    if self.args.show_loss:
+                        print("Loss  |  Actor loss {:.3f}  |  Critic loss {:.3f}  |  Encoder loss {:.3f}"
+                              .format(actor_loss, critic_loss, encoder_loss))
+
+                if done:
+                    if allow_train and self.args.use_checkpoints:
+                        self.agent.maybe_train_and_checkpoint(ep_timesteps, ep_total_reward)
+
+                    if self.total_steps >= self.args.start_steps:
+                        allow_train = True
+
+                # Evaluation.
+                if self.eval_flag and self.total_steps % self.eval_freq == 0:
+                    self.evaluate()
+
+                # Raise finish_flag.
+                if self.total_steps == self.max_steps:
+                    self.finish_flag = True
+
+
+
+
+
+
+
+
+
+

ud7.py

@@ -0,0 +1,202 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import copy
+from replay_memory import LAP
+from network import Policy, Encoder, EnsembleQNet
+from utils import hard_update, LAP_huber
+
+
+class UD7:
+    def __init__(self, state_dim, action_dim, action_bound, device, args):
+        self.args = args
+
+        self.state_dim = state_dim
+        self.action_dim = action_dim
+
+        self.device = device
+        self.buffer = LAP(self.state_dim, self.action_dim, device, args.buffer_size, normalize_action=True,
+                          max_action=action_bound[1], prioritized=True)
+        self.batch_size = args.batch_size
+
+        self.gamma = args.gamma
+        self.act_noise_scale = args.act_noise_scale
+
+        self.num_critics = args.num_critics
+
+        self.actor = Policy(self.state_dim, self.action_dim, self.device, args.zs_dim, args.policy_hidden_dims).to(self.device)
+        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=args.actor_lr)
+        self.target_actor = Policy(self.state_dim, self.action_dim, self.device, args.zs_dim, args.policy_hidden_dims).to(self.device)
+
+        self.critic = EnsembleQNet(self.num_critics,  self.state_dim, self.action_dim,
+                                   self.device, args.zs_dim, args.critic_hidden_dims).to(self.device)
+        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=args.critic_lr)
+        self.target_critic = EnsembleQNet(self.num_critics, self.state_dim, self.action_dim,
+                                   self.device, args.zs_dim, args.critic_hidden_dims).to(self.device)
+
+        self.encoder = Encoder(state_dim, action_dim, self.device, args.zs_dim, args.encoder_hidden_dims).to(self.device)
+        self.encoder_optimizer = torch.optim.Adam(self.encoder.parameters(), lr=args.encoder_lr)
+        self.fixed_encoder = copy.deepcopy(self.encoder)
+        self.fixed_encoder_target = copy.deepcopy(self.encoder)
+
+        self.checkpoint_actor = copy.deepcopy(self.actor)
+        self.checkpoint_encoder = copy.deepcopy(self.encoder)
+
+        self.training_steps = 0
+
+        self.max_action = action_bound[1]
+
+        # Checkpointing tracked values
+        self.eps_since_update = 0
+        self.timesteps_since_update = 0
+        self.max_eps_before_update = 1
+        self.min_return = 1e8
+        self.best_min_return = -1e8
+
+        # Value clipping tracked values
+        self.max = -1e8
+        self.min = 1e8
+        self.max_target = 0
+        self.min_target = 0
+
+        hard_update(self.actor, self.target_actor)
+        hard_update(self.critic, self.target_critic)
+
+    def get_action(self, state, use_checkpoint=False, add_noise=True):
+        with torch.no_grad():
+            if add_noise:
+                if use_checkpoint:
+                    zs = self.checkpoint_encoder.zs(state)
+                    action = self.checkpoint_actor(state, zs)
+                    action = action + torch.randn_like(action) * self.act_noise_scale
+                    action = np.clip(action.cpu().numpy()[0], -1, 1)
+                else:
+                    zs = self.fixed_encoder.zs(state)
+                    action = self.actor(state, zs)
+                    action = action + torch.randn_like(action) * self.act_noise_scale
+                    action = np.clip(action.cpu().numpy()[0], -1, 1)
+            else:
+                if use_checkpoint:
+                    zs = self.checkpoint_encoder.zs(state)
+                    action = self.checkpoint_actor(state, zs).cpu().numpy()[0]
+                else:
+                    zs = self.fixed_encoder.zs(state)
+                    action = self.actor(state, zs).cpu().numpy()[0]
+
+            action = action * self.max_action
+        return action
+
+    def train(self):
+        self.training_steps += 1
+
+        # Sample from LAP
+        states, actions, rewards, next_states, dones = self.buffer.sample(self.batch_size)
+
+        # Update Encoder
+        with torch.no_grad():
+            next_zs = self.encoder.zs(next_states)
+        zs = self.encoder.zs(states)
+        pred_zs = self.encoder.zsa(zs, actions)
+        encoder_loss = F.mse_loss(pred_zs, next_zs)
+
+        self.encoder_optimizer.zero_grad()
+        encoder_loss.backward()
+        self.encoder_optimizer.step()
+
+        # Update Critic
+        with torch.no_grad():
+            fixed_target_zs = self.fixed_encoder_target.zs(next_states)
+
+            target_act_noise = (torch.randn_like(actions) * self.args.target_noise_scale).clamp(-self.args.target_noise_clip, self.args.target_noise_clip).to(self.device)
+
+            if self.buffer.do_normalize_action is True:
+                next_target_actions = (self.target_actor(next_states, fixed_target_zs) + target_act_noise).clamp(-1, 1)
+            else:
+                next_target_actions = (self.target_actor(next_states, fixed_target_zs) + target_act_noise).clamp(-self.max_action, self.max_action)
+
+            fixed_target_zsa = self.fixed_encoder_target.zsa(fixed_target_zs, next_target_actions)
+
+            Q_target = self.target_critic(next_states, next_target_actions, fixed_target_zsa, fixed_target_zs)
+            m = Q_target.mean(dim=1, keepdim=True)  # Sample mean
+            b = Q_target.var(dim=1, unbiased=True, keepdim=True) # Sample variance
+            Bias_Corrected_Q_target = m - 0.5641896 * torch.sqrt(b)  # bias-corrected target Q
+
+            Q_target = rewards + (1 - dones) * self.gamma * Bias_Corrected_Q_target.clamp(self.min_target, self.max_target)
+
+            self.max = max(self.max, float(Q_target.max()))
+            self.min = min(self.min, float(Q_target.min()))
+
+            fixed_zs = self.fixed_encoder.zs(states)
+            fixed_zsa = self.fixed_encoder.zsa(fixed_zs, actions)
+
+        Q = self.critic(states, actions, fixed_zsa, fixed_zs)
+
+        td_loss = (Q - Q_target).abs()
+        critic_loss = LAP_huber(td_loss)
+
+        self.critic_optimizer.zero_grad()
+        critic_loss.backward()
+        self.critic_optimizer.step()
+
+        # Update LAP
+        priority = td_loss.max(1)[0].clamp(min=self.args.min_priority).pow(self.args.alpha)
+        self.buffer.update_priority(priority)
+
+        # Update Actor
+        if self.training_steps % self.args.policy_update_delay == 0:
+            actor_actions = self.actor(states, fixed_zs)
+            fixed_zsa = self.fixed_encoder.zsa(fixed_zs, actor_actions)
+            Q = self.critic(states, actor_actions, fixed_zsa, fixed_zs)
+
+            actor_loss = -Q.mean(dim=1, keepdim=True).mean()
+
+            self.actor_optimizer.zero_grad()
+            actor_loss.backward()
+            self.actor_optimizer.step()
+        else:
+            actor_loss = torch.tensor(0.0)
+
+        # Update Iteration
+        if self.training_steps % self.args.target_update_rate == 0:
+            self.target_actor.load_state_dict(self.actor.state_dict())
+            self.target_critic.load_state_dict(self.critic.state_dict())
+            self.fixed_encoder_target.load_state_dict(self.fixed_encoder.state_dict())
+            self.fixed_encoder.load_state_dict(self.encoder.state_dict())
+
+            self.buffer.reset_max_priority()
+
+            self.max_target = self.max
+            self.min_target = self.min
+
+        return actor_loss.item(), critic_loss.item(), encoder_loss.item()
+
+    def maybe_train_and_checkpoint(self, ep_timesteps, ep_return):
+        self.eps_since_update += 1
+        self.timesteps_since_update += ep_timesteps
+
+        self.min_return = min(self.min_return, ep_return)
+
+        # End evaluation of current policy early
+        if self.min_return < self.best_min_return:
+            self.train_and_reset()
+
+        # Update checkpoint
+        elif self.eps_since_update == self.max_eps_before_update:
+            self.best_min_return = self.min_return
+            self.checkpoint_actor.load_state_dict(self.actor.state_dict())
+            self.checkpoint_encoder.load_state_dict(self.fixed_encoder.state_dict())
+
+            self.train_and_reset()
+
+    # Batch training
+    def train_and_reset(self):
+        for _ in range(self.timesteps_since_update):
+            if self.training_steps == self.args.steps_before_checkpointing:
+                self.best_min_return *= self.args.reset_weight
+                self.max_eps_before_update = self.args.max_eps_when_checkpointing
+
+            self.train()
+
+        self.eps_since_update = 0
+        self.timesteps_since_update = 0
+        self.min_return = 1e8

utils.py

@@ -0,0 +1,63 @@
+import numpy as np
+import random
+import torch
+import torch.nn as nn
+
+
+def AvgL1Norm(x, eps=1e-8):
+    return x / x.abs().mean(-1, keepdim=True).clamp(min=eps)
+
+
+def LAP_huber(x, min_priority=1):
+    return torch.where(x < min_priority, 0.5 * x.pow(2), min_priority * x).sum(1).mean()
+
+
+def weight_init(m):
+    """Custom weight init for Conv2D and Linear layers.
+        Reference: https://github.com/MishaLaskin/rad/blob/master/curl_sac.py"""
+
+    if isinstance(m, nn.Linear):
+        nn.init.orthogonal_(m.weight.data)
+        m.bias.data.fill_(0.0)
+    elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+        # delta-orthogonal init from https://arxiv.org/pdf/1806.05393.pdf
+        assert m.weight.size(2) == m.weight.size(3)
+        m.weight.data.fill_(0.0)
+        m.bias.data.fill_(0.0)
+        mid = m.weight.size(2) // 2
+        gain = nn.init.calculate_gain('relu')
+        nn.init.orthogonal_(m.weight.data[:, :, mid, mid], gain)
+
+
+def hard_update(network, target_network):
+    for param, target_param in zip(network.parameters(), target_network.parameters()):
+        target_param.data.copy_(param.data)
+
+
+def soft_update(network, target_network, tau):
+    for param, target_param in zip(network.parameters(), target_network.parameters()):
+        target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)
+
+
+def set_seed(random_seed):
+    if random_seed <= 0:
+        random_seed = np.random.randint(1, 9999)
+    else:
+        random_seed = random_seed
+
+    torch.manual_seed(random_seed)
+    np.random.seed(random_seed)
+    random.seed(random_seed)
+    return random_seed
+
+
+def make_env(env_name, random_seed):
+    import gymnasium as gym
+    # openai gym
+    env = gym.make(env_name)
+    env.action_space.seed(random_seed)
+
+    eval_env = gym.make(env_name)
+    eval_env.action_space.seed(random_seed + 100)
+
+    return env, eval_env