Repository: mmcaulif/Cardio (language: python)

Issue:
Catching dev up 
Adding the stuff I'd forgotten about to dev, from here on out, all work is done on dev


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index feae5c1..0aa842d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,7 @@
 *.pyc
-__pycache__
\ No newline at end of file
+__pycache__
+.vscode
+build
+*.egg-info
+run
+dist
\ No newline at end of file
diff --git a/LICENSE.txt b/LICENSE.txt
new file mode 100644
index 0000000..c74aced
--- /dev/null
+++ b/LICENSE.txt
@@ -0,0 +1,19 @@
+Copyright (c) 2016 The Python Packaging Authority (PyPA)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/README.md b/README.md
index 11c4f11..4519235 100644
--- a/README.md
+++ b/README.md
@@ -1,1 +1,55 @@
-# GymCardio
\ No newline at end of file
+# GymCardio
+Cardio offers to replace much of the boiler plate code in deep reinforcement learnign algorithms; namely the data gathering and environment interaction, but also providing pre-written policies and replay buffers, all while aiming to do so in a modular fashion and allow users to implement their own algorithms or improvements to existing algorithms.
+
+The purpose of this library was initially to speed up my own self-implemented versions of algorithms and reducing the code overhead while also reducing the actual number of lines written per script for each algorithm. You'll often find that the actual algorithm details often corresponds to very little code and handling of the environment and buffer causes even simple implementations to feel somewhat bloated. Thus cardio hopes to allow for simple, easy to read, one file implementations. In an effort to showcase the library I have included a number of 'stubs' that are just that.
+
+This is still heavily a work in progress and even many of the (poorly organised) stubs do not work with the current versions of the runner. Going forward I will be chipping away at streamlining, organising and documenting this repository when I get the chance!
+
+# Cardio basics
+This section will be overhauled at a later date, but now the gist of Cardio is the Runner class that gives you a simple interface that wraps the environment and then with one method, will step through the environment, collect transitions, and process them in your defined way. The runner supports n-step transitions, custom policies and custom processing. Currently the runner class is biased with off-policy algorithms in mind (altough fully supporting on-policy approaches too) but going forward it will be better balanced between the two.
+
+## Immediate to do list
+* [ ] Create package!
+  * basic implementation done, need to do PyPi and look into further improvements
+
+* [x] Implement replay buffer class and move IET work into Cardio ecosystem
+
+* [ ] Remove warmup method in gatherer and make it a special call of the rollout method
+
+## Vectorised env work
+* [ ] Align VectorCollector with Collector
+  * Revisit A2C implementation
+  * Make Logger compatible with VectorCollector
+  * Learning doesnt seem to line up with stable baselines3 need to debug all aspects (collector, value estimation and logger)
+
+* [ ] Makes sure vector collector work as intended for off-policy methods and n-step collector work as intended for on-policy methods etc.
+
+## Lower priority
+* [ ] Offline gatherer
+  * on pause until mujoco sorted
+
+* [ ] Minor refactor to gatherer and runner, add default arg values, careful consideration needed
+  * mostly done, just some final decisions to make
+  * change collector name to gatherer, idk why its different
+
+* [ ] Move agent stubs outside of src and using package instead of src
+
+* [ ] Sort policies better, i.e. discrete, continuous
+
+* [ ] Benchmark each implementation wrt. SB3 (change logging to timestep based first though)
+ 
+## Completed
+* [x] Implement multibatch sampling for off-policy runner
+
+* [x] Add episode length to logger and use the same names as SB3 for easy integration!
+
+* [x] Parallel gatherer
+
+* [x] Change logging from episodic to timestep based
+  * include window and log_interval arguments to gatherer
+
+* [x] Implement 'step_for' function in Collector!
+
+* [x] Create dummy env for testing logging, collection and sampling!
+
+* [x] Implement 'reduce' argument for n-step learning that returns unsqueezed vectors (for DRQN)
\ No newline at end of file
diff --git a/cardio_rl/__init__.py b/cardio_rl/__init__.py
new file mode 100644
index 0000000..ba74248
--- /dev/null
+++ b/cardio_rl/__init__.py
@@ -0,0 +1,2 @@
+from .gatherer import Collector, VectorCollector
+from .runner import Runner
\ No newline at end of file
diff --git a/__init__.py b/cardio_rl/buffers/__init__.py
similarity index 100%
rename from __init__.py
rename to cardio_rl/buffers/__init__.py
diff --git a/cardio_rl/buffers/er_buffer.py b/cardio_rl/buffers/er_buffer.py
new file mode 100644
index 0000000..47d85bb
--- /dev/null
+++ b/cardio_rl/buffers/er_buffer.py
@@ -0,0 +1,20 @@
+import random
+from cardio_rl.transitions import TorchTransition
+from collections import deque
+
+
+class ErTable:
+	def __init__(self, capacity=1_000_000, transition_func=TorchTransition):
+		self.buffer = deque(maxlen=capacity)
+		self.transition_func = transition_func
+	
+	def __len__(self):
+		return len(self.buffer)
+
+	def store(self, transition):
+		self.buffer.append(transition)
+
+	def sample(self, batch_size):
+		batch = random.sample(list(self.buffer), int(min(batch_size, self.__len__())))
+		batch = self.transition_func(*zip(*batch))
+		return batch()
diff --git a/cardio_rl/buffers/iet_buffer.py b/cardio_rl/buffers/iet_buffer.py
new file mode 100644
index 0000000..a6cc789
--- /dev/null
+++ b/cardio_rl/buffers/iet_buffer.py
@@ -0,0 +1,71 @@
+import torch as th
+import torch as th
+import torch.nn as nn
+import numpy as np
+from sklearn.cluster import KMeans
+from cardio_rl.transitions import TorchTransition
+from cardio_rl.buffers.er_buffer import ErTable
+
+
+class IeTable:
+	def __init__(self, network, capacity=1_000_000, transition_func=TorchTransition, k_events=4, cluster_interval=10, percent=0.5):
+		self.tables = [ErTable(capacity//k_events, transition_func) for _ in range(k_events)]
+		self.buffer = ErTable(capacity, transition_func)
+		self.k_events = k_events
+		self.all_encoded_transitions = []
+		self.encoder = lambda x: x
+		self.cluster_interval = cluster_interval
+		self.percent = percent
+		self.t = 0
+		self.network = network
+
+	def __len__(self):
+		lengths = [len(self.tables[k]) for k in range(self.k_events)]
+		return sum(lengths)
+
+	def store(self, transition):
+		self.buffer.store(transition)
+		
+		if self.t % self.cluster_interval == 0:
+			self.kmeans = KMeans(n_clusters=self.k_events, n_init="auto").fit(self.all_encoded_transitions)
+
+		state, action, *_ = transition
+
+		'''With raw, concatonated vectors'''
+		# z = np.concatenate([state, action], axis=0)
+		'''With state embedding'''
+		z = self.network.embedding(th.from_numpy(state)).detach().numpy()
+
+		self.all_encoded_transitions.append(z)
+		cluster = self.kmeans.predict([z]).squeeze(0)
+
+		self.tables[cluster].store(transition)
+
+		self.t += 1
+
+	def sample(self, batch_size):
+		from_buffer = int(batch_size*(self.percent))
+		from_tables = int((batch_size-from_buffer)//self.k_events)
+		from_buffer += int(batch_size-(from_tables*self.k_events))	# ensure batch_size samples are received
+
+		s_batch, a_batch, r_batch, sp_batch, d_batch, _ = self.buffer.sample(from_buffer)
+
+		# lengths = []
+
+		for k in range(self.k_events):
+			# num_samples = int(min(, len(self.tables[k].table)))
+			
+			# Don't sample empty tables
+			if len(self.tables[k]) > 0:
+				s, a, r, s_p, d, _ = self.tables[k].sample(int(from_tables))
+				s_batch = th.cat([s_batch, s])
+				a_batch = th.cat([a_batch, a])
+				r_batch = th.cat([r_batch, r])
+				sp_batch = th.cat([sp_batch, s_p])
+				d_batch = th.cat([d_batch, d])
+
+			# lengths.append(len(self.tables[k]))
+		
+		# print(lengths)
+
+		return s_batch, a_batch, r_batch, sp_batch, d_batch, _
diff --git a/cardio_rl/gatherer.py b/cardio_rl/gatherer.py
new file mode 100644
index 0000000..d7f3a4c
--- /dev/null
+++ b/cardio_rl/gatherer.py
@@ -0,0 +1,227 @@
+from collections import deque
+from gymnasium import Env
+from cardio_rl.policies import BasePolicy
+from cardio_rl.logger import Logger
+import numpy as np
+import gymnasium as gym
+
+class Collector():
+    def __init__(
+            self,
+            rollout_len: int = 1,
+            warmup_len: int = 1_000,
+            n_step: int = 1,
+            take_every: int = 1,
+            logger_kwargs = dict(
+                    log_interval = 5_000,
+                    episode_window=50,
+                    tensorboard=False,
+                )
+        ) -> None:        
+
+        self.rollout_len = rollout_len        
+
+        if self.rollout_len == -1:
+            self.ret_if_term = True
+            self.rollout_len += 1000000
+        else: 
+            self.ret_if_term = False   
+
+        self.warmup_len = warmup_len
+        self.n_step = n_step   
+        self.take_every = take_every 
+        self.take_count = 0
+
+        if logger_kwargs:
+            self.logger = Logger(**logger_kwargs)
+        else:
+            self.logger = Logger()
+
+        # env initialisation
+        self.step_buffer = deque(maxlen=self.n_step)
+    
+    def _init_env(self, env: Env):
+        self.env = env
+        self.state, _ = self.env.reset()
+
+    def _init_policy(self, policy):
+        self.policy = policy
+
+    def _env_step(self, policy=None, warmup=False):
+        if warmup:
+            a = self.env.action_space.sample()
+        else:
+            a = policy(self.state, self.net)
+        s_p, r, d, t, info = self.env.step(a)
+        self.logger.step(r, d, t)
+        d = d or t
+        return (self.state, a, r, s_p, d, info), s_p, d, t
+
+    def warmup(            
+        self,
+        net = None,
+    ):        
+        # Maybe move this check to the runner?
+        if self.warmup_len == None:
+            return list(deque())
+        
+        self.net = net
+        gather_buffer = deque()
+        self.step_buffer = deque(maxlen=self.n_step)     
+
+        for _ in range(self.warmup_len):
+            transition, next_state, done, trun = self._env_step(self.policy, warmup=True)
+            self.step_buffer.append(transition)
+            
+            if len(self.step_buffer) == self.n_step:
+                
+                if self.take_count % self.take_every == 0:                    
+                    if self.n_step == 1:
+                        gather_buffer.append(*list(self.step_buffer))
+                    else:
+                        gather_buffer.append(list(self.step_buffer))
+ 
+                self.take_count += 1
+
+            self.state = next_state
+            if done or trun:
+                self.state, _ = self.env.reset()
+                self.step_buffer = deque(maxlen=self.n_step)
+
+        return list(gather_buffer)
+
+    def rollout(
+        self,
+        net,
+    ):        
+        self.net = net
+        gather_buffer = deque()   
+
+        for _ in range(self.rollout_len):
+            transition, next_state, done, trun = self._env_step(self.policy)
+            self.step_buffer.append(transition)
+            if len(self.step_buffer) == self.n_step:
+                
+                if self.take_count % self.take_every == 0:
+                    if self.n_step == 1:
+                        gather_buffer.append(*list(self.step_buffer))
+                    else:
+                        gather_buffer.append(list(self.step_buffer))
+ 
+                self.take_count += 1
+
+            self.state = next_state
+            if done or trun:
+                self.state, _ = self.env.reset()
+                self.step_buffer = deque(maxlen=self.n_step)
+                self.policy.reset()
+                if self.ret_if_term:
+                    return list(gather_buffer)
+
+        return list(gather_buffer)
+
+
+class VectorCollector():
+    def __init__(
+            self,
+            env,
+            num_envs = 2,
+            rollout_len = 1,
+            warmup_len = 0,
+            n_step = 1,
+            logger_kwargs = None
+        ) -> None:  
+        
+        # https://gymnasium.farama.org/api/vector/#async-vector-env            
+        env_list = [lambda: gym.wrappers.RecordEpisodeStatistics(env) for _ in range(num_envs)]
+        # Getting errors when using AsyncVectorEnv
+        self.env = gym.vector.SyncVectorEnv(env_list)
+        # self.env = gym.wrappers.RecordEpisodeStatistics(env)
+        
+        self.state, _ = self.env.reset()        
+
+        self.rollout_len = rollout_len        
+
+        if self.rollout_len == -1:
+            self.ret_if_term = True
+            self.rollout_len += 1000000
+        else: 
+            self.ret_if_term = False   
+
+        self.warmup_len = warmup_len
+        self.n_step = n_step        
+
+        # metrics
+        if logger_kwargs:
+            self.logger = Logger(n_envs=num_envs, **logger_kwargs)
+        else:
+            self.logger = Logger(n_envs=num_envs)
+
+    def init_policy(self, policy):
+        self.policy = policy
+
+    def _env_step(self, policy):
+        a = policy(self.state, self.net)
+        s_p, r, d, t, info = self.env.step(a)
+        self.logger.vector_step(r, d, t)
+        d_combined = [False] * len(d)
+
+        for i in range(len(d_combined)):
+            d_combined[i] = d[i] or t[i]
+            
+        return (self.state, a, r, s_p, d_combined, info), s_p, d, t
+
+    def warmup(            
+        self,
+        net = None,
+        policy = None,
+    ):
+        
+        # Maybe move this check to the runner?
+        if self.warmup_len == None:
+            return deque()
+        
+        warmup_policy = BasePolicy(self.env)
+        
+        self.net = net
+        gather_buffer = deque()
+        step_buffer = deque(maxlen=self.n_step)     
+
+        for _ in range(self.warmup_len):
+            self.total_steps += 1
+            a = policy(self.state, self.net)
+            s_p, r, d, t, info = self.env.step(a)
+
+            step_buffer.append([self.state, a, r, s_p, d])
+            if len(step_buffer) == self.n_step:
+
+                if self.n_step == 1:
+                    gather_buffer.append(*list(step_buffer))
+                else:
+                    gather_buffer.append(list(step_buffer))
+
+            self.state = s_p
+
+        return list(gather_buffer)
+
+    def rollout(
+        self,
+        net,
+        policy,
+    ):        
+        self.net = net
+        gather_buffer = deque()
+        step_buffer = deque(maxlen=self.n_step)   
+
+        for _ in range(self.rollout_len):
+            transition, next_state, _, _ = self._env_step(policy)
+            step_buffer.append(transition)
+            if len(step_buffer) == self.n_step:
+                if self.n_step == 1:
+                    gather_buffer.append(*list(step_buffer))
+                else:
+                    gather_buffer.append(list(step_buffer))
+
+            self.state = next_state
+
+        return list(gather_buffer)
\ No newline at end of file
diff --git a/cardio_rl/logger.py b/cardio_rl/logger.py
new file mode 100644
index 0000000..c59e320
--- /dev/null
+++ b/cardio_rl/logger.py
@@ -0,0 +1,87 @@
+from collections import deque
+from torch.utils.tensorboard import SummaryWriter
+import time
+import logging
+import numpy as np
+from datetime import datetime
+
+"""
+-Maybe change to Rich Logger
+-implement tensorboard next
+-make logger compatible with VectorCollector next
+
+# https://docs.python.org/3/howto/logging.html
+# https://www.tensorflow.org/tensorboard
+# https://pytorch.org/docs/stable/tensorboard.html
+"""
+
+class Logger():
+    def __init__(
+            self,
+            n_envs = 1,
+            log_interval=2000,
+            episode_window=20,
+            tensorboard=False,
+            log_dir=None,
+            exp_name='exp'
+        ) -> None:
+        
+        logging.basicConfig(format='%(asctime)s: %(message)s', datefmt=' %I:%M:%S %p', level=logging.INFO)
+        self.n_envs = n_envs
+        self.log_interval = log_interval
+        # self.episode_window = episode_window
+        self.tensorboard = tensorboard
+
+        if self.tensorboard:
+            dir = ''
+
+            if log_dir:
+                dir += log_dir + '/'
+
+            # time_key = str(int(time.time()//1))
+            
+            # Changed to the below to be more in line with Hydra
+            date_key = datetime.now().strftime('%Y-%m-%d')
+            time_key = datetime.now().strftime('%H-%M-%S')
+
+            dir += exp_name + '_' + date_key + '_' + time_key + '/'
+            self.writer = SummaryWriter(dir)
+        
+        self.timestep = 0
+        self.episodes = 0
+        if n_envs == 1:
+            self.running_reward = 0
+        else:
+            self.running_reward = np.zeros(n_envs)
+
+        self.episodic_rewards = deque(maxlen=episode_window)
+        
+    def step(self, reward, done, truncated):
+        self.timestep += 1
+        self.running_reward += reward
+
+        if done or truncated:
+            self.episodes += 1
+            self.episodic_rewards.append(self.running_reward)
+            self.running_reward = 0
+
+        if self.timestep % self.log_interval == 0:
+            logging.info(f'Timesteps: {self.timestep}, Episodes: {self.episodes}, Avg. reward is {np.mean(self.episodic_rewards)}')
+
+            if self.tensorboard:
+                self.writer.add_scalar('rollout/ep_rew_mean', np.mean(self.episodic_rewards), self.timestep)
+
+    def vector_step(self, rewards, dones, truns):
+        self.timestep += self.n_envs
+
+        self.running_reward += rewards
+
+        for i, (done, truncated) in enumerate(zip(dones, truns)):
+            if done or truncated:
+                self.episodes += 1
+                self.episodic_rewards.append(self.running_reward[i])
+                self.running_reward[i] = 0
+
+        if (self.timestep//self.n_envs) % self.log_interval == 0:
+            logging.info(f'Timesteps: {self.timestep}, Episodes: {self.episodes}, Avg. reward is {np.mean(self.episodic_rewards)}')
+
diff --git a/cardio_rl/policies/__init__.py b/cardio_rl/policies/__init__.py
new file mode 100644
index 0000000..3ef4901
--- /dev/null
+++ b/cardio_rl/policies/__init__.py
@@ -0,0 +1,19 @@
+from .basic import (BasePolicy, 
+                    WhitenoiseDeterministic, 
+                    EpsilonArgmax, 
+                    Gaussian, 
+                    NoisyNaf,
+                    Categorical,
+                    Beta)
+
+REGISTRY = {}
+
+REGISTRY["random"] = BasePolicy
+REGISTRY["whitenoise"] = WhitenoiseDeterministic
+REGISTRY["pinknoise"] = None
+REGISTRY["ounoise"] = None
+REGISTRY["argmax"] = EpsilonArgmax
+REGISTRY["gaussian"] = Gaussian
+REGISTRY["naf"] = NoisyNaf
+REGISTRY["categorical"] = Categorical
+REGISTRY["beta"] = Beta
diff --git a/src/policies/__init__.py b/cardio_rl/policies/basic.py
similarity index 54%
rename from src/policies/__init__.py
rename to cardio_rl/policies/basic.py
index 189102a..acd1c03 100644
--- a/src/policies/__init__.py
+++ b/cardio_rl/policies/basic.py
@@ -1,47 +1,37 @@
-# from .basic_policies import Epsilon_argmax_policy 
-
 import torch as th
 import numpy as np
-import random
-
-# should make as classes with __call__ functions
 
-class Basepolicy():
-    def __init__(self, env):
+class BasePolicy():
+    def __init__(self, env, recurrent=False, hidden_dims=0):
         self.env = env        
+        self.recurrent = recurrent
+        self.hidden_dims = hidden_dims
+        self.hidden = th.zeros(1, hidden_dims)
         
     def __call__(self, state, net):
         return self.env.action_space.sample()
     
+    def reset(self):
+        self.hidden = th.zeros(1, self.hidden_dims)
+        # for policies that require resets, like NoisyNet or OuNoise
+        # raise NotImplementedError
  
-class Epsilon_Deterministic_policy(Basepolicy):
+class WhitenoiseDeterministic(BasePolicy):
     def __init__(self, env):
         super().__init__(env)
-        self.eps = 0.9
-        self.noise = True 
         
     def __call__(self, state, net):
-        input = th.from_numpy(state).float()
-
-        if np.random.rand() > self.eps:
-            self.eps = max(0.05, self.eps*0.99)   
-            out = net(input)         
-
-            if self.noise:
-                mean = th.zeros_like(out)
-                noise = th.normal(mean=mean, std=0.1).clamp(-0.5, 0.5)
-                out = (out + noise)
-        
-            return out.clamp(-1, 1).detach().numpy()
-
-        else:
-            self.eps = max(0.1, self.eps*0.999)
-            return self.env.action_space.sample()
+        input = th.from_numpy(state).float()  
+        out = net(input)         
+        mean = th.zeros_like(out)
+        noise = th.normal(mean=mean, std=0.1)   # .clamp(-0.5, 0.5) # unsure if necessary... need to check other implementations
+        out = (out + noise)    
+        return out.clamp(-1, 1).detach().numpy()
                  
 
-class Epsilon_argmax_policy(Basepolicy):
-    def __init__(self, env, eps = 0.0, min_eps = 0.0, ann_coeff = 0.9):
-        super().__init__(env)
+class EpsilonArgmax(BasePolicy):
+    def __init__(self, env, eps=0.0, min_eps=0.0, ann_coeff=0.9, recurrent=False, hidden_dims=0):
+        super().__init__(env, recurrent, hidden_dims)
         self.eps = eps
         self.min_eps = min_eps
         self.ann_coeff = ann_coeff
@@ -51,16 +41,20 @@ def __call__(self, state, net):
 
         if np.random.rand() > self.eps:
             self.eps = max(self.min_eps, self.eps*self.ann_coeff)   
-            out = net(input).detach().numpy()
+            if self.recurrent:
+                out, self.hidden = net(input, hidden=self.hidden)
+                out = out.detach().numpy()
+            else: 
+                out = net(input).detach().numpy()
+
             return np.argmax(out)  
 
         else:
             self.eps = max(self.min_eps, self.eps*self.ann_coeff)
             return self.env.action_space.sample()
-        
-    
 
-class Gaussian_policy(Basepolicy):
+
+class Gaussian(BasePolicy):
     def __init__(self, env):
         super().__init__(env)
 
@@ -75,7 +69,7 @@ def __call__(self, state, net):
         return a_sampled.numpy() * self.env.action_space.high + 0
 
 
-class Noisy_naf_policy(Basepolicy):
+class NoisyNaf(BasePolicy):
     def __init__(self, env):
         super().__init__(env)
 
@@ -84,7 +78,8 @@ def __call__(self, state, net):
         out, _, _, _ = net(input)
         return out.detach().numpy()
 
-class Categorical_policy(Basepolicy):
+
+class Categorical(BasePolicy):
     def __init__(self, env):
         super().__init__(env)
 
@@ -93,12 +88,13 @@ def __call__(self, state, net):
         probs = net(input)
         dist = th.distributions.Categorical(probs)
         return dist.sample().detach().numpy()
+    
+class Beta(BasePolicy):
+    def __init__(self, env):
+        super().__init__(env)
 
-REGISTRY = {}
-
-REGISTRY["random"] = Basepolicy
-REGISTRY["epsilon_deterministic"] = Epsilon_Deterministic_policy
-REGISTRY["epsilon_argmax"] = Epsilon_argmax_policy
-REGISTRY["gaussian"] = Gaussian_policy
-REGISTRY["naf"] = Noisy_naf_policy
-REGISTRY["categorical"] = Categorical_policy
\ No newline at end of file
+    def __call__(self, state, net):
+        input = th.from_numpy(state).float()
+        alpha, beta = net(input)
+        dist = th.distributions.Beta(alpha, beta)
+        return dist.sample().detach().numpy()
\ No newline at end of file
diff --git a/src/rainbow_naf/noisy_linear.py b/cardio_rl/rainbow_naf/noisy_linear.py
similarity index 100%
rename from src/rainbow_naf/noisy_linear.py
rename to cardio_rl/rainbow_naf/noisy_linear.py
diff --git a/cardio_rl/runner.py b/cardio_rl/runner.py
new file mode 100644
index 0000000..4aeba1a
--- /dev/null
+++ b/cardio_rl/runner.py
@@ -0,0 +1,157 @@
+import logging
+from gymnasium import Env
+from .transitions import REGISTRY as tran_REGISTRY
+from .transitions import BaseTransition
+from .buffers.er_buffer import ErTable
+from .policies import BasePolicy, REGISTRY as pol_REGISTRY
+from cardio_rl.gatherer import Collector
+
+# https://stackoverflow.com/questions/40181284/how-to-get-random-sample-from-deque-in-python-3
+# faster replay memory
+
+class Runner():
+    def __init__(
+            self,
+            env: Env,
+            policy: BasePolicy = None,
+            capacity: int = None,
+            er_buffer: ErTable = None,
+            batch_size: int = 100,
+            collector: Collector = Collector(),
+            n_batches: int = 1,
+            reduce=True,
+            backend = 'numpy',
+        ) -> None:
+
+        # Can maybe remove environment as an argument of the runner
+        self.env = env
+
+        # Maybe combine sampler and capacity into one argument?
+        self.batch_size = batch_size
+        self.n_batches = n_batches
+
+        if er_buffer == None and capacity != None:
+            transition = self._set_up_transition(backend)
+            self.er_buffer = ErTable(capacity, transition)
+            self.sampler = True
+        elif capacity == None:            
+            self.sampler = False
+        else:
+            self.er_buffer = er_buffer
+            self.sampler = True
+
+        self.collector = collector
+        self.rollout_len = collector.rollout_len
+        self.warmup_len = collector.warmup_len
+        self.n_step = collector.n_step
+
+        self.reduce = reduce
+        self.backend = backend        
+
+        self.policy = self._set_up_policy(policy)
+        self.collector._init_policy(self.policy)
+        self.collector._init_env(self.env)
+        self._warm_start()            
+
+    def _warm_start(
+            self,
+            net=None,          
+        ):
+
+        batch = self.collector.warmup(net)
+
+        if self.sampler:
+            for transition in batch:
+                self.er_buffer.store(transition)
+        
+        logging.info('### Warm up finished ###')
+        pass
+
+    def _set_up_policy(self, policy):
+        if isinstance(policy, str):
+            return pol_REGISTRY[policy](self.env)
+
+        elif isinstance(policy, BasePolicy):
+            return policy
+
+        else:
+            # add warning
+            return 
+    
+    def _set_up_transition(self, backend):
+        """
+        Maybe change name of backend to transition_type, better discription
+        """
+        if isinstance(backend, str):
+            return tran_REGISTRY[backend]
+
+        # isinstance(A, B) didn't work, this is a temp workaround
+        elif backend.__base__ is BaseTransition:
+            return backend
+
+        else:
+            # add warning
+            return 
+
+    def get_batch(
+            self,
+            net,
+        ):
+        
+        self.net = net        
+        rollout_batch = self.collector.rollout(net)
+
+        if not self.sampler:            
+            return self.prep_batch(rollout_batch)
+
+        else:
+            for transition in rollout_batch:
+                self.er_buffer.store(transition)
+
+            k = min(self.batch_size, len(self.er_buffer))
+            
+            if self.n_batches == 1:		
+                return self.prep_batch(self.er_buffer.sample(k))
+
+            else:
+                batch_samples = []
+                for _ in range(self.n_batches):
+                    batch_samples.append(self.prep_batch(self.er_buffer.sample(k))) 
+
+                return batch_samples
+
+
+    def prep_batch(
+        self,
+        batch
+        ):
+        """
+        takes the batch (which will be a list of transitons) and processes them to be seperate etc.
+        """
+        # need to redo after implementing replay buffer class
+
+        if self.n_step == 1:
+            return batch
+
+        elif self.reduce == False:
+            processed_batch = []
+            for n_step_transition in batch:
+                transition = n_step_transition
+                processed_batch.append([*transition])
+
+            return processed_batch
+        
+        else:
+            processed_batch = []
+            for n_step_transition in batch:
+                s, a, r, s_p, d, i = n_step_transition
+                s = s[0]
+                a = a[0]
+                r_list = list(r)
+                s_p = s_p[-1]
+                d = any(d)
+                i = i
+                processed_batch.append([s, a, r_list, s_p, d, i])
+
+            return self.transition(*zip(*processed_batch))
+        
\ No newline at end of file
diff --git a/cardio_rl/transitions/__init__.py b/cardio_rl/transitions/__init__.py
new file mode 100644
index 0000000..708d6a2
--- /dev/null
+++ b/cardio_rl/transitions/__init__.py
@@ -0,0 +1,9 @@
+from .transitions import (BaseTransition,
+                          TorchTransition,
+                          JaxTransition)
+
+REGISTRY = {}
+
+REGISTRY["numpy"] = BaseTransition
+REGISTRY["pytorch"] = TorchTransition
+REGISTRY["jax"] = JaxTransition
\ No newline at end of file
diff --git a/cardio_rl/transitions/transitions.py b/cardio_rl/transitions/transitions.py
new file mode 100644
index 0000000..65d2840
--- /dev/null
+++ b/cardio_rl/transitions/transitions.py
@@ -0,0 +1,54 @@
+from typing import NamedTuple
+import torch as th
+import numpy as np
+import jax.numpy as jnp
+
+"""
+All transition functions should output each vairable as shape [N, d]
+where N is the batch size and d is the dimension of that specific variable,
+i.e. d = [N, 1], r = [N, 1], s = [N, obs_dimensions]
+"""
+
+"""
+Transitions need to be revisited some time later, I got extermely 
+confused when trying to use them so it is likely unintuitive, each batch ends up
+having variables named s,a,r,s_p,d but you need to call the batch function itself to
+convert them to their respective back end? Maybe have the backend converstion as a method
+of the Transition class itself? That could be a good fix
+"""
+
+class BaseTransition(NamedTuple):
+    s: list # state
+    a: float    # action
+    r: float    # reward
+    s_p: list   # next state
+    d: int  # done
+    i: dict #info
+    
+    def __call__(self):
+        s = self.s
+        a = self.a
+        r = self.r
+        s_p = self.s_p
+        d = self.d
+        i = self.i
+        return s, a, r, s_p, d, i
+
+class TorchTransition(BaseTransition):
+    def __call__(self):
+        s = th.from_numpy(np.array(self.s)).float()
+        a = th.from_numpy(np.array(self.a)).unsqueeze(1).float()
+        r = th.from_numpy(np.array(self.r)).unsqueeze(1).float()
+        s_p = th.from_numpy(np.array(self.s_p)).float()
+        d = th.from_numpy(np.array(self.d)).unsqueeze(1).int()
+        return s, a, r, s_p, d, self.i
+    
+class JaxTransition(BaseTransition):
+    def __call__(self):
+        s = jnp.asarray(np.array(self.s))
+        a = jnp.asarray(np.array(self.a))
+        r = jnp.expand_dims(jnp.asarray(np.array(self.r)), -1)
+        s_p = jnp.asarray(np.array(self.s_p))
+        d = jnp.expand_dims(jnp.asarray(np.array(self.d)), -1)
+        return s, a, r, s_p, d, self.i
+    
\ No newline at end of file
diff --git a/cardio_rl/wrappers/__init__.py b/cardio_rl/wrappers/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/debug_env.py b/debug_env.py
new file mode 100644
index 0000000..d041f8a
--- /dev/null
+++ b/debug_env.py
@@ -0,0 +1,44 @@
+import gymnasium as gym
+from gymnasium import spaces
+import numpy as np
+
+class CardioDebugEnv(gym.Env):
+    def __init__(self, maxlen=10, discrete=True) -> None:
+        self.maxlen = maxlen
+        self.discrete = discrete
+        self._reset_env_vars()
+
+        self.action_space = spaces.Discrete(2)
+        self.observation_space = spaces.Box(0, self.maxlen, dtype=np.float32)
+
+    def step(self, action):
+        self.count +=1        
+        state = np.ones(5)*self.count
+        state[-1] = action
+        if self.count == self.maxlen:
+            return np.array(state), 1, True, False, {}
+        
+        return np.array(state), 1, False, False, {}        
+
+    def reset(self):
+        self._reset_env_vars()
+        return np.ones(5)*self.count, {}
+
+    def _reset_env_vars(self):
+        self.count = 0
+
+def main():
+    env = CardioDebugEnv()
+
+    s_t = env.reset()
+    while True:
+        print(s_t)
+        a_t = env.action_space.sample()
+        s_tp1, r, d, t, i = env.step(a_t)
+        print(s_t, a_t, s_tp1, r, d)
+        s_t = s_tp1
+        if d:
+            break
+
+if __name__ == '__main__':
+    main()
diff --git a/debugging.py b/debugging.py
new file mode 100644
index 0000000..ecb2ea5
--- /dev/null
+++ b/debugging.py
@@ -0,0 +1,54 @@
+import gymnasium as gym
+import numpy as np
+import torch as th
+from cardio_rl import Runner, Collector
+from cardio_rl.transitions import BaseTransition
+from typing import NamedTuple
+from debug_env import CardioDebugEnv
+
+class TestTransition(BaseTransition):
+    def __call__(self):
+        s = th.from_numpy(np.array(self.s)).float()
+        a = th.from_numpy(np.array(self.a)).unsqueeze(1).float()
+        r = th.from_numpy(np.array(self.r)).unsqueeze(1).float()
+        s_p = th.from_numpy(np.array(self.s_p)).float()
+        d = th.from_numpy(np.array(self.d)).unsqueeze(1).int()
+        return s, a, r, s_p, d, self.i
+
+env = CardioDebugEnv()
+
+"""
+N-step collection currently works in what may be an unexpected way,
+rollout steps are spent filling the n-step buffer, which is then flushed
+when the episode is terminated. i.e.
+In the DebugEnv, with n_step=4 and a rollout_len=10-13, only 7 transitions 
+are collected as rollout steps 0-3 are spent filling the buffer (you actually
+an error if rollout_steps < n_step), and then when the episode terminates, steps
+10-13 are spent filling it again. 
+Need to consider options/possibilities...
+"""
+
+runner = Runner(
+	env=env,
+	policy='random',
+    sampler=True,
+    capacity=1_000,
+    batch_size=2,
+	collector=Collector(
+		env=env,
+		rollout_len=2,
+        warmup_len=6,
+        n_step=4,
+        take_every=2,
+		),
+	reduce=False,
+	backend='pytorch'
+)
+
+batch: NamedTuple = runner.get_batch(net=None)
+
+# Needing to call batch as a function is awkward and unnecessary, should try find a workaround
+print(len(runner.er_buffer))
+# exit()
+s, a, *_ = batch()
+print(s, a)
diff --git a/implicit_event_tables/sb3_benchmarking.py b/implicit_event_tables/sb3_benchmarking.py
new file mode 100644
index 0000000..fa6b771
--- /dev/null
+++ b/implicit_event_tables/sb3_benchmarking.py
@@ -0,0 +1,49 @@
+import gymnasium as gym
+import numpy as np
+from stable_baselines3 import TD3
+from stable_baselines3.common.noise import NormalActionNoise
+from gymnasium.wrappers import RescaleAction
+
+
+env_name = 'BipedalWalker-v3'
+env = gym.make(env_name)
+env = RescaleAction(env, -1.0, 1.0)
+
+n_actions = env.action_space.shape[-1]
+action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions))
+
+"""
+SB3 zoo hyperparams:
+  n_timesteps: !!float 3e5
+  policy: 'MlpPolicy'
+  gamma: 0.98
+  buffer_size: 200000
+  learning_starts: 10000
+  noise_type: 'normal'
+  noise_std: 0.1
+  gradient_steps: -1
+  train_freq: [1, "episode"]
+  learning_rate: !!float 1e-3
+  policy_kwargs: "dict(net_arch=[400, 300])"
+"""
+
+model = TD3(
+    "MlpPolicy",
+    env,
+    verbose=1,
+    action_noise=action_noise, 
+    gamma=0.98,
+    train_freq=1,
+    buffer_size=200_000,
+    learning_starts=10_000,
+    learning_rate=1e-3,
+    policy_kwargs=dict(net_arch=[400, 300]),
+    stats_window_size=50,
+    tensorboard_log=f'implicit_event_tables/tb_logs/{env_name}/')
+
+model.learn(
+    total_timesteps=500_000,
+    log_interval=3,
+    tb_log_name='sb3_td3_benchmark_long_cap200k',
+    progress_bar=True
+    )
\ No newline at end of file
diff --git a/implicit_event_tables/td3_w_ietable.py b/implicit_event_tables/td3_w_ietable.py
new file mode 100644
index 0000000..a49e712
--- /dev/null
+++ b/implicit_event_tables/td3_w_ietable.py
@@ -0,0 +1,175 @@
+import copy
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from gymnasium.wrappers import RescaleAction
+from cardio_rl import Runner, Collector
+from cardio_rl.buffers.iet_buffer import IeTable
+
+
+class Q_critic(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Q_critic, self).__init__()
+
+		self.net1 = nn.Sequential(
+			nn.Linear(state_dim+action_dim, 400),
+			nn.ReLU(),
+			nn.Linear(400, 300),
+			nn.ReLU(),
+			nn.Linear(300, 1))
+		
+		self.net2 = nn.Sequential(
+			nn.Linear(state_dim+action_dim, 400),
+			nn.ReLU(),
+			nn.Linear(400, 300),
+			nn.ReLU(),
+			nn.Linear(300, 1))
+
+	def forward(self, state, action):
+		sa = th.concat([state, action], dim=-1)
+		return self.net1(sa), self.net2(sa)
+
+class Policy(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Policy, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(state_dim, 400),
+			nn.ReLU(),
+			nn.Linear(400, 300),
+			nn.ReLU(),
+			nn.Linear(300, action_dim),
+			nn.Tanh())
+		
+	def forward(self, state):
+		return self.net(state)
+
+class NguEncoder(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(NguEncoder, self).__init__()
+
+		self.embedding = nn.Sequential(
+			nn.Linear(state_dim, 64),
+			nn.ReLU(),
+			nn.Linear(64, 32),
+			nn.ReLU(),
+			nn.Linear(32, 16))
+		
+		self.regressor = nn.Linear(32, action_dim)
+		
+	def forward(self, state, next_state):
+		z_st = self.embedding(state)
+		z_stp1 = self.embedding(next_state)
+		z_joint = th.concat([z_st, z_stp1], dim=-1)
+		return F.tanh(self.regressor(z_joint))
+
+env_name = 'BipedalWalker-v3'
+env = gym.make(env_name)
+env = RescaleAction(env, -1.0, 1.0)
+
+observation_dim = env.observation_space.shape[0]
+action_dim = env.action_space.shape[0]
+
+emb_net = NguEncoder(observation_dim, action_dim)
+
+k_events = 4
+interval = 50
+implicit_event_table = IeTable(network=emb_net, capacity=200_000, k_events=k_events, cluster_interval=interval)
+iet_warmup = 100
+
+emb_optimizer = th.optim.Adam(implicit_event_table.network.parameters(), lr=1e-3)
+
+s_t, _ = env.reset()
+for _ in range(iet_warmup):
+	"""
+	Try to think of a way to remove needing to warmup (train emb network and fill the representation buffer)
+	"""
+	a_t = env.action_space.sample()
+	s_tp1, r_t, d_t, t_t, _ = env.step(a_t)
+	transition = s_t, a_t, r_t, s_tp1, d_t
+	'''With raw, concatonated vectors'''
+	# sa = np.concatenate([s_t, a_t], axis=0)
+	# implicit_event_table.all_encoded_transitions.append(sa)
+	'''With state embedding'''
+	z = implicit_event_table.network.embedding(th.from_numpy(s_t)).detach().numpy()
+	implicit_event_table.all_encoded_transitions.append(z)
+	s_t = s_tp1
+	if d_t or t_t:
+		s_t, _ = env.reset()
+
+s_t, _ = env.reset()
+
+runner = Runner(
+	env=env,
+	policy='whitenoise',
+	sampler=True,
+	# capacity=1_000_000,
+	batch_size=100,	# batch_size needs to exceed value of k_events
+	collector=Collector(
+		env=env,
+		warmup_len=10_000,
+		logger_kwargs=dict(
+			log_interval = 5_000,
+			episode_window=50,
+            tensorboard=True,
+			log_dir = 'implicit_event_tables/tb_logs/' + env_name,
+		    exp_name = f'ngu_embedding_cap200k_iet_k{k_events}_int{interval}_warm{iet_warmup}'
+		)
+	),
+	er_buffer=implicit_event_table
+)
+
+# https://stable-baselines3.readthedocs.io/en/master/modules/ddpg.html
+critic = Q_critic(observation_dim, action_dim)
+actor = Policy(observation_dim, action_dim)
+targ_critic = copy.deepcopy(critic)
+targ_actor = copy.deepcopy(actor)
+c_optimizer = th.optim.Adam(critic.parameters(), lr=1e-3)
+a_optimizer = th.optim.Adam(actor.parameters(), lr=1e-3)
+
+for steps in range(590_000):
+	s, a, r, s_p, d, _ = runner.get_batch(actor)
+
+	with th.no_grad():
+		a_p = targ_actor(s_p)
+
+		mean = th.zeros_like(a_p)
+		noise = th.normal(mean=mean, std=0.2).clamp(-0.5, 0.5)	# target policy noise value in SB3
+		
+		a_p = (a_p + noise).clamp(-1, 1)
+
+		q_p = th.min(*targ_critic(s_p, a_p))
+		y = r + 0.98 * q_p * (1 - d)
+
+	q1, q2 = critic(s, a.squeeze(1))
+		
+	loss = F.mse_loss(q1, y) + F.mse_loss(q2, y)
+
+	c_optimizer.zero_grad()
+	loss.backward()
+	th.nn.utils.clip_grad_norm_(critic.parameters(), 0.5)
+	c_optimizer.step()
+
+	'''Maybe implement a mean and std. var instead of explicit action predictions'''
+	a_pred = implicit_event_table.network(s, s_p)
+	emb_loss = F.mse_loss(a.squeeze(1), a_pred)
+	emb_optimizer.zero_grad()
+	emb_loss.backward()
+	emb_optimizer.step()
+
+	if steps % 2 == 0:
+		sa = th.concat([s_p, actor(s_p)], dim=-1)
+		policy_loss =  -(critic.net1(sa).mean())	# critic net1 as in sb3
+
+		a_optimizer.zero_grad()
+		policy_loss.backward()
+		th.nn.utils.clip_grad_norm_(actor.parameters(), 0.5)
+		a_optimizer.step()	
+
+		for targ_params, params in zip(targ_critic.parameters(), critic.parameters()):
+			targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
+
+		for targ_params, params in zip(targ_actor.parameters(), actor.parameters()):
+			targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
\ No newline at end of file
diff --git a/implicit_event_tables/td3_w_uerbuffer.py b/implicit_event_tables/td3_w_uerbuffer.py
new file mode 100644
index 0000000..4dbfd05
--- /dev/null
+++ b/implicit_event_tables/td3_w_uerbuffer.py
@@ -0,0 +1,119 @@
+import copy
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from gymnasium.wrappers import RescaleAction
+from cardio_rl import Runner, Collector
+from cardio_rl.buffers.iet_buffer import IeTable
+
+
+class Q_critic(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Q_critic, self).__init__()
+
+		self.net1 = nn.Sequential(
+			nn.Linear(state_dim+action_dim, 400),
+			nn.ReLU(),
+			nn.Linear(400, 300),
+			nn.ReLU(),
+			nn.Linear(300, 1))
+		
+		self.net2 = nn.Sequential(
+			nn.Linear(state_dim+action_dim, 400),
+			nn.ReLU(),
+			nn.Linear(400, 300),
+			nn.ReLU(),
+			nn.Linear(300, 1))
+
+	def forward(self, state, action):
+		sa = th.concat([state, action], dim=-1)
+		return self.net1(sa), self.net2(sa)
+
+class Policy(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Policy, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(state_dim, 400),
+			nn.ReLU(),
+			nn.Linear(400, 300),
+			nn.ReLU(),
+			nn.Linear(300, action_dim),
+			nn.Tanh())
+		
+	def forward(self, state):
+		return self.net(state)
+
+env_name = 'BipedalWalker-v3'
+env = gym.make(env_name)
+env = RescaleAction(env, -1.0, 1.0)
+
+observation_dim = env.observation_space.shape[0]
+action_dim = env.action_space.shape[0]
+
+runner = Runner(
+	env=env,
+	policy='whitenoise',
+	sampler=True,
+	capacity=200_000,
+	batch_size=100,	# batch_size needs to exceed value of k_events
+	collector=Collector(
+		env=env,
+		warmup_len=10_000,
+		logger_kwargs=dict(
+			log_interval = 5_000,
+			episode_window=50,
+            tensorboard=True,
+			log_dir = 'implicit_event_tables/tb_logs/' + env_name,
+		    exp_name = f'vanilla_td3'
+		)
+	),
+)
+
+# https://stable-baselines3.readthedocs.io/en/master/modules/ddpg.html
+critic = Q_critic(observation_dim, action_dim)
+actor = Policy(observation_dim, action_dim)
+targ_critic = copy.deepcopy(critic)
+targ_actor = copy.deepcopy(actor)
+c_optimizer = th.optim.Adam(critic.parameters(), lr=1e-3)
+a_optimizer = th.optim.Adam(actor.parameters(), lr=1e-3)
+
+for steps in range(590_000):
+	s, a, r, s_p, d, _ = runner.get_batch(actor)
+
+	with th.no_grad():
+		a_p = targ_actor(s_p)
+
+		mean = th.zeros_like(a_p)
+		noise = th.normal(mean=mean, std=0.2).clamp(-0.5, 0.5)	# target policy noise value in SB3
+		
+		a_p = (a_p + noise).clamp(-1, 1)
+
+		q_p = th.min(*targ_critic(s_p, a_p))
+		y = r + 0.98 * q_p * (1 - d)
+
+	q1, q2 = critic(s, a.squeeze(1))
+		
+	loss = F.mse_loss(q1, y) + F.mse_loss(q2, y)
+
+	c_optimizer.zero_grad()
+	loss.backward()
+	th.nn.utils.clip_grad_norm_(critic.parameters(), 0.5)
+	c_optimizer.step()
+
+	if steps % 2 == 0:
+		sa = th.concat([s_p, actor(s_p)], dim=-1)
+		policy_loss =  -(critic.net1(sa).mean())	# critic net1 as in sb3
+
+		a_optimizer.zero_grad()
+		policy_loss.backward()
+		th.nn.utils.clip_grad_norm_(actor.parameters(), 0.5)
+		a_optimizer.step()	
+
+		for targ_params, params in zip(targ_critic.parameters(), critic.parameters()):
+			targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
+
+		for targ_params, params in zip(targ_actor.parameters(), actor.parameters()):
+			targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
\ No newline at end of file
diff --git a/implicit_event_tables/to-do.md b/implicit_event_tables/to-do.md
new file mode 100644
index 0000000..2a5c676
--- /dev/null
+++ b/implicit_event_tables/to-do.md
@@ -0,0 +1,11 @@
+# To do list
+* [ ] Benchmark TD3 implementation with SB3!
+* [ ] set up proper experimenting pipeline, benchmark TD3 implementation etc.
+* [x] add central buffer that all transitions are added to and create a sampling ratio between central buffer and event tables
+* [ ] implement encoder architecture and gradient updates
+    * [x] Siamese action regression network (only state)
+    * [ ] VAE (can be everything)
+    * [ ] TD7-like (state and action)
+
+# Notes
+* Platueaing could be due to replay buffer being too large, SB3 uses a size of 200,000 for BipedalWalker, test this hypothesis with SB3
\ No newline at end of file
diff --git a/main.py b/main.py
new file mode 100644
index 0000000..e6e98c8
--- /dev/null
+++ b/main.py
@@ -0,0 +1,91 @@
+import copy
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from cardio_rl import Runner
+from cardio_rl import Collector
+from cardio_rl.policies import EpsilonArgmax
+
+"""
+A simple DQN implementation!
+"""
+
+class ImplicitQ(nn.Module):
+	"""
+	Implicit Quantile Network: https://arxiv.org/abs/1806.06923
+	-this helped a lot: https://datascience.stackexchange.com/questions/40874/how-does-implicit-quantile-regression-network-iqn-differ-from-qr-dqn
+	-still need to implement quantile huber loss!
+	"""
+	def __init__(self, state_dim, action_dim, emb_dim=64, n_samples=10):
+		super(ImplicitQ, self).__init__()
+		self.l1 = nn.Linear(state_dim, 128)
+		self.l2 = nn.Linear(128, 128)
+		self.l3 = nn.Linear(128, action_dim)
+
+		self.emb_dim = emb_dim
+		self.embedding = nn.Linear(self.emb_dim, 128)
+		self.n_samples = n_samples
+
+	def forward(self, state, reduce=True):
+		q = F.relu(self.l1(state))
+
+		tau = th.zeros(self.n_samples, 1).uniform_()
+		embedding = th.arange(self.emb_dim) * th.pi * tau
+		embedding = th.cos(embedding)
+
+		embedding = self.embedding(embedding)
+		q = q.unsqueeze(-2) * embedding
+
+		q = F.relu(self.l2(q))
+		q = self.l3(q)
+
+		if reduce:
+			q = th.mean(q, dim=-2)
+		
+		return q
+	
+def main():
+	env = gym.make('CartPole-v1')
+
+	runner = Runner(
+		env=env,
+		policy=EpsilonArgmax(env, 0.5, 0.05, 0.9),
+		sampler=True,
+		capacity=100000,
+		batch_size=64,
+		collector=Collector(
+			env=env,
+			rollout_len=4,
+			warmup_len=1000,
+		),
+		backend='pytorch'
+	)
+
+	critic = ImplicitQ(4, 2)
+	targ_net: nn.Module = copy.deepcopy(critic)
+	optimizer = th.optim.Adam(critic.parameters(), lr=2.3e-3)
+	gamma = 0.99
+	target_update = 10
+
+	for t in range(10000):
+		batch = runner.get_batch(critic)
+		s, a, r, s_p, d, _ = batch()
+
+		with th.no_grad():
+			q_p = th.max(targ_net(s_p), keepdim=True, dim=-1).values
+			y = r + gamma * q_p * (1 - d)
+
+		q = critic(s).gather(1, a.long())
+
+		loss = F.mse_loss(q, y)
+
+		optimizer.zero_grad()
+		loss.backward()
+		optimizer.step()
+
+		if t % target_update == 0:        
+			targ_net = copy.deepcopy(critic)
+			
+if __name__ == '__main__':
+	main()
\ No newline at end of file
diff --git a/outdated_a2c.py b/outdated_a2c.py
new file mode 100644
index 0000000..724cf08
--- /dev/null
+++ b/outdated_a2c.py
@@ -0,0 +1,94 @@
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Categorical
+import gymnasium as gym
+from cardio_rl import Runner, VectorCollector
+
+class Model(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Model, self).__init__()
+
+		self.critic = nn.Sequential(
+			nn.Linear(state_dim, 64),
+			nn.ReLU(),
+			nn.Linear(64, 64),
+			nn.ReLU(),
+			nn.Linear(64, 1))
+
+		self.actor = nn.Sequential(
+			nn.Linear(state_dim, 64),
+			nn.ReLU(),
+			nn.Linear(64, 64),
+			nn.ReLU(),
+			nn.Linear(64, action_dim),
+			nn.Softmax(dim=-1))
+
+	def forward(self, state):
+		return self.actor(state)
+
+env = gym.make('CartPole-v1')
+
+runner = Runner(
+	env=env,
+	policy='categorical',
+	collector=VectorCollector(
+		env=env,
+		num_envs=8,
+		rollout_len=5,
+	)
+)
+
+net = Model(4, 2)
+optimizer = th.optim.RMSprop(net.parameters(), lr=7e-4, eps=1e-5)
+
+"""
+Need to implement parallel environments
+-Implement timestep-based logging!
+-Debug and benchmark, doesn't seem like it works rn 
+-Maybe merge actor and critic into one nn.Module
+
+https://github.com/sweetice/Deep-reinforcement-learning-with-pytorch/blob/master/Char04%20A2C/A2C.py
+"""
+
+timesteps = 0
+
+for rollout_steps in range(10_000):
+
+	timesteps += 5*2
+
+	batch = runner.get_batch(net)
+
+	s, a, r, s_p, d = batch()
+
+	values = net.critic(s)
+	returns = th.zeros_like(values)
+
+	ret = th.zeros_like(r)
+	running_r = net.critic(s_p[-1]).squeeze(-1).detach()
+
+	for t in reversed(range(len(r))):
+		running_r = r[t] + 0.99 * running_r * (1 - d[t])
+		ret[t] = running_r
+
+	critic_loss = F.mse_loss(returns, values)
+	
+	adv = (returns - values).squeeze(-1).detach()
+
+	dist = Categorical(net(s))	
+	log_probs = dist.log_prob(a.squeeze(1))
+	ent = dist.entropy().sum(-1).mean()
+
+	policy_loss = -(log_probs * adv.detach()).mean()
+
+	if timesteps % 2000 == 0:
+		# Check if the trnopy loss recorded in sb3 is before or after being multiplied by the 
+		# entropy coefficient
+		print(f'{timesteps} = Policy loss: {policy_loss}, Critic loss: {critic_loss}, Entropy loss: {ent}')
+
+	loss = policy_loss +  (0.5 * critic_loss) # - (0.001 * ent) 
+
+	optimizer.zero_grad()
+	loss.backward()
+	nn.utils.clip_grad_norm_(net.parameters(), 0.5)
+	optimizer.step()
diff --git a/proof_mpo.py b/proof_mpo.py
new file mode 100644
index 0000000..bfd7fe2
--- /dev/null
+++ b/proof_mpo.py
@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+import sys
+
+n = nn.Parameter(torch.ones(1)*100)
+a = nn.Parameter()
+e = 0.2
+
+optim = torch.optim.Adam(params=[n], lr=1.)
+
+# example_pi = torch.tensor([0.6, 0.4]).float()
+
+example_pi = nn.Sequential(
+    nn.Linear(2, 32),
+    nn.Tanh(),
+    nn.Linear(32, 32),
+    nn.Tanh(),
+    nn.Linear(32, 2),
+    nn.Softmax(dim=-1)
+)
+
+example_Q = torch.randint(low=-5, high=10, size=[3,2]).float()
+
+# print(example_Q)
+
+# print(example_pi(example_Q))
+
+# sys.exit()
+
+def expectation(Q, policy):
+    return torch.sum(Q * policy, dim=-1, keepdim=True)
+
+# print(f"Expectation of random policy: {expectation(example_Q, example_pi(example_Q)).squeeze()}")
+# print(f"Expectation of q_ij policy: {expectation(example_Q, torch.softmax(example_Q/n, -1)).squeeze()}")
+n = 10
+for i in range(50):
+    # n.zero_grad()
+
+    def dual(n):
+        avg_over_actions = torch.exp(example_Q/n)
+        # print(example_Q)
+        # print(avg_over_actions)
+        avg_over_actions = avg_over_actions.mean(dim=-1)
+        # print(avg_over_actions)
+
+        avg_over_states = avg_over_actions.log()
+        # print(avg_over_states)
+        avg_over_states = avg_over_states.mean(dim=-1)
+        # print(avg_over_states)
+        loss = n*e + n*avg_over_states
+
+        return nn.functional.softplus(loss)
+    
+    loss = dual(n=n)
+    print(f'{i}: {loss.item(), n}')
+
+    n = loss # *0.1
+
+    continue
+    
+    # print(f'{i}: {loss.item(), n.item(), act_weights, expectation(example_Q, act_weights)}')
+    loss.backward()
+    optim.step()
+
+    n = loss
+
+print(f"Expectation of random policy: {expectation(example_Q, example_pi(example_Q)).squeeze()}")
+print(f"Expectation of q_ij policy: {expectation(example_Q, torch.softmax(example_Q/n, -1)).squeeze()}")
diff --git a/sb3_benchmarking.py b/sb3_benchmarking.py
new file mode 100644
index 0000000..e7862c4
--- /dev/null
+++ b/sb3_benchmarking.py
@@ -0,0 +1,44 @@
+import gymnasium as gym
+import numpy as np
+from stable_baselines3 import A2C, DDPG, TD3, SAC
+from stable_baselines3.common.noise import NormalActionNoise
+
+# env = gym.make("CartPole-v1")
+# print('test')
+# model = A2C("MlpPolicy", env, verbose=1)
+# model.learn(total_timesteps=50_000)
+
+# env = gym.make('MountainCarContinuous-v0')
+# print('test')
+# n_actions = env.action_space.shape[-1]
+# action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions)
+# model = DDPG("MlpPolicy", env, verbose=1, action_noise=action_noise)
+# model.learn(total_timesteps=300_000, log_interval=10_000)
+
+env_name = 'BipedalWalker-v3'
+env = gym.make(env_name)
+print('test')
+n_actions = env.action_space.shape[-1]
+action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions))
+model = SAC(
+    "MlpPolicy",
+    env,
+    verbose=1,
+    action_noise=action_noise, 
+    gamma=0.98,
+    tau=0.02,
+    train_freq=64,
+    gradient_steps=64,
+    buffer_size=300_000,
+    learning_starts=10_000,
+    learning_rate=7.3e-4,
+    policy_kwargs=dict(net_arch=[400, 300]),
+    # stats_window_size=40,
+    tensorboard_log='BipedalWalker-v3/')
+
+model.learn(
+    total_timesteps=300_000,
+    log_interval=6,
+    tb_log_name='vanilla',
+    progress_bar=True
+    )
\ No newline at end of file
diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000..a9048ee
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,34 @@
+import os
+from setuptools import setup, find_packages
+
+def read(fname):
+    return open(os.path.join(os.path.dirname(__file__), fname)).read()
+
+setup(
+    name = 'cardio_rl',
+    packages=find_packages(),
+    install_requires=[
+        "numpy>=1.21.0",
+        "gymnasium>=0.28.0",
+        'tensorboard'
+    ],
+    description = "Cardio RL. In development...",  
+    author = 'Manus McAuliffe',
+    url = 'https://github.com/mmcaulif/GymCardio',
+    author_email = "mmcaulif@tcd.ie",  
+    license = 'MIT',
+    long_description=read('README.md'),
+    long_description_content_type="text/markdown",
+    version = '0.0.4',
+    python_requires=">=3.7",
+    
+    # PyPI package information.
+    classifiers=[
+        'Development Status :: 3 - Alpha',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
+        'Programming Language :: Python :: 3.10',
+    ],
+)
diff --git a/src/gatherer.py b/src/gatherer.py
deleted file mode 100644
index 9de27c8..0000000
--- a/src/gatherer.py
+++ /dev/null
@@ -1,94 +0,0 @@
-from collections import deque
-import torch as th
-import numpy as np
-
-class Collector():
-    def __init__(
-            self,
-            env,
-            capacity,
-        ) -> None:
-
-        self.env = env
-        self.capacity = capacity
-        
-        # environment init (move to function)
-        self.state, _ = self._reset()
-
-        # metrics
-        self.episodes = 0
-        self.total_steps = 0
-        self.ep_rew = 0
-        self.epsiode_window = deque(maxlen=50)
-
-        pass
-
-    def _step(
-        self,
-        a
-    ):
-        return self.env.step(a)
-
-    def _reset(
-        self
-    ):
-        return self.env.reset()
-    
-    def agent_step(
-        self,
-        policy,
-        state=None
-    ):        
-        return policy(self.state, self.net)
-
-    def rollout(
-        self,
-        net,
-        policy,
-        length,
-        n_step
-    ):
-        
-        self.net = net
-
-        gather_buffer = deque()
-        step_buffer = deque(maxlen=n_step)
-
-        if length == -1:
-            ret_when_term = True
-            length += 1000000
-        else: 
-            ret_when_term = False        
-
-        for _ in range(length):
-            self.total_steps += 1
-            a = self.agent_step(policy)
-            s_p, r, d, t, info = self._step(a)
-
-            # metrics
-            self.ep_rew += r
-
-            step_buffer.append([self.state, a, r, s_p, d])
-            if len(step_buffer) == n_step:
-                # print(*list(step_buffer))
-
-                if n_step == 1:
-                    gather_buffer.append(*list(step_buffer))
-                else:
-                    gather_buffer.append(list(step_buffer))
-
-            self.state = s_p
-            if d or t:
-                self.episodes += 1
-                self.epsiode_window.append(self.ep_rew)
-                self.ep_rew = 0
-                self.state, _ = self._reset()
-                step_buffer = deque(maxlen=n_step)
-
-                if self.episodes % 10 == 0:
-                    print(f"Average reward after {self.episodes} episodes or {self.total_steps} timesteps: {np.mean(self.epsiode_window)}")
-
-                if ret_when_term:
-                    return list(gather_buffer)
-
-        return list(gather_buffer)
\ No newline at end of file
diff --git a/src/main.py b/src/main.py
deleted file mode 100644
index 075741f..0000000
--- a/src/main.py
+++ /dev/null
@@ -1,105 +0,0 @@
-import copy
-import torch as th
-import torch.nn as nn
-import torch.nn.functional as F
-import gymnasium as gym
-import numpy as np
-from policies import Basepolicy
-from runner import get_offpolicy_runner
-
-"""
-Duelling Double DQN with n-step returns used for verification of runner
-"""
-
-class Q_duelling(nn.Module):
-	def __init__(self, state_dim, action_dim, dueling):
-		super(Q_duelling, self).__init__()
-		self.dueling = dueling
-
-		self.l1 = nn.Linear(state_dim, 128)
-
-		self.adv = nn.Sequential(
-			nn.Linear(128, 128),
-			nn.ReLU(),
-			nn.Linear(128, action_dim)
-		)
-
-		self.val = nn.Sequential(
-			nn.Linear(128, 128),
-			nn.ReLU(),
-			nn.Linear(128, 1)
-		)
-
-	def forward(self, state):
-		q = F.relu(self.l1(state))
-		v = self.val(q)
-		a = self.adv(q)
-
-		if self.dueling:
-			return v + (a - a.mean())
-
-		return a
-
-class Epsilon_argmax_policy(Basepolicy):
-    def __init__(self, env, eps = 0.0, min_eps = 0.0, ann_coeff = 0.9):
-        super().__init__(env)
-        self.eps = eps
-        self.min_eps = min_eps
-        self.ann_coeff = ann_coeff
-
-    def __call__(self, state, net):
-        input = th.from_numpy(state).float()
-
-        if np.random.rand() > self.eps:
-            self.eps = max(self.min_eps, self.eps*self.ann_coeff)   
-            out = net(input).detach().numpy()
-            return np.argmax(out)  
-
-        else:
-            self.eps = max(self.min_eps, self.eps*self.ann_coeff)
-            return self.env.action_space.sample()
-
-env = gym.make('CartPole-v1')
-runner = get_offpolicy_runner(
-	env, 
-	Epsilon_argmax_policy(env, 0.5, 0.05, 0.9),
-	length=4,
-	capacity=100000, 
-	batch_size=512, 
-	train_after=100,
-	n_step=3)
-
-net = Q_duelling(4, 2, True)
-targ_net = copy.deepcopy(net)
-optimizer = th.optim.Adam(net.parameters(), lr=2.3e-3)
-gamma = 0.99
-target_update = 10
-
-for t in range(10000):
-	batch = runner.get_batch(net)
-	s, a, r, s_p, d = batch()
-
-	if runner.n_step !=  1:
-		r_list = r.squeeze(1)
-		r_nstep = th.zeros(len(r))
-		for n in reversed(range(runner.n_step)):
-			r_nstep += r_list[:,n] * gamma
-		r = r_nstep.unsqueeze(-1)
-
-	q = net(s).gather(1, a.long())
-
-	with th.no_grad():
-		a_p = net(s_p).argmax(dim=1).unsqueeze(1)		
-		q_p = targ_net(s_p).gather(1, a_p.long())		
-		y = r + gamma * q_p * (1 - d)
-
-	loss = F.mse_loss(q, y)
-
-	optimizer.zero_grad()
-	loss.backward()
-	th.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
-	optimizer.step()
-
-	if t % target_update == 0:        
-		targ_net = copy.deepcopy(net)
-		
\ No newline at end of file
diff --git a/src/policies/basic_policies b/src/policies/basic_policies
deleted file mode 100644
index 0b5a73a..0000000
--- a/src/policies/basic_policies
+++ /dev/null
@@ -1,30 +0,0 @@
-import torch as th
-import numpy as np
-
-# should make as classes with __call__ functions
-
-class Random_policy():
-    def __init__(self, env):
-        self.env = env        
-        
-    def __call__(self, state):
-        return self.env.action_space.sample()
-    
-class Epsilon_argmax_policy(Random_policy):
-    def __init__(self):
-        super().__init__()
-
-    def __call__(self, state, net, epsilon=None):
-        input = th.from_numpy(state).float()
-        out = net(input).detach().numpy()
-        return np.argmax(out)
-
-def gaussian_policy(net, state, env):
-    input = th.from_numpy(state).float()
-    out = net(input)
-
-    mean, std = out[0].detach().numpy(), out[1].detach().numpy()
-
-    noise = np.random.uniform()
-
-    return (np.tanh(mean + (std * noise)) * env.action_space.high)
\ No newline at end of file
diff --git a/src/runner.py b/src/runner.py
deleted file mode 100644
index 4007236..0000000
--- a/src/runner.py
+++ /dev/null
@@ -1,148 +0,0 @@
-from collections import deque
-from gatherer import Collector
-from transitions import TorchTransition
-from policies import REGISTRY as pol_REGISTRY
-from policies import Basepolicy
-import random
-import numpy as np
-
-# https://stackoverflow.com/questions/40181284/how-to-get-random-sample-from-deque-in-python-3
-# faster replay memory
-
-def get_offpolicy_runner(
-        env,
-        policy,
-        length,
-        capacity,
-        batch_size,
-        n_step=1,
-        collector=Collector,
-        train_after=10000):
-    
-    return Runner(
-        env,
-        policy,
-        length,
-        True,
-        capacity,
-        batch_size,
-        n_step,
-        collector,
-        train_after)
-
-def get_onpolicy_runner(
-        env,
-        policy,
-        length,
-        collector=Collector,):
-    
-    return Runner(
-        env,
-        policy,
-        length,
-        False,
-        None,
-        None,
-        collector,
-        0)
-
-class Runner():
-    def __init__(
-            self,
-            env,
-            policy,
-            length,
-            sampler,
-            capacity,
-            batch_size,
-            n_step,
-            collector=Collector,
-            train_after=None
-        ) -> None:
-
-        self.env = env
-        self.length = length
-        self.sampler = sampler
-        self.capacity = capacity
-        self.batch_size = batch_size
-        self.n_step = n_step
-
-        self.train_after = train_after
-
-        self.er_buffer = deque(maxlen=self.capacity)
-
-        self.collector = collector(env, 100000)
-
-        if self.train_after:
-            self._warm_start(self.train_after)
-            print('\n### Warm up finished ###\n')
-            
-        self.policy = self._set_up_policy(policy)
-
-
-    def _warm_start(
-            self,
-            steps
-        ):
-
-        batch = self.collector.rollout(None,  pol_REGISTRY['random'](self.env), steps, self.n_step)
-
-        if self.sampler:
-            for transition in batch:
-                self.er_buffer.append(transition)
-
-        pass
-
-    def _set_up_policy(self, policy):
-        if isinstance(policy, str):
-            return pol_REGISTRY[policy](self.env)
-
-        elif isinstance(policy, Basepolicy):
-            return policy
-    
-        return 
-
-    def get_batch(
-            self,
-            net,
-        ):
-        
-        self.net = net        
-        batch = self.collector.rollout(net, self.policy, self.length, self.n_step)
-
-        if self.sampler:
-            for transition in batch:
-                self.er_buffer.append(transition)
-
-            k = min(self.batch_size, len(self.er_buffer))
-            batch = random.sample(list(self.er_buffer), k)
-
-        return self.prep_batch(batch)
-
-    def prep_batch(
-        self,
-        batch
-        ):
-        """
-        takes the batch (which will be a list of transitons) and processes them to be seperate etc.
-        """
-
-        if self.n_step == 1:
-            return TorchTransition(*zip(*batch))    # try return it as dict maybe? or just something fancier and seperated
-        
-        else:
-            processed_batch = []
-            # process the n_step transition
-            for n_step_transition in batch:
-                # print(f"{i}: ", n_step_transition, '\n')
-
-                transition = TorchTransition(*zip(*n_step_transition))
-                s = np.array(transition.s[0], dtype=np.float32)
-                a = transition.a[0]
-                r_list = list(transition.r)
-                s_p = np.array(transition.s_p[-1], dtype=np.float32)
-                d = any(transition.d)
-
-                processed_batch.append([s, a, r_list, s_p, d])
-
-            return TorchTransition(*zip(*processed_batch))
\ No newline at end of file
diff --git a/src/transitions.py b/src/transitions.py
deleted file mode 100644
index aa2641f..0000000
--- a/src/transitions.py
+++ /dev/null
@@ -1,21 +0,0 @@
-from typing import NamedTuple
-import torch as th
-import numpy as np
-
-class Transition(NamedTuple):
-    s: list  # state
-    a: float  # action
-    r: float  # reward
-    s_p: list  # next state
-    d: int  # done
-
-class TorchTransition(Transition):
-
-    def __call__(self):
-
-        s = th.from_numpy(np.array(self.s)).float()
-        a = th.from_numpy(np.array(self.a)).float().unsqueeze(1)
-        r = th.FloatTensor(self.r).unsqueeze(1)
-        s_p = th.from_numpy(np.array(self.s_p)).float()
-        d = th.IntTensor(self.d).unsqueeze(1)
-        return s, a, r, s_p, d
\ No newline at end of file
diff --git a/stub_awr.py b/stub_awr.py
new file mode 100644
index 0000000..59b31b5
--- /dev/null
+++ b/stub_awr.py
@@ -0,0 +1,122 @@
+import copy
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from cardio_rl import Runner
+from cardio_rl import Collector
+from cardio_rl.policies import EpsilonArgmax
+
+class Critic(nn.Module):
+	def __init__(self):
+		super(Critic, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(4, 64),
+			nn.ReLU(),
+			nn.Linear(64, 64),
+			nn.ReLU(),
+			nn.Linear(64, 2))
+
+	def forward(self, state):
+		return self.net(state)
+
+class Actor(nn.Module):
+	def __init__(self):
+		super(Actor, self).__init__()		
+		self.net = nn.Sequential(
+			nn.Linear(4, 128),
+			nn.ReLU(),
+			nn.Linear(128, 64),
+			nn.ReLU(),
+			nn.Linear(64, 2),
+			nn.Softmax(dim=-1)
+		)
+
+	def forward(self, state):
+		return self.net(state)
+
+
+def main():
+	env = gym.make('CartPole-v1')
+
+	runner = Runner(
+		env=env,
+		policy='categorical',
+		sampler=True,
+		capacity=100_000,
+		batch_size=256,
+		collector=Collector(
+			env=env,
+			rollout_len=16,
+			warmup_len=10_000,
+		),
+		backend='pytorch'
+	)
+
+	actor = Actor()
+	targ_actor: nn.Module = copy.deepcopy(actor)
+
+	critic = Critic()
+	targ_critic: nn.Module = copy.deepcopy(critic)
+
+	# use the default a2c optimiser
+	a_optimizer = th.optim.RMSProp(actor.parameters(), lr=5e-4)
+	c_optimizer = th.optim.RMSProp(critic.parameters(), lr=1e-3)
+	
+	gamma = 0.99
+	target_update = 10
+
+	for t in range(10000):
+		batch = runner.get_batch(actor)
+		s, a, r, s_p, d, _ = batch()
+
+		for _ in range(20):
+			with th.no_grad():
+				q_p = th.max(targ_critic(s_p), keepdim=True, dim=-1).values
+				y = r + gamma * q_p * (1 - d)
+
+			q = critic(s).gather(1, a.long())
+
+			c_loss = F.mse_loss(q, y)
+
+			c_optimizer.zero_grad()
+			c_loss.backward()
+			# nn.utils.clip_grad_norm_(critic.parameters(), 5.0)
+			c_optimizer.step()
+
+
+		for _ in range(100):
+			with th.no_grad():
+				q_adv = critic(s).gather(1, a.long())
+
+				# a_sampled = th.distributions.Categorical(actor(s)).sample()				
+				# v_adv = critic(s).gather(1, a_sampled.unsqueeze(-1).long())
+
+				# below is what the q-net suggests, not what the policy net suggests 
+				# should fix
+				v_adv = th.max(critic(s), keepdim=True, dim=-1).values
+				adv = q_adv - v_adv
+
+				weight = adv.exp()
+				weight = th.clamp_max(weight, 20)
+			
+			probs = actor(s)
+			dist = th.distributions.Categorical(probs)
+			log_probs = dist.log_prob(a.squeeze(-1))
+			policy_loss = -th.mean(log_probs * weight.squeeze(-1))
+
+			a_optimizer.zero_grad()
+			policy_loss.backward()
+			nn.utils.clip_grad_norm_(actor.parameters(), 5.0)
+			a_optimizer.step()
+
+
+		for targ_params, params in zip(targ_critic.parameters(), critic.parameters()):
+			targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
+
+		for targ_params, params in zip(targ_actor.parameters(), actor.parameters()):
+			targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
+			
+if __name__ == '__main__':
+	main()
\ No newline at end of file
diff --git a/stub_ddpg.py b/stub_ddpg.py
new file mode 100644
index 0000000..eda27bc
--- /dev/null
+++ b/stub_ddpg.py
@@ -0,0 +1,99 @@
+import copy
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from gymnasium.wrappers import RescaleAction
+from cardio_rl import Runner, Collector
+
+env_name = 'Pendulum-v1'
+env = gym.make(env_name)
+env = RescaleAction(env, -1.0, 1.0)
+
+runner = Runner(
+	env=env,
+	policy='whitenoise',
+	capacity=1_000_000,
+	collector=Collector(
+		warmup_len=100,
+		logger_kwargs=dict(
+			log_interval = 1_000
+			# log_dir = env_name,
+		)
+	),
+	backend='pytorch'
+)
+
+class Q_critic(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Q_critic, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(state_dim+action_dim, 256),
+			nn.ReLU(),
+			nn.Linear(256, 256),
+			nn.ReLU(),
+			nn.Linear(256, 1))
+
+	def forward(self, state, action):
+		sa = th.concat([state, action], dim=-1)
+		return self.net(sa)
+
+class Policy(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Policy, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(state_dim, 256),
+			nn.ReLU(),
+			nn.Linear(256, 256),
+			nn.ReLU(),
+			nn.Linear(256, action_dim),
+			nn.Tanh())
+		
+	def forward(self, state):
+		return self.net(state)
+
+# https://stable-baselines3.readthedocs.io/en/master/modules/ddpg.html
+critic = Q_critic(3, 1)
+actor = Policy(3, 1)
+targ_critic = copy.deepcopy(critic)
+targ_actor = copy.deepcopy(actor)
+c_optimizer = th.optim.Adam(critic.parameters(), lr=1e-3)
+a_optimizer = th.optim.Adam(actor.parameters(), lr=1e-3)
+
+for steps in range(300_000):
+	s, a, r, s_p, d, _ = runner.get_batch(actor)
+
+	with th.no_grad():
+		a_p = targ_actor(s_p)
+
+		mean = th.zeros_like(a_p)
+		noise = th.normal(mean=mean, std=0.2).clamp(-0.5, 0.5)
+		
+		a_p = (a_p + noise).clamp(-1, 1)
+
+		q_p = targ_critic(s_p, a_p)
+		y = r + 0.99 * q_p * (1 - d)
+
+	q = critic(s, a.squeeze(1))
+		
+	loss = F.mse_loss(q, y)
+
+	c_optimizer.zero_grad()
+	loss.backward()
+	th.nn.utils.clip_grad_norm_(critic.parameters(), 0.5)
+	c_optimizer.step()
+
+	policy_loss =  -(critic(s, actor(s_p))).mean()
+
+	a_optimizer.zero_grad()
+	policy_loss.backward()
+	th.nn.utils.clip_grad_norm_(actor.parameters(), 0.5)
+	a_optimizer.step()	
+
+	for targ_params, params in zip(targ_critic.parameters(), critic.parameters()):
+		targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
+
+	for targ_params, params in zip(targ_actor.parameters(), actor.parameters()):
+		targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
\ No newline at end of file
diff --git a/stub_drqn.py b/stub_drqn.py
new file mode 100644
index 0000000..4dd01f9
--- /dev/null
+++ b/stub_drqn.py
@@ -0,0 +1,83 @@
+import copy
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from cardio_rl import Runner
+from cardio_rl import Collector
+from cardio_rl.policies import EpsilonArgmax
+
+# to impelment:
+# -hidden state burn in
+# -sequence overlap
+# R2D2 paper for help: https://openreview.net/pdf?id=r1lyTjAqYX
+
+class RecurrentQ(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(RecurrentQ, self).__init__()
+		self.gru = nn.GRU(state_dim, 64, batch_first=True)
+		self.h1 = nn.Linear(64, 64)
+		self.h2 = nn.Linear(64, action_dim)
+
+	def forward(self, state, hidden, only_q=False):
+		if len(state.shape) < 2:
+			state = state.unsqueeze(0)
+		hidden, _ = self.gru(state, hidden)
+		x = F.relu(hidden)
+		x = F.relu(self.h1(x))
+		q = self.h2(x)
+		if only_q:
+			return q
+		return q, hidden
+	
+def main():
+	env = gym.make('CartPole-v1')
+
+	runner = Runner(
+		env=env,
+		policy=EpsilonArgmax(env, 0.5, 0.05, 0.9, True, 64),
+		sampler=True,
+		capacity=100_000,
+		batch_size=256,
+		collector=Collector(
+			env=env,
+			rollout_len=4,
+			warmup_len=10_000,
+			n_step=16,
+			take_every=8,
+		),
+		reduce=False,
+		backend='pytorch'
+	)
+
+	critic = RecurrentQ(4, 2)
+	targ_net: nn.Module = copy.deepcopy(critic)
+	optimizer = th.optim.Adam(critic.parameters(), lr=1.0e-3)
+	gamma = 0.99
+	target_update = 10
+
+	for t in range(10000):
+		batch = runner.get_batch(critic)
+		s, a, r, s_p, d, _ = batch()
+
+		# You are pretty close, just need to cross reference with this:
+		# https://github.com/mynkpl1998/Recurrent-Deep-Q-Learning/blob/master/LSTM%2C%20BPTT%3D8.ipynb
+		# also, implement the 'warm-start' use
+
+		with th.no_grad():
+			q_p = th.max(targ_net(s_p, th.zeros(1, runner.batch_size, 64), only_q=True), keepdim=True, dim=-1).values.transpose(1,2)
+			y = r + gamma * q_p * (1 - d)
+
+		q = critic(s, th.zeros(1, runner.batch_size, 64), only_q=True).gather(2, a.long())
+
+		loss = F.mse_loss(q, y)
+
+		optimizer.zero_grad()
+		loss.backward()
+		optimizer.step()
+
+		if t % target_update == 0:        
+			targ_net = copy.deepcopy(critic)
+			
+if __name__ == '__main__':
+	main()
\ No newline at end of file
diff --git a/stub_qrdqn.py b/stub_qrdqn.py
new file mode 100644
index 0000000..f4978c1
--- /dev/null
+++ b/stub_qrdqn.py
@@ -0,0 +1,85 @@
+import copy
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from cardio_rl import Runner
+from cardio_rl import Collector
+from cardio_rl.policies import EpsilonArgmax
+
+
+class QuantileregressionQ(nn.Module):
+	def __init__(self, state_dim, action_dim, quantiles=32):
+		super(QuantileregressionQ, self).__init__()
+		self.action_dim = action_dim
+		self.quantiles = quantiles
+
+		self.net = nn.Sequential(
+			nn.Linear(state_dim, 128),
+			nn.ReLU(),
+			nn.Linear(128, 128),
+			nn.ReLU(),
+			nn.Linear(128, action_dim*quantiles)
+		)
+
+	def forward(self, state, reduce=True):
+		q = self.net(state).view(-1, self.action_dim, self.quantiles)
+
+		if reduce:
+			q = th.mean(q, dim=-1)
+
+		return q
+	
+def main():
+	env = gym.make('CartPole-v1')
+
+	runner = Runner(
+		env=env,
+		policy=EpsilonArgmax(env, 0.9, 0.05, 0.9),
+		sampler=True,
+		capacity=100_000,
+		batch_size=256,
+		collector=Collector(
+			env=env,
+			rollout_len=4,
+			warmup_len=500,
+		),
+		backend='pytorch'
+	)
+
+	critic = QuantileregressionQ(4, 2, 32)
+	targ_net: nn.Module = copy.deepcopy(critic)
+	optimizer = th.optim.Adam(critic.parameters(), lr=2.3e-3)
+	gamma = 0.99
+	target_update = 300
+
+	tau = th.FloatTensor([(n/32) for n in range(1, 32+1)])
+
+	for t in range(10000):
+		batch = runner.get_batch(critic)
+		s, a, r, s_p, d, _ = batch()
+
+		with th.no_grad():
+			a_p = critic(s_p).argmax(dim=-1).view(-1, 1, 1)
+			q_p = targ_net(s_p, False).gather(1, a_p.expand(256, 1, 32).long())
+			y = r.view(-1, 1, 1) + gamma * q_p * (1 - d.view(-1, 1, 1))
+
+		q = critic(s, False).gather(1, a.unsqueeze(-1).expand(256, 1, 32).long())
+		
+		kappa = 1.0
+		u = y - q
+
+		huber_loss = th.where(u.abs() <= kappa, 0.5 * u.pow(2), kappa * (u.abs() - 0.5 * kappa))
+		qh_loss = abs(tau - (u.detach() < 0).float()) * huber_loss
+		loss = qh_loss.sum(-1).mean()
+
+		optimizer.zero_grad()
+		loss.backward()
+		th.nn.utils.clip_grad_norm_(critic.parameters(), 10.0)
+		optimizer.step()
+
+		if t % target_update == 0:        
+			targ_net = copy.deepcopy(critic)
+			
+if __name__ == '__main__':
+	main()
\ No newline at end of file
diff --git a/stub_reinforce.py b/stub_reinforce.py
new file mode 100644
index 0000000..9c492d9
--- /dev/null
+++ b/stub_reinforce.py
@@ -0,0 +1,57 @@
+from cardio_rl import Runner, Collector
+import torch as th
+import gymnasium as gym
+import torch.nn as nn
+from torch.distributions import Categorical
+
+class Actor(nn.Module):
+	def __init__(self):
+		super(Actor, self).__init__()
+		
+		self.net = nn.Sequential(
+			nn.Linear(4, 128),
+			nn.ReLU(),
+			nn.Linear(128, 128),
+			nn.ReLU(),
+			nn.Linear(128, 2),
+			nn.Softmax(dim=-1)
+		)
+
+	def forward(self, state):
+		return self.net(state)
+
+env = gym.make('CartPole-v1')
+
+runner = Runner(
+	env=env,
+	policy='categorical',
+	collector=Collector(
+		env=env,
+		rollout_len=-1,
+	),
+	backend='pytorch'
+)
+
+actor = Actor()
+optimizer = th.optim.Adam(actor.parameters(), lr=2e-3)
+
+for _ in range(1_000):
+	batch = runner.get_batch(actor)
+	s, a, r, s_p, d, _ = batch()
+
+	returns = th.zeros_like(r)
+	running_r = 0
+	for t in reversed(range(len(r))):
+		running_r = r[t] + 0.99 * running_r
+		returns[t] = running_r
+
+	returns = (returns - returns.mean()) / returns.std()
+
+	probs = actor(s)
+	dist = Categorical(probs)
+	policy_loss = -th.mean(dist.log_prob(a.squeeze(-1)) * returns.squeeze(-1))
+
+	optimizer.zero_grad()
+	policy_loss.backward()
+	nn.utils.clip_grad_norm_(actor.parameters(), 0.5)
+	optimizer.step()	
diff --git a/stub_rwr.py b/stub_rwr.py
new file mode 100644
index 0000000..555094d
--- /dev/null
+++ b/stub_rwr.py
@@ -0,0 +1,99 @@
+from cardio_rl import Runner, Collector
+import torch as th
+import gymnasium as gym
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Categorical
+import copy
+import numpy as np
+
+class Critic(nn.Module):
+	def __init__(self):
+		super(Critic, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(4, 64),
+			nn.ReLU(),
+			nn.Linear(64, 64),
+			nn.ReLU(),
+			nn.Linear(64, 1))
+
+	def forward(self, state):
+		return self.net(state)
+
+class Actor(nn.Module):
+	def __init__(self):
+		super(Actor, self).__init__()
+		
+		self.net = nn.Sequential(
+			nn.Linear(4, 128),
+			nn.ReLU(),
+			nn.Linear(128, 128),
+			nn.ReLU(),
+			nn.Linear(128, 2),
+			nn.Softmax(dim=-1)
+		)
+
+	def forward(self, state):
+		return self.net(state)
+
+env = gym.make('CartPole-v1')
+
+runner = Runner(
+	env=env,
+	policy='categorical',
+	collector=Collector(
+		env=env,
+		rollout_len=-1,
+	),
+	backend='pytorch'
+)
+
+actor = Actor()
+targ_actor = copy.deepcopy(actor)
+critic = Critic()
+a_optimizer = th.optim.Adam(actor.parameters(), lr=1e-3)
+
+log_alpha = nn.Parameter(th.zeros(1))
+alpha_optim = th.optim.Adam([log_alpha], lr=1e-3)
+alpha = log_alpha.exp().detach().item()
+
+for i in range(1_000):
+	batch = runner.get_batch(actor)
+	s, a, r, s_p, d, _ = batch()
+
+	returns = th.zeros_like(r)
+	running_r = 0
+	for t in reversed(range(len(r))):
+		running_r = r[t] + 0.99 * running_r
+		returns[t] = running_r
+
+	adv: th.Tensor = (returns - 40).detach()
+
+	weights = th.clamp_max(adv.exp(), 20)
+
+	for _ in range(20):
+		probs = actor(s)
+		dist = Categorical(probs)
+
+		targ_probs = targ_actor(s).detach()
+
+		kl = th.mean((probs * th.log(probs / targ_probs)).sum(dim=-1))
+
+		alpha_loss = -log_alpha.exp() * (0.01 - kl).detach().item()
+		alpha_optim.zero_grad()
+		alpha_loss.backward()
+		alpha_optim.step()
+
+		alpha = log_alpha.exp().detach().item()
+
+		policy_loss = -th.mean(dist.log_prob(a.squeeze(-1)) * (weights)) + (alpha * (0.01 - kl))
+
+		a_optimizer.zero_grad()
+		policy_loss.backward()
+		nn.utils.clip_grad_norm_(actor.parameters(), 0.1)
+		a_optimizer.step()	
+
+	targ_actor = copy.deepcopy(actor)
+
+	
diff --git a/stub_td3_rnd.py b/stub_td3_rnd.py
new file mode 100644
index 0000000..9a5b078
--- /dev/null
+++ b/stub_td3_rnd.py
@@ -0,0 +1,203 @@
+import copy
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+import numpy as np
+from gymnasium.wrappers import RescaleAction
+from cardio_rl.policies import BasePolicy
+from cardio_rl import Runner, Collector
+
+class RndWhitenoiseDeterministic(BasePolicy):
+	def __init__(self, env: gym.Env, obs_dims, output_dims):
+		super().__init__(env)
+		# architecture can maybe be improved on
+		self.rnd_net = nn.Sequential(
+			nn.Linear(obs_dims, 16),
+			nn.Tanh(),
+			nn.Linear(16, 16),
+			nn.Tanh(),
+			nn.Linear(16, output_dims),
+		)
+
+		self.targ_net = nn.Sequential(
+			nn.Linear(obs_dims, 16),
+			nn.Tanh(),
+			nn.Linear(16, 16),
+			nn.Tanh(),
+			nn.Linear(16, output_dims),
+		)
+		
+	def __call__(self, state, net):
+		input = th.from_numpy(state).float()  
+		out = net(input)         
+		mean = th.zeros_like(out)
+		noise = th.normal(mean=mean, std=0.1).clamp(-0.5, 0.5)
+		out = (out + noise)    
+		return out.clamp(-1, 1).detach().numpy()
+	
+	def predictor(self, next_state):
+		next_state = th.from_numpy(next_state).float()
+		pred = self.rnd_net(next_state)
+		targ = self.targ_net(next_state)
+		intrinsic_r = th.pow(th.mean(pred-targ, dim=-1, keepdim=True), 2).detach().item()
+		return intrinsic_r
+
+	def update(self, rnd_net):
+		self.rnd_net = rnd_net
+
+class RndCollector(Collector):
+	def _env_step(self, policy=None, warmup=False):
+		if warmup:
+			a = self.env.action_space.sample()
+		else:
+			a = policy(self.state, self.net)
+		s_p, r, d, t, info = self.env.step(a)
+		self.logger.step(r, d, t)
+		d = d or t
+
+		r_i = 1.0 * policy.predictor(s_p)
+		r_e = r
+
+		return (self.state, a, [r_e, r_i], s_p, d, info), s_p, d, t
+
+env_name = 'MountainCarContinuous-v0'
+# env_name = 'Pendulum-v1'
+# env_name = 'LunarLanderContinuous-v2'
+env = gym.make(env_name)
+env = RescaleAction(env, -1.0, 1.0)
+
+beta_rnd = 0.1
+
+runner = Runner(
+	env=env,
+	policy=RndWhitenoiseDeterministic(env, obs_dims=2, output_dims=16),
+	sampler=True,
+	capacity=1_000_000,
+	batch_size=512,
+	n_batches=32,
+	collector=RndCollector(
+		env=env,
+		rollout_len=32,
+		warmup_len=0,
+		logger_kwargs=dict(
+			log_interval=5_000,
+			episode_window=20,
+			tensorboard=True,
+            log_dir='run',
+			exp_name = env_name + f'_RND_separate_nets_beta{beta_rnd}'
+		)
+	),
+	backend='pytorch'
+)
+
+class Q_critic(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Q_critic, self).__init__()
+
+		self.net1 = nn.Sequential(
+			nn.Linear(state_dim+action_dim, 64),
+			nn.Dropout(0.01),
+			nn.LayerNorm(64),
+			nn.ReLU(),
+			nn.Linear(64, 64),
+			nn.Dropout(0.01),
+			nn.LayerNorm(64),
+			nn.ReLU(),
+			nn.Linear(64, 1))
+		
+		self.net2 = nn.Sequential(
+			nn.Linear(state_dim+action_dim, 64),
+			nn.Dropout(0.01),
+			nn.LayerNorm(64),
+			nn.ReLU(),
+			nn.Linear(64, 64),
+			nn.Dropout(0.01),
+			nn.LayerNorm(64),
+			nn.ReLU(),
+			nn.Linear(64, 1))
+
+	def forward(self, state, action):
+		sa = th.concat([state, action], dim=-1)
+		return self.net1(sa), self.net2(sa)
+
+class Policy(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Policy, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(state_dim, 64),
+			nn.ReLU(),
+			nn.Linear(64, 64),
+			nn.ReLU(),
+			nn.Linear(64, action_dim),
+			nn.Tanh())
+		
+	def forward(self, state):
+		return self.net(state)
+
+# https://stable-baselines3.readthedocs.io/en/master/modules/ddpg.html
+critic = Q_critic(2, 1)
+int_critic = Q_critic(2, 1)
+actor = Policy(2, 1)
+targ_critic = copy.deepcopy(critic)
+int_targ_critic = copy.deepcopy(int_critic)
+targ_actor = copy.deepcopy(actor)
+c_optimizer = th.optim.Adam(list(critic.parameters()) + list(int_critic.parameters()), lr=1e-3)
+a_optimizer = th.optim.Adam(actor.parameters(), lr=1e-3)
+rnd_optimizer = th.optim.Adam(runner.policy.rnd_net.parameters(), lr=1e-3)
+
+tau = 0.01
+
+for steps in range(120_000):
+	all_batches = runner.get_batch(actor)
+	for batch in all_batches:
+		s, a, r, s_p, d, _ = batch()
+		
+		r_e, r_i = r[:, :, 0], r[:, :, 1]
+
+		with th.no_grad():
+			a_p = targ_actor(s_p)
+
+			noise = th.normal(mean=th.zeros_like(a_p), std=0.2).clamp(-0.5, 0.5)		
+			a_p = (a_p + noise).clamp(-1, 1)
+
+			q_p = th.min(*targ_critic(s_p, a_p))
+			y = r_e + 0.9999 * q_p * (1 - d)
+
+			i_q_p = th.min(*int_targ_critic(s_p, a_p))
+			i_y = r_i + 0.99 * q_p * (1 - d)
+
+		q1, q2 = critic(s, a.squeeze(1))
+		i_q1, i_q2 = int_critic(s, a.squeeze(1))
+			
+		critic_loss = (F.mse_loss(q1, y) + F.mse_loss(q2, y)) + (F.mse_loss(i_q1, i_y) + F.mse_loss(i_q2, i_y))
+
+		c_optimizer.zero_grad()
+		critic_loss.backward()
+		c_optimizer.step()
+
+		pred = runner.policy.rnd_net(s_p)
+		targ = runner.policy.targ_net(s_p).detach()
+
+		rnd_loss = F.mse_loss(pred, targ)
+
+		rnd_optimizer.zero_grad()
+		rnd_loss.backward()
+		rnd_optimizer.step()
+
+		if steps % 2 == 0:
+			policy_loss = -(th.min(*critic(s_p, actor(s_p))) + beta_rnd * th.min(*int_critic(s_p, actor(s_p)))).mean()
+
+			a_optimizer.zero_grad()
+			policy_loss.backward()
+			a_optimizer.step()	
+
+			for targ_params, params in zip(targ_critic.parameters(), critic.parameters()):
+				targ_params.data.copy_(params.data * tau + targ_params.data * (1.0 - tau))
+
+			for targ_params, params in zip(int_targ_critic.parameters(), int_critic.parameters()):
+				targ_params.data.copy_(params.data * tau + targ_params.data * (1.0 - tau))
+
+			for targ_params, params in zip(targ_actor.parameters(), actor.parameters()):
+				targ_params.data.copy_(params.data * tau + targ_params.data * (1.0 - tau))


Test patch:
diff --git a/src/__init__.py b/cardio_rl/offline_rl/testing_bcq.py
similarity index 100%
rename from src/__init__.py
rename to cardio_rl/offline_rl/testing_bcq.py
diff --git a/cardio_rl/offline_rl/testing_cql.py b/cardio_rl/offline_rl/testing_cql.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/testing_naf.py b/cardio_rl/testing_naf.py
similarity index 100%
rename from src/testing_naf.py
rename to cardio_rl/testing_naf.py
diff --git a/jax_testing.py b/jax_testing.py
new file mode 100644
index 0000000..352448f
--- /dev/null
+++ b/jax_testing.py
@@ -0,0 +1,136 @@
+import copy
+import jax
+import jax.numpy as jnp
+import optax
+import flax.linen as nn
+import distrax
+import gymnasium as gym
+
+from cardio_rl import Runner, Collector
+
+env = gym.make('Pendulum-v1')
+
+runner = Runner(
+	env=env,
+	policy='random',	# need to implement jax policies
+	sampler=True,
+	capacity=100000,
+	batch_size=256,
+	collector=Collector(
+		env=env,
+		warmup_len=100,
+	),
+	backend='jax'
+)
+
+class Actor(nn.Module):
+	@nn.compact
+	def __call__(self, state):
+		x = nn.Dense(128)(state)
+		x = nn.relu(x)
+
+		# Deterministic policy for DDPG
+		x = nn.Dense(128)(x)
+		x = nn.relu(x)
+		x = nn.Dense(1)(x)
+		x = nn.relu(x)
+		return x
+
+		# Beta policy for SAC
+		"""alpha = nn.Dense(128)(x)
+		alpha = nn.relu(alpha)
+		alpha = nn.Dense(1)(alpha)
+		alpha = nn.softplus(alpha)
+
+		beta = nn.Dense(128)(x)
+		beta = nn.relu(beta)
+		beta = nn.Dense(1)(beta)
+		beta = nn.softplus(beta)
+		return alpha, beta"""
+
+class Critic(nn.Module):
+	@nn.compact
+	def __call__(self, state, action):
+		sa = jnp.concatenate([state, action], axis=-1)
+		x = nn.Dense(128)(sa)
+		x = nn.relu(x)
+		x = nn.Dense(128)(x)
+		x = nn.relu(x)
+		x = nn.Dense(1)(x)
+		return x
+
+
+actor = Actor()
+a_batch = jnp.ones(3)
+a_params = actor.init(jax.random.PRNGKey(0), a_batch)
+a_optimizer = optax.adam(learning_rate=1e-3)
+a_opt_state = a_optimizer.init(a_params)
+output = actor.apply(a_params, a_batch)
+
+critic = Critic()
+c_batch = (jnp.ones(3), jnp.ones(1))
+c_params = critic.init(jax.random.PRNGKey(0), *c_batch)
+targ_c_params = copy.deepcopy(c_params)
+c_optimizer = optax.adam(learning_rate=1e-3)
+c_opt_state = a_optimizer.init(c_params)
+output = critic.apply(c_params, *c_batch)
+
+ent_coeff = 0.1
+
+for steps in range(20_000):
+	if steps % 100 == 0 and steps > 0:
+		print(steps, runner.policy)
+
+	batch = runner.get_batch(actor)
+
+	# first step is implementing Transition/backend selection
+	s, a, r, s_p, d = batch()
+
+	a_p = actor.apply(a_params, s_p)
+	q_p = critic.apply(targ_c_params, s_p, a_p)
+
+	"""
+	alpha, beta = actor.apply(params, x)
+	dist = distrax.Beta(alpha, beta)
+	a_p = (dist.sample(jax.random.PRNGKey(0)) * 4) - 2
+	return -jnp.mean(critic.apply(q_params, x, a_p))
+	q_p = critic.apply(targ_c_params, s_p, a_p) - ent_coeff*dist.entropy()
+	y = jax.lax.stop_gradient(r + 0.99 *  q_p * (1 - d))	
+	"""
+
+	y = jax.lax.stop_gradient(r + 0.99 * q_p * (1 - d))	
+
+	@jax.jit
+	@jax.value_and_grad
+	def td_error(params, y):
+		q = critic.apply(params, s, a)
+		return jnp.mean(optax.l2_loss(q, y))	
+	
+	c_loss, c_grads = td_error(c_params, y)
+	updates, c_opt_state = c_optimizer.update(c_grads, c_opt_state)
+	c_params = optax.apply_updates(c_params, updates)
+
+	if steps % 2 == 0:
+		@jax.jit
+		@jax.value_and_grad
+		def policy_loss(params, q_params, x):
+			a_sampled = actor.apply(params, x)
+			return -jnp.mean(critic.apply(q_params, x, a_sampled))
+
+		a_loss, a_grads = policy_loss(a_params, c_params, s)
+		updates, a_opt_state, = a_optimizer.update(a_grads, a_opt_state)
+		a_params = optax.apply_updates(a_params, updates)
+
+		continue
+		@jax.jit
+		@jax.value_and_grad
+		def policy_loss(params, q_params, x):
+			alpha, beta = actor.apply(params, x)
+			dist = distrax.Beta(alpha, beta)
+			a_sampled = (dist.sample(jax.random.PRNGKey(0)) * 4) - 2
+			return -jnp.mean(critic.apply(q_params, x, a_sampled) - (ent_coeff*dist.entropy()))
+		
+
+	# see if you can do this more smoothly
+	targ_c_params = jax.tree_util.tree_map(lambda x, y: (1 - 0.005) * x + 0.005 * y, targ_c_params, c_params)
+	
\ No newline at end of file
diff --git a/minari_testing.py b/minari_testing.py
new file mode 100644
index 0000000..31e7faa
--- /dev/null
+++ b/minari_testing.py
@@ -0,0 +1,95 @@
+from collections import deque
+from cardio_rl import Collector, Runner
+from cardio_rl.policies import BasePolicy
+from minari import DataCollectorV0
+import minari
+import gymnasium as gym
+import d4rl
+
+class OfflineCollector(Collector):
+    def __init__(self, env, capacity) -> None:
+        super().__init__(env, capacity)
+        self.env = env
+        self.warmup()
+
+    def warmup(self, net=None, policy=None, length=0, n_step=1):
+
+        if isinstance(self.env, gym.Env):
+            return self._warmup_gym(net, policy, length, n_step)
+        
+        self._warumup_d4rl()
+
+        return 
+    
+    def _warumup_d4rl(self):
+        dataset = d4rl.qlearning_dataset(self.env)
+        print(dataset)
+
+    def _warmup_gym(self, net=None, policy=None, length=1000000, n_step=1):
+
+        if policy == None:
+            policy = BasePolicy(self.env)
+
+        self.env = DataCollectorV0(
+            self.env,
+            record_infos=False,
+            max_buffer_steps=length)
+        
+        self.net = net
+        gather_buffer = deque()
+        step_buffer = deque(maxlen=n_step)
+
+        if length < 1:
+            print('cant have n-step less than 1 when using offline gatherer')      
+
+        for _ in range(length):
+            self.total_steps += 1
+            a = policy(self.state, self.net)
+            s_p, r, d, t, info = self.env.step(a)
+
+            step_buffer.append([self.state, a, r, s_p, d])
+            if len(step_buffer) == n_step:
+
+                if n_step == 1:
+                    gather_buffer.append(*list(step_buffer))
+                else:
+                    gather_buffer.append(list(step_buffer))
+
+            self.state = s_p
+            if d or t:
+                self.episodes += 1
+                self.epsiode_window.append(self.ep_rew)
+                self.state, _ = self.env.reset()
+                step_buffer = deque(maxlen=n_step)
+
+        self._set_up_loader(data=list(gather_buffer))
+
+        return list(gather_buffer)
+    
+    def _set_up_loader(self, data):
+        print(len(data))
+        dataset = minari.create_dataset_from_collector_env(
+            dataset_id="CartPole-v1-test-v0", 
+            collector_env=self.env,
+            author='Manus',
+            author_email='mmcaulif@tcd.ie')
+
+def get_offline_runner(env, capacity):
+    return Runner(
+        env,
+        BasePolicy(env),
+        length=1,
+        sampler=False,
+        capacity=0,
+        batch_size=256,
+        n_step=1,
+        collector=OfflineCollector,
+        train_after=capacity,
+    )
+
+def main():
+    env = gym.make('maze2d-umaze-v1')
+    runner = get_offline_runner(env, 200)
+
+print('here')
+main()
diff --git a/mpo_testing.py b/mpo_testing.py
new file mode 100644
index 0000000..2a1202b
--- /dev/null
+++ b/mpo_testing.py
@@ -0,0 +1,161 @@
+import copy
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from cardio_rl import Runner
+from cardio_rl import Collector
+
+"""
+To do:
+-figure out how many e-pochs to do for E and M step
+-look into what hyperparams to use, maybe good time to try optuna
+-look into Acme impl. 
+Resources:
+-https://github.com/ethanluoyc/magi/pull/59
+-https://github.com/vwxyzjn/cleanrl/pull/392
+-https://github.com/daisatojp/mpo/blob/master/mpo/mpo.py
+-https://github.com/deepmind/acme/tree/master/acme/agents/jax/mpo
+"""
+
+class Critic(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Critic, self).__init__()
+		self.net = nn.Sequential(
+			nn.Linear(state_dim, 256),
+			nn.ReLU(),
+			nn.Linear(256, 256),
+			nn.ReLU(),
+			nn.Linear(256, action_dim))
+
+	def forward(self, state):
+		return self.net(state)
+
+class Actor(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Actor, self).__init__()
+		self.net = nn.Sequential(
+			nn.Linear(state_dim, 256),
+			nn.ReLU(),
+			nn.Linear(256, 256),
+			nn.ReLU(),
+			nn.Linear(256, action_dim),
+			nn.Softmax(dim=-1))
+
+	def forward(self, state):
+		return self.net(state)
+
+env = gym.make('LunarLander-v2')
+
+runner = Runner(
+	env=env,
+	policy='categorical',
+	sampler=True,
+	capacity=50_000,	# 1_000_000,
+	batch_size=256,
+	collector=Collector(
+		env=env,
+		rollout_len=4,
+		warmup_len=1_000,
+		logger_kwargs=dict(
+			log_interval=5000,
+		)
+	),
+	backend='pytorch'
+)
+
+critic = Critic(8, 4)
+actor = Actor(8, 4)
+
+targ_critic = copy.deepcopy(critic)
+targ_actor = copy.deepcopy(actor)
+
+# taken from acme
+lr = 1e-3
+dual_lr = 1e-3
+
+critic_optim = th.optim.Adam(critic.parameters(), lr=lr)
+actor_optim = th.optim.Adam(actor.parameters(), lr=lr)
+gamma = 0.99
+
+eps = 0.1
+init_logs = 0.0
+def from_log(x): 
+	# projection suggested in paper and acme impl.
+	# return F.softplus(th.clamp(x, min=-18)) + 1e-8
+	return x.exp()
+
+# check init values of lagrangian multipliers, acme impl. uses clamping to restrict the min value
+# update: apparently acme uses 1.0 as init values? these are very unstable when testing...
+log_eta = th.tensor(init_logs, requires_grad=True)
+eta_optim = th.optim.Adam([log_eta], lr=dual_lr)
+eta = from_log(log_eta).item()
+
+log_alpha = th.tensor(init_logs, requires_grad=True)
+alpha_optim = th.optim.Adam([log_alpha], lr=dual_lr)
+alpha = from_log(log_alpha).item()
+
+for t in range(12500):
+	batch = runner.get_batch(actor)
+	s, a, r, s_p, d, _ = batch()
+
+	### Policy evaluation
+	with th.no_grad():
+		# ddqn
+		a_p = critic(s_p).argmax(dim=1).unsqueeze(1)		
+		q_p = targ_critic(s_p).gather(1, a_p.long())		
+		y = r + gamma * q_p * (1 - d)
+
+	q = critic(s).gather(1, a.long())
+
+	loss = F.mse_loss(q, y)
+
+	critic_optim.zero_grad()
+	loss.backward()
+	nn.utils.clip_grad_norm_(critic.parameters(), 0.5)
+	critic_optim.step()
+	
+	### E-step
+	for _ in range(5):
+		energies = th.exp(critic(s).detach() / F.softplus(log_eta)).mean(-1)
+		logits = th.log(energies).mean()
+		eta_loss = (from_log(log_eta) * (eps + logits))
+
+		eta_optim.zero_grad()
+		eta_loss.backward()		
+		nn.utils.clip_grad_norm_([log_eta], 40.0)	
+		eta_optim.step()
+
+	eta = from_log(log_eta).item()
+	
+	qij = F.softmax(targ_critic(s) / eta, dim=-1).detach()
+	prior_probs = targ_actor(s).detach()
+
+	### M-step
+	for _ in range(5):
+		probs = actor(s)
+		
+		# check kl div calculation
+		loss_p = th.mean(qij * th.log(probs))
+
+		# check dimensions of kl-div!
+		kl = th.mean((probs * th.log(probs / prior_probs)).sum(dim=-1))
+		total_loss_p = -(loss_p + alpha * (eps - kl))
+
+		actor_optim.zero_grad()
+		total_loss_p.backward()
+		nn.utils.clip_grad_norm_(actor.parameters(), 0.5)
+		actor_optim.step()
+
+		# is alpha updated inside or outside of M-steps?
+		alpha_loss = from_log(log_alpha) * (eps - kl).detach()
+		alpha_optim.zero_grad()
+		alpha_loss.backward()
+		nn.utils.clip_grad_norm_([log_alpha], 40.0)
+		alpha_optim.step()
+		alpha = from_log(log_alpha).item()
+
+	# implement soft/delayed updates maybe?
+	if t % 100 == 0:	# 100 is used in acme impl.
+		targ_critic = copy.deepcopy(critic)
+		targ_actor = copy.deepcopy(actor)
diff --git a/src/testing_a2c.py b/src/testing_a2c.py
deleted file mode 100644
index 4519ca2..0000000
--- a/src/testing_a2c.py
+++ /dev/null
@@ -1,83 +0,0 @@
-import torch as th
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.distributions import Categorical
-import gymnasium as gym
-from runner import get_onpolicy_runner
-
-env = gym.make('CartPole-v1')
-
-runner = get_onpolicy_runner(
-	env, 
-	'categorical',
-	length=5,)
-
-class Critic(nn.Module):
-	def __init__(self, state_dim):
-		super(Critic, self).__init__()
-
-		self.net = nn.Sequential(
-			nn.Linear(state_dim, 64),
-			nn.ReLU(),
-			nn.Linear(64, 64),
-			nn.ReLU(),
-			nn.Linear(64, 1))
-
-	def forward(self, state):
-		return self.net(state)
-
-class Policy(nn.Module):
-	def __init__(self, state_dim, action_dim):
-		super(Policy, self).__init__()
-
-		self.net = nn.Sequential(
-			nn.Linear(state_dim, 64),
-			nn.ReLU(),
-			nn.Linear(64, 64),
-			nn.ReLU(),
-			nn.Linear(64, action_dim),
-			nn.Softmax())
-
-	def forward(self, state):
-		return self.net(state)
-
-net = Critic(4)
-policy = Policy(4, 2)
-c_optimizer = th.optim.Adam(net.parameters(), lr=7e-4)
-pi_optimizer = th.optim.Adam(policy.parameters(), lr=7e-4)
-
-"""
-Need to implement parallel environments and better returns estimation!
-"""
-
-for rollout_steps in range(50000):
-	batch = runner.get_batch(policy)
-
-	s, a, r, s_p, d = batch()
-
-	# Need to improve returns estimates with future expected value
-	values = net(s)
-	returns = th.zeros_like(r)
-	ret = 0
-		
-	for i in reversed(range(len(r))):
-		ret += 0.99 * r[i]
-		ret *= (1 - d[i])
-		returns[i] = ret
-	
-	loss = F.mse_loss(returns.detach(), values)
-	
-	c_optimizer.zero_grad()
-	loss.backward()
-	th.nn.utils.clip_grad_norm_(net.parameters(), 0.5)
-	c_optimizer.step()
-	
-	adv = (values - returns).detach()
-	log_probs = Categorical(policy(s)).log_prob(a.squeeze(-1))
-
-	policy_loss = -(log_probs * adv).mean()
-
-	pi_optimizer.zero_grad()
-	policy_loss.backward()
-	th.nn.utils.clip_grad_norm_(policy.parameters(), 0.5)
-	pi_optimizer.step()
diff --git a/src/testing_sac.py b/src/testing_sac.py
deleted file mode 100644
index 596b01c..0000000
--- a/src/testing_sac.py
+++ /dev/null
@@ -1,160 +0,0 @@
-import copy
-import torch as th
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.distributions import Normal
-import numpy as np
-import gymnasium as gym
-from runner import get_offpolicy_runner
-
-env = gym.make('gym_cartpole_continuous:CartPoleContinuous-v1')
-
-runner = get_offpolicy_runner(
-	env, 
-	'gaussian',
-	freq=1,
-	capacity=1000000, 
-	batch_size=256, 
-	train_after=1000)
-
-class Q_critic(nn.Module):
-	def __init__(self, state_dim, action_dim):
-		super(Q_critic, self).__init__()
-
-		self.net1 = nn.Sequential(
-			nn.Linear(state_dim+action_dim, 400),
-			nn.ReLU(),
-			nn.Linear(400, 300),
-			nn.ReLU(),
-			nn.Linear(300, 1))
-		
-		self.net2 = nn.Sequential(
-			nn.Linear(state_dim+action_dim, 400),
-			nn.ReLU(),
-			nn.Linear(400, 300),
-			nn.ReLU(),
-			nn.Linear(300, 1))
-
-	def forward(self, state, action):
-		sa = th.concat([state, action], dim=-1)
-		q1, q2 = self.net1(sa), self.net2(sa)
-		return q1, q2
-
-class Policy(nn.Module):
-	def __init__(self, state_dim, action_dim):
-		super(Policy, self).__init__()
-
-		self.net = nn.Sequential(
-			nn.Linear(state_dim, 400),
-			nn.ReLU(),
-			nn.Linear(400, 300),
-			nn.ReLU(),)
-		
-		self.mean = nn.Sequential(
-			nn.Linear(300, action_dim)
-		)
-
-		self.std = nn.Sequential(
-			nn.Linear(300, action_dim),
-			nn.Tanh()
-		)
-		
-	def forward(self, state):
-		h = self.net(state)
-		mean = self.mean(h)
-		log_std = self.std(h)
-		log_std = -5 + 0.5 * (2 - -5) * (log_std + 1)  # From SpinUp / Denis Yarats
-		return mean, log_std
-
-net = Q_critic(4, 1)
-policy = Policy(4, 1)
-targ_net = copy.deepcopy(net)
-c_optimizer = th.optim.Adam(net.parameters(), lr=2e-3)
-pi_optimizer = th.optim.Adam(policy.parameters(), lr=2e-3)
-
-autotune = False
-
-if autotune:
-	target_entropy = -th.prod(th.Tensor(env.action_space.shape)).item()
-	log_alpha = th.zeros(1, requires_grad=True)
-	alpha = log_alpha.exp().item()
-	a_optimizer = th.optim.Adam([log_alpha], lr=2e-3)
-
-else:
-	alpha = 0.15
-
-# https://stable-baselines3.readthedocs.io/en/master/modules/sac.html
-
-for steps in range(50000):
-	batch = runner.get_batch(policy)
-
-	s, a, r, s_p, d = batch()
-
-	with th.no_grad():
-
-		mu, log_std = policy(s_p)
-		std = log_std.exp()
-
-		dist_p = Normal(mu, std)
-		x_t = dist_p.rsample()
-		y_t = nn.Tanh()(x_t).detach()
-		a_p = y_t * env.action_space.high + 0
-
-		q1_p, q2_p = targ_net(s_p, a_p)
-
-		log_prob = dist_p.log_prob(x_t)
-		log_prob -= th.log(th.from_numpy(env.action_space.high) * (1 - y_t.pow(2)) + 1e-6)
-		log_prob = log_prob.sum(-1, keepdim=True)
-
-		q_p = th.min(q1_p, q2_p) - (alpha * log_prob)
-
-		y = r + 0.99 * q_p * (1 - d)
-
-	q1, q2 = net(s, a.squeeze(1))
-		
-	loss = F.mse_loss(q1, y) + F.mse_loss(q2, y)
-
-	c_optimizer.zero_grad()
-	loss.backward()
-	th.nn.utils.clip_grad_norm_(net.parameters(), 0.5)
-	c_optimizer.step()
-
-	if steps % 2 == 0:
-		mu, log_std = policy(s)
-		std = log_std.exp()
-
-		dist = Normal(mu, std)
-		x_t = dist.rsample()
-		y_t = nn.Tanh()(x_t).detach()
-		a_sampled = y_t * env.action_space.high + 0
-
-		log_prob = dist.log_prob(x_t)
-		log_prob -= th.log(th.from_numpy(env.action_space.high) * (1 - y_t.pow(2)) + 1e-6)
-		log_prob = log_prob.sum(-1, keepdim=True)
-
-		policy_loss =  -(th.min(*net(s, a_sampled)) - (alpha * log_prob)).mean()
-
-		pi_optimizer.zero_grad()
-		policy_loss.backward()
-		th.nn.utils.clip_grad_norm_(policy.parameters(), 0.5)
-		pi_optimizer.step()
-
-		if autotune:
-			# Entropy tuning
-			with th.no_grad():
-				mu, log_std = policy(s)
-				x_t = Normal(mu, log_std.exp()).rsample()
-
-				log_prob = dist.log_prob(x_t)
-				log_prob -= th.log(th.from_numpy(env.action_space.high) * (1 - y_t.pow(2)) + 1e-6)
-				log_prob = log_prob.sum(-1, keepdim=True)
-				
-			alpha_loss = -(log_alpha * (log_prob + target_entropy).detach()).mean()
-
-			a_optimizer.zero_grad()
-			alpha_loss.backward()
-			a_optimizer.step()
-			alpha = log_alpha.exp().item()
-	
-		for targ_params, params in zip(targ_net.parameters(), net.parameters()):
-			targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
\ No newline at end of file
diff --git a/src/testing_td3.py b/src/testing_td3.py
deleted file mode 100644
index 7d923be..0000000
--- a/src/testing_td3.py
+++ /dev/null
@@ -1,110 +0,0 @@
-import copy
-import torch as th
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.distributions import Normal
-import numpy as np
-import gymnasium as gym
-from runner import get_offpolicy_runner
-
-env = gym.make('gym_cartpole_continuous:CartPoleContinuous-v1')
-
-runner = get_offpolicy_runner(
-	env, 
-	'deterministic',
-	length=1,
-	capacity=1000000, 
-	batch_size=100, 
-	train_after=10000)
-
-class Q_critic(nn.Module):
-	def __init__(self, state_dim, action_dim):
-		super(Q_critic, self).__init__()
-
-		self.net1 = nn.Sequential(
-			nn.Linear(state_dim+action_dim, 400),
-			nn.ReLU(),
-			nn.Linear(400, 300),
-			nn.ReLU(),
-			nn.Linear(300, 1))
-		
-		self.net2 = nn.Sequential(
-			nn.Linear(state_dim+action_dim, 400),
-			nn.ReLU(),
-			nn.Linear(400, 300),
-			nn.ReLU(),
-			nn.Linear(300, 1))
-
-	def forward(self, state, action):
-		sa = th.concat([state, action], dim=-1)
-		q1, q2 = self.net1(sa), self.net2(sa)
-		return q1, q2
-	
-	def q1(self, state, action):
-		sa = th.concat([state, action], dim=-1)
-		q1 = self.net1(sa)
-		return q1
-
-class Policy(nn.Module):
-	def __init__(self, state_dim, action_dim):
-		super(Policy, self).__init__()
-
-		self.net = nn.Sequential(
-			nn.Linear(state_dim, 400),
-			nn.ReLU(),
-			nn.Linear(400, 300),
-			nn.ReLU(),
-			nn.Linear(300, action_dim))
-		
-	def forward(self, state):
-		return self.net(state)
-
-net = Q_critic(4, 1)
-targ_net = copy.deepcopy(net)
-policy = Policy(4, 1)
-targ_policy = copy.deepcopy(policy)
-c_optimizer = th.optim.Adam(net.parameters(), lr=1e-3)
-pi_optimizer = th.optim.Adam(policy.parameters(), lr=1e-3)
-
-# https://stable-baselines3.readthedocs.io/en/master/modules/sac.html
-# Need to debug, not performing well enough at all, learning collapses
-
-for steps in range(50000):
-	batch = runner.get_batch(policy)
-
-	s, a, r, s_p, d = batch()
-
-	with th.no_grad():
-		a_p = targ_policy(s_p)
-
-		mean = th.zeros_like(a_p)
-		noise = th.normal(mean=mean, std=0.1).clamp(-0.5, 0.5)
-
-		q1_p, q2_p = targ_net(s_p, (a_p + noise).clamp(-1, 1))
-		q_p = th.min(q1_p, q2_p)
-
-		y = r + 0.98 * q_p * (1 - d)
-
-	q1, q2 = net(s, a.squeeze(1))
-		
-	loss = F.mse_loss(q1, y) + F.mse_loss(q2, y)
-
-	c_optimizer.zero_grad()
-	loss.backward()
-	th.nn.utils.clip_grad_norm_(net.parameters(), 0.5)
-	c_optimizer.step()
-
-	if steps % 2 == 0:
-
-		policy_loss = -net.q1(s, policy(s)).mean()
-
-		pi_optimizer.zero_grad()
-		policy_loss.backward()
-		th.nn.utils.clip_grad_norm_(policy.parameters(), 0.5)
-		pi_optimizer.step()
-	
-		for targ_params, params in zip(targ_net.parameters(), net.parameters()):
-			targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
-
-		for targ_params, params in zip(targ_policy.parameters(), policy.parameters()):
-			targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
\ No newline at end of file
diff --git a/test_q_lambda.py b/test_q_lambda.py
new file mode 100644
index 0000000..f9b5641
--- /dev/null
+++ b/test_q_lambda.py
@@ -0,0 +1,100 @@
+import copy
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from cardio_rl import Runner
+from cardio_rl import Collector
+
+
+class Qfunc(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Qfunc, self).__init__()
+		self.l1 = nn.Linear(state_dim, 64)
+		self.l2 = nn.Linear(64, 64)
+		self.l3 = nn.Linear(64, action_dim)
+
+	def forward(self, state):
+		q = F.relu(self.l1(state))
+		q = F.relu(self.l2(q))
+		q = self.l3(q)		
+		return q
+	
+def main():
+	env = gym.make('CartPole-v1')
+
+	runner = Runner(
+		env=env,
+		policy='argmax',
+		sampler=True,
+		capacity=1_000_000,
+		batch_size=64,
+		collector=Collector(
+			env=env,
+			rollout_len=4,
+			warmup_len=1_000,
+			n_step=4
+		),
+		reduce=False,
+		backend='pytorch'
+	)
+
+	critic = Qfunc(4, 2)
+	targ_net: nn.Module = copy.deepcopy(critic)
+	optimizer = th.optim.Adam(critic.parameters(), lr=3e-4)
+	gamma = 0.99
+	target_update = 10
+
+	for t in range(10000):
+		batch = runner.get_batch(critic)
+		s, a, r, s_p, d, _ = batch()
+
+		a = a.transpose(2,1)
+		r = r.transpose(2,1)
+		d = d.transpose(2,1)
+
+		"""
+		Q(lambda) implementation, need to look furtehr into watkins and pengs and the difference of the two
+		"""
+
+		lam = 0.8
+
+		y_lam = th.zeros([runner.batch_size,1])
+		
+		with th.no_grad():
+
+			for l in range(r.shape[1]):
+
+				n_rets = th.zeros([runner.batch_size, 1])
+
+				# print(np.power(lam, l), r.shape[1] - l)
+
+				for i in reversed(range(r.shape[1] - l)):
+					
+					if i == (r.shape[1] - (1 + l)):
+						q_p = th.max(targ_net(s_p[:,-(1 + l)]), keepdim=True, dim=-1).values
+						n_rets = r[:,i] + (0.99 * q_p) * (1 - d[:,i])
+					else:
+						n_rets = r[:,i] + (0.99 * n_rets) * (1 - d[:,i])
+
+				y_lam += np.power(lam, l) * n_rets
+
+		y_lam *= (1 - lam)
+
+		s = s[:,0]
+		a = a[:,0]
+
+		q = critic(s).gather(1, a.long())
+
+		loss = F.mse_loss(q, y_lam)
+
+		optimizer.zero_grad()
+		loss.backward()
+		optimizer.step()
+
+		if t % target_update == 0:        
+			targ_net = copy.deepcopy(critic)
+			
+if __name__ == '__main__':
+	main()
\ No newline at end of file
diff --git a/testing_a2c.py b/testing_a2c.py
new file mode 100644
index 0000000..7d94cb3
--- /dev/null
+++ b/testing_a2c.py
@@ -0,0 +1,107 @@
+from cardio_rl import Runner, VectorCollector
+import torch as th
+import gymnasium as gym
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Categorical
+
+
+class Critic(nn.Module):
+	def __init__(self):
+		super(Critic, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(4, 64),
+			nn.ReLU(),
+			nn.Linear(64, 64),
+			nn.ReLU(),
+			nn.Linear(64, 1))
+
+	def forward(self, state):
+		return self.net(state)
+
+class Actor(nn.Module):
+	def __init__(self):
+		super(Actor, self).__init__()
+		
+		self.net = nn.Sequential(
+			nn.Linear(4, 128),
+			nn.ReLU(),
+			nn.Linear(128, 128),
+			nn.ReLU(),
+			nn.Linear(128, 2),
+			nn.Softmax(dim=-1)
+		)
+
+	def forward(self, state):
+		return self.net(state)
+
+env = gym.make('CartPole-v1')
+
+runner = Runner(
+	env=env,
+	policy='categorical',
+	collector=VectorCollector(
+		env=env,
+		num_envs=4,
+		rollout_len=15,
+		logger_kwargs=dict(
+			log_interval=1_000,
+			episode_window=500
+		)
+	),
+	backend='pytorch'
+)
+
+actor = Actor()
+critic = Critic()
+a_optimizer = th.optim.RMSprop(actor.parameters(), lr=0.0007, eps=1e-5)
+c_optimizer = th.optim.RMSprop(critic.parameters(), lr=0.0007, eps=1e-5)
+
+for _ in range(100_000):
+	batch = runner.get_batch(actor)
+	s, a, r, s_p, d, _ = batch()
+
+	r = r.squeeze(-2)
+	a = a.squeeze(-2)
+	d = d.squeeze(-2)
+
+	returns = th.zeros_like(r)
+	running_r = critic(s_p[-1]).squeeze(-1).detach()
+	for t in reversed(range(len(r))):
+		running_r = r[t] + 0.99 * running_r * (1 - d[t])
+		returns[t] = running_r
+
+	values = critic(s).squeeze()
+	critic_loss = F.mse_loss(returns, values)
+
+	c_optimizer.zero_grad()
+	critic_loss.backward()
+	nn.utils.clip_grad_norm_(critic.parameters(), 0.5)
+	c_optimizer.step()
+
+	adv: th.Tensor = (returns - values).detach()
+
+	# adv = ((adv - adv.mean()) / (adv.std() + 1e-5)).detach()
+
+	probs = actor(s)
+	dist = Categorical(probs)
+	policy_loss = -th.mean(dist.log_prob(a) * adv)
+
+	a_optimizer.zero_grad()
+	policy_loss.backward()
+	nn.utils.clip_grad_norm_(actor.parameters(), 0.5)
+	a_optimizer.step()
+	continue
+
+	# Below is for Advantage Weighted Regression (need to explore)
+	# https://github.com/jcwleo/awr-pytorch/tree/master
+	for _ in range(10):
+		probs = actor(s)
+		dist = Categorical(probs)
+		policy_loss = -th.mean(dist.log_prob(a.squeeze(-1)) * th.exp(norm_adv/1.0))
+
+		a_optimizer.zero_grad()
+		policy_loss.backward()
+		# nn.utils.clip_grad_norm_(actor.parameters(), 0.5)
+		a_optimizer.step()
diff --git a/testing_qtopt.py b/testing_qtopt.py
new file mode 100644
index 0000000..9296f82
--- /dev/null
+++ b/testing_qtopt.py
@@ -0,0 +1,133 @@
+import copy
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from gymnasium.wrappers import RescaleAction
+from cardio_rl import Runner, Collector
+from cardio_rl.policies import BasePolicy
+
+
+env_name = 'LunarLanderContinuous-v2'
+env = gym.make(env_name)
+env = RescaleAction(env, -1.0, 1.0)
+
+class CEM():
+	''' 
+	cross-entropy method, as optimization of the action policy 
+	'''
+	def __init__(self, theta_dim, ini_mean_scale=0.0, ini_std_scale=1.0):
+		self.theta_dim = theta_dim
+		self.initialize(ini_mean_scale=ini_mean_scale, ini_std_scale=ini_std_scale)
+
+	def initialize(self, ini_mean_scale=0.0, ini_std_scale=1.0):
+		self.mean = ini_mean_scale*th.ones(self.theta_dim)
+		self.std = ini_std_scale*th.ones(self.theta_dim)
+
+	def sample(self, n):
+		theta_list = self.mean + th.normal(mean=th.zeros([n, self.theta_dim]), std=th.zeros([n, self.theta_dim])) * self.std
+		theta_list = th.clip(theta_list, -1.0, 1.0)
+		return theta_list
+
+	def update(self, selected_samples):
+		self.mean = th.mean(selected_samples, dim=0)
+		# print('mean: ', self.mean)
+		self.std = th.std(selected_samples, dim=0)  # plus the entropy offset, or else easily get 0 std
+		# print('std: ', self.std)
+		# return self.mean, self.std
+
+# https://github.com/quantumiracle/QT_Opt/blob/master/qt_opt_v3.py
+class CemPolicy(BasePolicy):
+	def __init__(self, env, recurrent=False, hidden_dims=0):
+		super().__init__(env, recurrent, hidden_dims)
+		self.cem = CEM(2)
+		self.cem_update_iter = 2
+		self.num_samples = 64
+		self.select_num = 6
+
+	def __call__(self, state, net):
+		state = th.FloatTensor(state)
+		return self.cem_optimal_action(state, net)
+	
+	def cem_optimal_action(self, state, net):
+		states = state.expand(self.num_samples, -1)
+		self.cem.initialize() # every time use a new cem, cem is only for deriving the argmax_a'
+		for _ in range(self.cem_update_iter):
+			actions = self.cem.sample(self.num_samples)
+			q_values = net(states, actions).detach().cpu().reshape(-1) # 2 dim to 1 dim
+			# max_idx=q_values.argsort()[-1]  # select one maximal q
+			idx = q_values.argsort()[-int(self.select_num):]  # select top maximum q
+			selected_actions = actions[idx]
+			self.cem.update(selected_actions)
+	
+		optimal_action = actions[idx[-1]]
+		return optimal_action.numpy()
+
+runner = Runner(
+	env=env,
+	policy=CemPolicy(env),
+	sampler=True,
+	capacity=1_000_000,
+	batch_size=100,
+	collector=Collector(
+		env=env,
+		warmup_len=10_000,
+		logger_kwargs=dict(
+			log_dir = env_name,
+			log_interval = 2_500
+		)
+	),
+	backend='pytorch'
+)
+
+class Q_critic(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Q_critic, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(state_dim+action_dim, 256),
+			nn.ReLU(),
+			nn.Linear(256, 256),
+			nn.ReLU(),
+			nn.Linear(256, 1))
+
+	def forward(self, state, action):
+		sa = th.concat([state, action], dim=-1)
+		return self.net(sa)
+
+# https://stable-baselines3.readthedocs.io/en/master/modules/ddpg.html
+critic = Q_critic(8, 2)
+targ_critic = copy.deepcopy(critic)
+c_optimizer = th.optim.Adam(critic.parameters(), lr=1e-3)
+
+for steps in range(150_000):
+	s, a, r, s_p, d, _ = runner.get_batch(critic)
+
+	with th.no_grad():
+		"""
+		Convert this to parallel instead of using list comprehension
+		"""
+		a_p = np.array([runner.policy.cem_optimal_action(next_state, targ_critic) for next_state in s_p])
+		a_p = th.from_numpy(a_p).float()
+		# exit(a_p)
+
+		mean = th.zeros_like(a_p)
+		noise = th.normal(mean=mean, std=0.2).clamp(-0.5, 0.5)
+		
+		a_p = (a_p + noise).clamp(-1, 1)
+
+		q_p = targ_critic(s_p, a_p)
+		y = r + 0.98 * q_p * (1 - d)
+
+	q = critic(s, a.squeeze(1))
+		
+	loss = F.mse_loss(q, y)
+
+	c_optimizer.zero_grad()
+	loss.backward()
+	th.nn.utils.clip_grad_norm_(critic.parameters(), 0.5)
+	c_optimizer.step()
+
+	for targ_params, params in zip(targ_critic.parameters(), critic.parameters()):
+		targ_params.data.copy_(params.data * 0.005 + targ_params.data * (1.0 - 0.005))
diff --git a/testing_rainbow.py b/testing_rainbow.py
new file mode 100644
index 0000000..1a3239b
--- /dev/null
+++ b/testing_rainbow.py
@@ -0,0 +1,105 @@
+import copy
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import gymnasium as gym
+from cardio_rl import Runner
+from cardio_rl import Collector
+from cardio_rl.rainbow_naf.noisy_linear import NoisyLinear
+
+"""
+Duelling Double DQN with n-step returns used for verification of runner functionality
+-double check NoisyNet implementation
+-should maybe implement distributional q-learning (qr or iqn) and PER to fully make rainbow
+-PER will be interesting, need to figure out how to do it with current design structure
+"""
+
+class Q_duelling(nn.Module):
+	def __init__(self, state_dim, action_dim, dueling):
+		super(Q_duelling, self).__init__()
+		self.dueling = dueling
+
+		self.l1 = NoisyLinear(state_dim, 128)
+
+		self.a1 = NoisyLinear(128, 128)
+		self.a2 = NoisyLinear(128, action_dim)
+
+		self.v1 = NoisyLinear(128, 128)
+		self.v2 = NoisyLinear(128, 1)
+
+	def forward(self, state):
+		q = F.relu(self.l1(state))
+
+		a = F.relu(self.a1(q))
+		a = self.a2(a)
+
+		v = F.relu(self.v1(q))
+		v = self.v2(v)
+
+		if self.dueling:
+			return v + (a - a.mean())
+
+		return a   
+	
+	def reset_noise(self):
+		"""Reset all noisy layers."""
+		self.l1.reset_noise()
+		self.a1.reset_noise()
+		self.a2.reset_noise()
+		self.v1.reset_noise()
+		self.v2.reset_noise()
+
+env = gym.make('CartPole-v1')
+
+runner = Runner(
+	env=env,
+	policy='argmax',	# Epsilon_argmax_policy(env, 0.5, 0.05, 0.9),
+	sampler=True,
+	capacity=100000,
+	batch_size=64,
+	collector=Collector(
+		env=env,
+		rollout_len=4,
+		warmup_len=1000,
+		n_step=3
+	),
+	backend='pytorch'
+)
+
+critic = Q_duelling(4, 2, True)
+targ_net = copy.deepcopy(critic)
+optimizer = th.optim.Adam(critic.parameters(), lr=2.3e-3)
+gamma = 0.99
+target_update = 10
+
+for t in range(5000):
+	batch = runner.get_batch(critic)
+	s, a, r, s_p, d, _ = batch()
+
+	if runner.n_step !=  1:
+		r_nstep = th.zeros([len(r), 1])
+		for n in reversed(range(runner.n_step)):
+			r_nstep += r[:,:,n] * gamma
+
+		r = r_nstep
+
+	q = critic(s).gather(1, a.long())
+
+	with th.no_grad():
+		a_p = critic(s_p).argmax(dim=1).unsqueeze(1)		
+		q_p = targ_net(s_p).gather(1, a_p.long())		
+		y = r + pow(gamma, runner.n_step) * q_p * (1 - d)
+
+	loss = F.mse_loss(q, y)
+
+	optimizer.zero_grad()
+	loss.backward()
+	th.nn.utils.clip_grad_norm_(critic.parameters(), 1.0)
+	optimizer.step()
+	
+	critic.reset_noise()
+	targ_net.reset_noise()
+
+	if t % target_update == 0:        
+		targ_net = copy.deepcopy(critic)
+		
\ No newline at end of file
diff --git a/testing_sac.py b/testing_sac.py
new file mode 100644
index 0000000..2a9974a
--- /dev/null
+++ b/testing_sac.py
@@ -0,0 +1,209 @@
+import copy
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Beta
+import gymnasium as gym
+from gymnasium.wrappers import RescaleAction
+
+from cardio_rl import Runner, Collector
+
+# env_name = 'LunarLanderContinuous-v2'
+env_name = 'BipedalWalker-v3'
+# env_name = 'Pendulum-v1'
+env = gym.make(env_name)
+env = RescaleAction(env, 0, 1)
+
+"""
+Soft-actor critic with beta policy
+
+# https://stable-baselines3.readthedocs.io/en/master/modules/sac.html
+# https://github.com/DLR-RM/rl-baselines3-zoo/blob/master/hyperparams/sac.yml
+"""
+
+runner = Runner(
+	env=env,
+	policy='beta',
+	sampler=True,
+	capacity=300_000,
+	batch_size=256,
+	n_batches=64,
+	collector=Collector(
+		env=env,
+		rollout_len=64,
+		warmup_len=10_000,
+		logger_kwargs=dict(
+            tensorboard=True,
+	    	log_dir=env_name,
+		    log_interval=5_000,
+		    exp_name = 'vanilla_selection_awr_test_0.05_policy_decay'	# _no_WD'
+		)
+	),
+	backend='pytorch'
+)
+
+class DroQ(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(DroQ, self).__init__()
+
+		# self.net = nn.Sequential(
+		# 	nn.Linear(state_dim+action_dim, 400),
+		# 	nn.Dropout(0.01),
+		# 	nn.LayerNorm(400),
+		# 	nn.ReLU(),
+		# 	nn.Linear(400, 300),
+		# 	# nn.Dropout(0.01),
+		# 	# nn.LayerNorm(300),
+		# 	# nn.ReLU(),
+		# 	# nn.Linear(300, 300),
+		# 	nn.Dropout(0.01),
+		# 	nn.LayerNorm(300),
+		# 	nn.ReLU(),
+		# 	nn.Linear(300, 1))
+		
+		self.net = nn.Sequential(
+			nn.Linear(state_dim+action_dim, 400),
+			nn.ReLU(),
+			nn.Linear(400, 300),
+			nn.ReLU(),
+			nn.Linear(300, 1)
+		)
+
+	def forward(self, sa):
+		return self.net(sa)
+
+class Q_critic(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Q_critic, self).__init__()
+
+		self.net1 = DroQ(state_dim, action_dim)
+		self.net2 = DroQ(state_dim, action_dim)
+
+	def forward(self, state, action):
+		sa = th.concat([state, action], dim=-1)
+		q1, q2 = self.net1(sa), self.net2(sa)
+		return q1, q2
+
+class Policy(nn.Module):
+	def __init__(self, state_dim, action_dim):
+		super(Policy, self).__init__()
+
+		self.net = nn.Sequential(
+			nn.Linear(state_dim, 400),
+			nn.ReLU(),
+			nn.Linear(400, 300),
+			nn.ReLU(),)
+		
+		self.alpha = nn.Sequential(
+			nn.Linear(300, action_dim),
+			nn.Softplus()
+		)
+
+		self.beta = nn.Sequential(
+			nn.Linear(300, action_dim),
+			nn.Softplus()
+		)
+		
+	def forward(self, state):
+		h = self.net(state)
+		a = self.alpha(h)
+		b = self.beta(h)
+		return 1+a, 1+b
+
+critic = Q_critic(24, 4)
+actor = Policy(24, 4)
+targ_critic= copy.deepcopy(critic)
+targ_actor = copy.deepcopy(actor)
+c_optimizer = th.optim.Adam(critic.parameters(), lr=7.3e-4)	#, weight_decay=2.4e-4)	# weight_decay value from Relu to the rescue paper
+a_optimizer = th.optim.Adam(actor.parameters(), lr=7.3e-4, weight_decay=2.4e-4)	# weight_decay value from Relu to the rescue paper
+
+# cleanrl optimises the log of alpha and computes alpha through it each iteration
+#^can apply this to AWAC dual variables?
+# also check out what cleanrl implements for atari in SAC: https://docs.cleanrl.dev/rl-algorithms/sac/#implementation-details_1
+log_alpha = nn.Parameter(th.tensor([1.0], requires_grad=True))
+ent_coeff = log_alpha.exp().item()	# nn.Parameter(th.tensor([0.2], requires_grad=True))
+ent_optim = th.optim.Adam([log_alpha], lr=7.3e-4)	# check if lr is same as critic and actor
+H = -4
+tau = 0.02
+
+for steps in range(4_532):
+	n_batches = runner.get_batch(targ_actor)	# using target actor, prevents policy-churn
+	# print(len(n_batches))
+	# exit()
+	for batch in n_batches:
+
+		s, a, r, s_p, d, _ = batch()
+
+		for _ in range(1):
+			with th.no_grad():
+				# use target actor net here?
+				alpha, beta = actor(s_p)
+				dist_p = Beta(alpha, beta)
+				a_p = dist_p.rsample()
+
+				q1_p, q2_p = targ_critic(s_p, a_p)
+
+				log_prob = dist_p.log_prob(a_p)
+				log_prob = log_prob.sum(-1, keepdim=True)
+
+				q_p = th.min(q1_p, q2_p) - (ent_coeff * log_prob)
+
+				y = r + 0.98 * q_p * (1 - d)
+
+			q1, q2 = critic(s, a.squeeze(1))
+				
+			loss = F.mse_loss(q1, y) + F.mse_loss(q2, y)
+
+			c_optimizer.zero_grad()
+			loss.backward()
+			th.nn.utils.clip_grad_norm_(critic.parameters(), 10.0)
+			c_optimizer.step()
+
+		if steps % 2 == 0:
+			alpha, beta = actor(s_p)
+			dist = Beta(alpha, beta)
+			a_sampled = dist.rsample()
+
+			log_prob = dist.log_prob(a_sampled)
+			log_prob = log_prob.sum(-1, keepdim=True)
+
+			# AWR-like policy updates for contraining current policy to behavioural 
+			with th.no_grad():
+				# Need to implement ADV-max like in CRR
+				q_adv = th.min(*critic(s, a.squeeze(1)))
+				actions = dist.sample((10,)).transpose(0,1)
+				# print(a_sampled.shape, actions.shape)
+				j_samples = critic(s.unsqueeze(1).expand(-1, 10, -1), actions)[0]
+				v_adv = th.max(j_samples, dim=1).values
+
+				# v_adv = th.min(*critic(s, a_sampled))
+				adv = q_adv - v_adv
+				# weight = adv.exp()
+				# weight = th.clamp_max(weight, 20)
+
+				# Binary
+				weight = (adv>0.0).float()
+
+			awr_loss = (dist.log_prob(a.squeeze(1)).sum(-1, keepdim=True) * weight) * 0.05
+			
+			policy_loss =  -(th.min(*critic(s, a_sampled)) - (ent_coeff * log_prob) + awr_loss).mean()
+
+			a_optimizer.zero_grad()
+			policy_loss.backward()
+			th.nn.utils.clip_grad_norm_(actor.parameters(), 10.0)
+			a_optimizer.step()		
+
+			ent_loss = -log_alpha.exp() * (log_prob.detach() + H).mean()
+			ent_optim.zero_grad()
+			ent_loss.backward()
+			ent_optim.step()	
+
+			ent_coeff = log_alpha.exp().item()
+
+			# update targets only when actor is updated or not? TD3 paper only updates targets during policy updates
+			# in contrast, CleanRL's SAC does so every update step (could be a matter of performance vs stability)
+		for targ_params, params in zip(targ_actor.parameters(), actor.parameters()):
+			targ_params.data.copy_(params.data * tau + targ_params.data * (1.0 - tau))
+
+		for targ_params, params in zip(targ_critic.parameters(), critic.parameters()):
+			targ_params.data.copy_(params.data * tau + targ_params.data * (1.0 - tau))