diff --git a/basic_tf.py b/basic_tf.py
new file mode 100644
index 0000000..e274431
--- /dev/null
+++ b/basic_tf.py
@@ -0,0 +1,35 @@
+# TensorFlow and tf.keras sample to make sure the tf install is running
+import tensorflow as tf
+
+# Helper libraries
+import numpy as np
+import matplotlib.pyplot as plt
+
+fashion_mnist = tf.keras.datasets.fashion_mnist
+
+(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
+
+class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
+               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
+
+train_images = train_images / 255.0
+
+test_images = test_images / 255.0
+
+model = tf.keras.Sequential([
+    tf.keras.layers.Flatten(input_shape=(28, 28)),
+    tf.keras.layers.Dense(128, activation='relu'),
+    tf.keras.layers.Dense(10)
+])
+
+model.compile(optimizer='adam',
+              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+              metrics=['accuracy'])
+
+model.fit(train_images, train_labels, epochs=10)
+
+test_loss, test_acc = model.evaluate(test_images,  test_labels, verbose=2)
+
+print('\nTest accuracy:', test_acc)
+
+
diff --git a/clean_sac.py b/clean_sac.py
new file mode 100644
index 0000000..922c032
--- /dev/null
+++ b/clean_sac.py
@@ -0,0 +1,361 @@
+# docs and experiment results can be found at https://docs.cleanrl.dev/rl-algorithms/sac/#sac_continuous_actionpy
+import os
+import random
+import time
+from dataclasses import dataclass
+
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import tyro
+from stable_baselines3.common.buffers import ReplayBuffer
+from torch.utils.tensorboard import SummaryWriter
+
+from continualworld.envs import get_cl_env, get_single_env
+from continualworld.tasks import TASK_SEQS
+
+
+@dataclass
+class Args:
+    exp_name: str = os.path.basename(__file__)[: -len(".py")]
+    """the name of this experiment"""
+    seed: int = 1
+    """seed of the experiment"""
+    torch_deterministic: bool = True
+    """if toggled, `torch.backends.cudnn.deterministic=False`"""
+    cuda: bool = True
+    """if toggled, cuda will be enabled by default"""
+    track: bool = False
+    """if toggled, this experiment will be tracked with Weights and Biases"""
+    wandb_project_name: str = "cleanRL"
+    """the wandb's project name"""
+    wandb_entity: str = None
+    """the entity (team) of wandb's project"""
+    capture_video: bool = False
+    """whether to capture videos of the agent performances (check out `videos` folder)"""
+
+    # Algorithm specific arguments
+    env_id: str = "Hopper-v4"
+    """the environment id of the task"""
+    # num_envs: int = 2
+    # """the environment id of the task"""
+    total_timesteps: int = 1000000
+    """total timesteps of the experiments"""
+    buffer_size: int = int(1e6)
+    """the replay memory buffer size"""
+    gamma: float = 0.99
+    """the discount factor gamma"""
+    tau: float = 0.005
+    """target smoothing coefficient (default: 0.005)"""
+    batch_size: int = 256
+    """the batch size of sample from the reply memory"""
+    learning_starts: int = 5e3
+    """timestep to start learning"""
+    policy_lr: float = 3e-4
+    """the learning rate of the policy network optimizer"""
+    q_lr: float = 1e-3
+    """the learning rate of the Q network network optimizer"""
+    policy_frequency: int = 2
+    """the frequency of training policy (delayed)"""
+    target_network_frequency: int = 1  # Denis Yarats' implementation delays this by 2.
+    """the frequency of updates for the target nerworks"""
+    noise_clip: float = 0.5
+    """noise clip parameter of the Target Policy Smoothing Regularization"""
+    alpha: float = 0.2
+    """Entropy regularization coefficient."""
+    autotune: bool = True
+    """automatic tuning of the entropy coefficient"""
+
+
+def make_env(env_id, seed, idx, capture_video, run_name):
+    def thunk():
+        if capture_video and idx == 0:
+            env = gym.make(env_id, render_mode="rgb_array")
+            env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
+        else:
+            env = gym.make(env_id)
+        env = gym.wrappers.RecordEpisodeStatistics(env)
+        env.action_space.seed(seed)
+        return env
+
+    return thunk
+
+
+# ALGO LOGIC: initialize agent here:
+class SoftQNetwork(nn.Module):
+    def __init__(self, env):
+        super().__init__()
+        self.fc1 = nn.Linear(np.array(env.single_observation_space.shape).prod() + np.prod(env.single_action_space.shape), 256)
+        self.fc2 = nn.Linear(256, 256)
+        self.fc3 = nn.Linear(256, 1)
+
+    def forward(self, x, a):
+        x = torch.cat([x, a], 1)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+
+LOG_STD_MAX = 2
+LOG_STD_MIN = -5
+
+
+class Actor(nn.Module):
+    def __init__(self, env):
+        super().__init__()
+        ## TODO balloch: if making continual world vectorized, change "action_space" and "observation_space" back to "single_action_space" and"single_observation_space"
+        self.fc1 = nn.Linear(np.array(env.single_observation_space.shape).prod(), 256)
+        self.fc2 = nn.Linear(256, 256)
+        self.fc_mean = nn.Linear(256, np.prod(env.single_action_space.shape))
+        self.fc_logstd = nn.Linear(256, np.prod(env.single_action_space.shape))
+        # action rescaling
+        self.register_buffer(
+            "action_scale", torch.tensor((env.action_space.high - env.action_space.low) / 2.0, dtype=torch.float32)
+        )
+        self.register_buffer(
+            "action_bias", torch.tensor((env.action_space.high + env.action_space.low) / 2.0, dtype=torch.float32)
+        )
+
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        mean = self.fc_mean(x)
+        log_std = self.fc_logstd(x)
+        log_std = torch.tanh(log_std)
+        log_std = LOG_STD_MIN + 0.5 * (LOG_STD_MAX - LOG_STD_MIN) * (log_std + 1)  # From SpinUp / Denis Yarats
+
+        return mean, log_std
+
+    def get_action(self, x):
+        mean, log_std = self(x)
+        std = log_std.exp()
+        normal = torch.distributions.Normal(mean, std)
+        x_t = normal.rsample()  # for reparameterization trick (mean + std * N(0,1))
+        y_t = torch.tanh(x_t)
+        action = y_t * self.action_scale + self.action_bias
+        log_prob = normal.log_prob(x_t)
+        # Enforcing Action Bound
+        log_prob -= torch.log(self.action_scale * (1 - y_t.pow(2)) + 1e-6)
+        if len(log_prob.shape) < 2:  # for flat vectors
+            log_prob = log_prob.sum()
+        else:
+            log_prob = log_prob.sum(1, keepdim=True)
+        mean = torch.tanh(mean) * self.action_scale + self.action_bias
+        return action, log_prob, mean
+
+
+if __name__ == "__main__":
+    import stable_baselines3 as sb3
+
+    if sb3.__version__ < "2.0":
+        raise ValueError(
+            """Ongoing migration: run the following command to install the new dependencies:
+poetry run pip install "stable_baselines3==2.0.0a1"
+"""
+        )
+
+    args = tyro.cli(Args)
+    run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
+    if args.track:
+        import wandb
+
+        wandb.init(
+            project=args.wandb_project_name,
+            entity=args.wandb_entity,
+            sync_tensorboard=True,
+            config=vars(args),
+            name=run_name,
+            monitor_gym=True,
+            save_code=True,
+        )
+    writer = SummaryWriter(f"runs/{run_name}")
+    writer.add_text(
+        "hyperparameters",
+        "|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])),
+    )
+
+    # TRY NOT TO MODIFY: seeding
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    torch.backends.cudnn.deterministic = args.torch_deterministic
+
+    device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
+
+    # env setup
+    if args.env_id == 'CW10':
+        tasks='CW10'  # TODO unhardcode this
+        if tasks is not None:
+            tasks = TASK_SEQS[tasks]
+        steps_per_task = 1000
+        envs = get_cl_env(tasks, steps_per_task)
+        # Consider normalizing test envs in the future.
+        num_tasks = len(tasks)
+        test_envs = [
+            get_single_env(task, one_hot_idx=i, one_hot_len=num_tasks) for i, task in enumerate(tasks)
+        ]
+        args.total_timesteps = steps_per_task * len(tasks)
+        envs.single_observation_space = envs.observation_space
+        envs.single_action_space = envs.action_space
+        num_envs = 1  # One at a time
+
+    else:
+        num_envs = 1  # should be an arg for vectorization...
+        envs = gym.vector.SyncVectorEnv([make_env(args.env_id, args.seed, 0, args.capture_video, run_name)])
+        assert isinstance(envs.single_action_space, gym.spaces.Box), "only continuous action space is supported"
+
+    # Ensure matching types, has "kind" syntax per https://numpy.org/doc/stable/reference/generated/numpy.dtype.kind.html
+    envs.single_observation_space.dtype = np.dtype('f4')
+    if args.env_id == 'CW10':
+        envs.single_action_space.dtype = np.dtype('f4')
+
+
+
+    actor = Actor(envs).to(device)
+    qf1 = SoftQNetwork(envs).to(device)
+    qf2 = SoftQNetwork(envs).to(device)
+    qf1_target = SoftQNetwork(envs).to(device)
+    qf2_target = SoftQNetwork(envs).to(device)
+    qf1_target.load_state_dict(qf1.state_dict())
+    qf2_target.load_state_dict(qf2.state_dict())
+    q_optimizer = optim.Adam(list(qf1.parameters()) + list(qf2.parameters()), lr=args.q_lr)
+    actor_optimizer = optim.Adam(list(actor.parameters()), lr=args.policy_lr)
+
+    # Automatic entropy tuning
+    if args.autotune:
+        target_entropy = -torch.prod(torch.Tensor(envs.single_action_space.shape).to(device)).item()
+        log_alpha = torch.zeros(1, requires_grad=True, device=device)
+        alpha = log_alpha.exp().item()
+        a_optimizer = optim.Adam([log_alpha], lr=args.q_lr)
+    else:
+        alpha = args.alpha
+
+    rb = ReplayBuffer(
+        buffer_size=args.buffer_size,
+        observation_space=envs.single_observation_space,
+        action_space=envs.single_action_space,
+        device=device,
+        n_envs=num_envs, #envs.num_envs
+        handle_timeout_termination=False,
+    )
+
+
+    start_time = time.time()
+
+    # TRY NOT TO MODIFY: start the game
+    obs, _ = envs.reset(seed=args.seed)
+    for global_step in range(args.total_timesteps):
+        # ALGO LOGIC: put action logic here
+        if global_step < args.learning_starts:
+            actions = np.array([envs.single_action_space.sample() for _ in range(num_envs)]) 
+        else:
+            actions, _, _ = actor.get_action(torch.Tensor(obs).to(device))
+            actions = actions.detach().cpu().numpy()
+
+        # TRY NOT TO MODIFY: execute the game and log data.
+        if args.env_id == 'CW10':
+            next_obs, rewards, terminations, truncations, infos = envs.step(actions.flatten())
+        else:
+            next_obs, rewards, terminations, truncations, infos = envs.step(actions)
+
+        # TRY NOT TO MODIFY: record rewards for plotting purposes
+        if args.env_id == 'CW10':  # Non-vector envs
+            if 'episode' in infos:
+                print(f"global_step={global_step}, episodic_return={info['episode']['r']}")
+                writer.add_scalar("charts/episodic_return", info["episode"]["r"], global_step)
+                writer.add_scalar("charts/episodic_length", info["episode"]["l"], global_step)
+                break
+
+        if "final_info" in infos:
+            for info in infos["final_info"]:
+                if info:
+                    print(f"global_step={global_step}, episodic_return={info['episode']['r']}")
+                    writer.add_scalar("charts/episodic_return", info["episode"]["r"], global_step)
+                    writer.add_scalar("charts/episodic_length", info["episode"]["l"], global_step)
+                    break
+
+        # TRY NOT TO MODIFY: save data to reply buffer; handle `final_observation`
+        real_next_obs = next_obs.copy()
+        if args.env_id != 'CW10':  ## only for vectorized environments
+            for idx, trunc in enumerate(truncations):
+                if trunc:
+                    real_next_obs[idx] = infos["final_observation"][idx]
+        else:
+            if truncations:
+                envs.reset()
+        rb.add(obs, real_next_obs, actions, rewards, terminations, infos)
+
+        # TRY NOT TO MODIFY: CRUCIAL step easy to overlook
+        obs = next_obs
+
+        # ALGO LOGIC: training.
+        if global_step > args.learning_starts:
+            data = rb.sample(args.batch_size)
+            with torch.no_grad():
+                next_state_actions, next_state_log_pi, _ = actor.get_action(data.next_observations)
+                qf1_next_target = qf1_target(data.next_observations, next_state_actions)
+                qf2_next_target = qf2_target(data.next_observations, next_state_actions)
+                min_qf_next_target = torch.min(qf1_next_target, qf2_next_target) - alpha * next_state_log_pi
+                next_q_value = data.rewards.flatten() + (1 - data.dones.flatten()) * args.gamma * (min_qf_next_target).view(-1)
+
+            qf1_a_values = qf1(data.observations, data.actions).view(-1)
+            qf2_a_values = qf2(data.observations, data.actions).view(-1)
+            qf1_loss = F.mse_loss(qf1_a_values, next_q_value)
+            qf2_loss = F.mse_loss(qf2_a_values, next_q_value)
+            qf_loss = qf1_loss + qf2_loss
+
+            # optimize the model
+            q_optimizer.zero_grad()
+            qf_loss.backward()
+            q_optimizer.step()
+
+            if global_step % args.policy_frequency == 0:  # TD 3 Delayed update support
+                for _ in range(
+                    args.policy_frequency
+                ):  # compensate for the delay by doing 'actor_update_interval' instead of 1
+                    pi, log_pi, _ = actor.get_action(data.observations)
+                    qf1_pi = qf1(data.observations, pi)
+                    qf2_pi = qf2(data.observations, pi)
+                    min_qf_pi = torch.min(qf1_pi, qf2_pi)
+                    actor_loss = ((alpha * log_pi) - min_qf_pi).mean()
+
+                    actor_optimizer.zero_grad()
+                    actor_loss.backward()
+                    actor_optimizer.step()
+
+                    if args.autotune:
+                        with torch.no_grad():
+                            _, log_pi, _ = actor.get_action(data.observations)
+                        alpha_loss = (-log_alpha.exp() * (log_pi + target_entropy)).mean()
+
+                        a_optimizer.zero_grad()
+                        alpha_loss.backward()
+                        a_optimizer.step()
+                        alpha = log_alpha.exp().item()
+
+            # update the target networks
+            if global_step % args.target_network_frequency == 0:
+                for param, target_param in zip(qf1.parameters(), qf1_target.parameters()):
+                    target_param.data.copy_(args.tau * param.data + (1 - args.tau) * target_param.data)
+                for param, target_param in zip(qf2.parameters(), qf2_target.parameters()):
+                    target_param.data.copy_(args.tau * param.data + (1 - args.tau) * target_param.data)
+
+            if global_step % 100 == 0:
+                writer.add_scalar("losses/qf1_values", qf1_a_values.mean().item(), global_step)
+                writer.add_scalar("losses/qf2_values", qf2_a_values.mean().item(), global_step)
+                writer.add_scalar("losses/qf1_loss", qf1_loss.item(), global_step)
+                writer.add_scalar("losses/qf2_loss", qf2_loss.item(), global_step)
+                writer.add_scalar("losses/qf_loss", qf_loss.item() / 2.0, global_step)
+                writer.add_scalar("losses/actor_loss", actor_loss.item(), global_step)
+                writer.add_scalar("losses/alpha", alpha, global_step)
+                print("SPS:", int(global_step / (time.time() - start_time)))
+                writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)
+                if args.autotune:
+                    writer.add_scalar("losses/alpha_loss", alpha_loss.item(), global_step)
+
+    envs.close()
+    writer.close()
\ No newline at end of file
diff --git a/continualworld/envs.py b/continualworld/envs.py
index f018743..cd6ca1f 100644
--- a/continualworld/envs.py
+++ b/continualworld/envs.py
@@ -1,10 +1,10 @@
 from copy import deepcopy
 from typing import Any, Dict, List, Tuple, Union
 
-import gym
+import gymnasium as gym
 import metaworld
 import numpy as np
-from gym.wrappers import TimeLimit
+from gymnasium.wrappers import TimeLimit
 
 from continualworld.utils.wrappers import OneHotAdder, RandomizationWrapper, SuccessCounter
 
@@ -126,7 +126,7 @@ def pop_successes(self) -> List[bool]:
 
     def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict]:
         self._check_steps_bound()
-        obs, reward, done, info = self.envs[self.cur_seq_idx].step(action)
+        obs, reward, terminated, truncated, info = self.envs[self.cur_seq_idx].step(action)
         info["seq_idx"] = self.cur_seq_idx
 
         self.cur_step += 1
@@ -138,7 +138,7 @@ def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict]:
 
             self.cur_seq_idx += 1
 
-        return obs, reward, done, info
+        return obs, reward, terminated, truncated, info
 
     def reset(self) -> np.ndarray:
         self._check_steps_bound()
@@ -214,20 +214,20 @@ def pop_successes(self) -> List[bool]:
 
     def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict]:
         self._check_steps_bound()
-        obs, reward, done, info = self.envs[self._cur_seq_idx].step(action)
+        obs, reward, terminated, truncated, info = self.envs[self._cur_seq_idx].step(action)
         info["mt_seq_idx"] = self._cur_seq_idx
         if self.cycle_mode == "step":
             self._cur_seq_idx = (self._cur_seq_idx + 1) % self.num_envs
         self.cur_step += 1
 
-        return obs, reward, done, info
+        return obs, reward, terminated, truncated, info
 
     def reset(self) -> np.ndarray:
         self._check_steps_bound()
         if self.cycle_mode == "episode":
             self._cur_seq_idx = (self._cur_seq_idx + 1) % self.num_envs
-        obs = self.envs[self._cur_seq_idx].reset()
-        return obs
+        obs, info = self.envs[self._cur_seq_idx].reset()
+        return obs, info
 
 
 def get_mt_env(
diff --git a/continualworld/methods/vcl.py b/continualworld/methods/vcl.py
index 61de631..788fca4 100644
--- a/continualworld/methods/vcl.py
+++ b/continualworld/methods/vcl.py
@@ -1,6 +1,6 @@
 from typing import Callable, Iterable, List, Tuple
 
-import gym
+import gymnasium as gym
 import tensorflow as tf
 import tensorflow.keras as tfk
 from tensorflow.keras import Input, Model
diff --git a/continualworld/sac/models.py b/continualworld/sac/models.py
index c0dbd7e..f3f78e0 100644
--- a/continualworld/sac/models.py
+++ b/continualworld/sac/models.py
@@ -1,6 +1,6 @@
 from typing import Callable, Iterable, List, Tuple
 
-import gym
+import gymnasium as gym
 import numpy as np
 import tensorflow as tf
 from tensorflow.keras import Input, Model
diff --git a/continualworld/sac/sac.py b/continualworld/sac/sac.py
index 627e641..994889d 100644
--- a/continualworld/sac/sac.py
+++ b/continualworld/sac/sac.py
@@ -4,7 +4,7 @@
 import time
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
-import gym
+import gymnasium as gym
 import numpy as np
 import tensorflow as tf
 
@@ -180,8 +180,8 @@ def __init__(
             + self.critic1.common_variables
             + self.critic2.common_variables
         )
-
-        self.optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
+        self.optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=lr)
+        # self.optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
 
         # For reference on automatic alpha tuning, see
         # "Automating Entropy Adjustment for Maximum Entropy" section
@@ -232,7 +232,7 @@ def get_episodic_batch(self, current_task_idx: int) -> Optional[Dict[str, tf.Ten
     def get_log_alpha(self, obs: tf.Tensor) -> tf.Tensor:
         return tf.squeeze(tf.linalg.matmul(obs[:, -self.num_tasks :], self.all_log_alpha))
 
-    @tf.function
+    # @tf.function
     def get_action(self, o: tf.Tensor, deterministic: tf.Tensor = tf.constant(False)) -> tf.Tensor:
         mu, log_std, pi, logp_pi = self.actor(tf.expand_dims(o, 0))
         if deterministic:
@@ -246,7 +246,10 @@ def get_action_test(
         return self.get_action(o, deterministic)
 
     def get_learn_on_batch(self, current_task_idx: int) -> Callable:
-        @tf.function
+        # TODO : decorator causes error:
+        # : CommandLine Error: Option 'help-list' registered more than once!
+        # LLVM ERROR: inconsistency in registered CommandLine options
+        # @tf.function  
         def learn_on_batch(
             seq_idx: tf.Tensor,
             batch: Dict[str, tf.Tensor],
@@ -407,9 +410,12 @@ def test_agent(self, deterministic, num_episodes) -> None:
             self.on_test_start(seq_idx)
 
             for j in range(num_episodes):
-                obs, done, episode_return, episode_len = test_env.reset(), False, 0, 0
+                obs, info = test_env.reset()
+                done = False
+                episode_return = 0
+                episode_len = 0
                 while not (done or (episode_len == self.max_episode_len)):
-                    obs, reward, done, _ = test_env.step(
+                    obs, reward, terminated, truncated, info = test_env.step(
                         self.get_action_test(tf.convert_to_tensor(obs), tf.constant(deterministic))
                     )
                     episode_return += reward
@@ -532,16 +538,19 @@ def _handle_task_change(self, current_task_idx: int):
     def run(self):
         """A method to run the SAC training, after the object has been created."""
         self.start_time = time.time()
-        obs, episode_return, episode_len = self.env.reset(), 0, 0
+        obs, info = self.env.reset()
+        episode_return = 0
+        episode_len = 0
 
         # Main loop: collect experience in env and update/log each epoch
         current_task_timestep = 0
         current_task_idx = -1
-        self.learn_on_batch = self.get_learn_on_batch(current_task_idx)
+        # self.learn_on_batch = self.get_learn_on_batch(current_task_idx)
 
         for global_timestep in range(self.steps):
             # On task change
             if current_task_idx != getattr(self.env, "cur_seq_idx", -1):
+                print("if statement 1")
                 current_task_timestep = 0
                 current_task_idx = getattr(self.env, "cur_seq_idx")
                 self._handle_task_change(current_task_idx)
@@ -552,20 +561,22 @@ def run(self):
             if current_task_timestep > self.start_steps or (
                 self.agent_policy_exploration and current_task_idx > 0
             ):
+                print("if not exploring")
                 action = self.get_action(tf.convert_to_tensor(obs))
             else:
                 action = self.env.action_space.sample()
 
             # Step the env
-            next_obs, reward, done, info = self.env.step(action)
+            next_obs, reward, terminated, truncated, info = self.env.step(action)
             episode_return += reward
             episode_len += 1
 
             # Ignore the "done" signal if it comes from hitting the time
             # horizon (that is, when it's an artificial terminal signal
             # that isn't based on the agent's state)
+            done = np.logical_or(terminated,truncated)
             done_to_store = done
-            if episode_len == self.max_episode_len or info.get("TimeLimit.truncated"):
+            if episode_len == self.max_episode_len or truncated:  # updated for gymnasium
                 done_to_store = False
 
             # Store experience to replay buffer
@@ -579,20 +590,20 @@ def run(self):
             if done or (episode_len == self.max_episode_len):
                 self.logger.store({"train/return": episode_return, "train/ep_length": episode_len})
                 episode_return, episode_len = 0, 0
-                if global_timestep < self.steps - 1:
-                    obs = self.env.reset()
+                if global_timestep < self.steps - 1:  # This may not work with mujoco anymore
+                    obs, info = self.env.reset()
 
             # Update handling
             if (
                 current_task_timestep >= self.update_after
                 and current_task_timestep % self.update_every == 0
             ):
-
                 for j in range(self.update_every):
                     batch = self.replay_buffer.sample_batch(self.batch_size)
 
                     episodic_batch = self.get_episodic_batch(current_task_idx)
 
+                    ### TODO LLVM ERROR COMES FROM HERE
                     results = self.learn_on_batch(
                         tf.convert_to_tensor(current_task_idx), batch, episodic_batch
                     )
diff --git a/continualworld/tasks.py b/continualworld/tasks.py
index 9a34208..75a1973 100644
--- a/continualworld/tasks.py
+++ b/continualworld/tasks.py
@@ -1,15 +1,15 @@
 TASK_SEQS = {
     "CW10": [
-        "hammer-v1",
-        "push-wall-v1",
-        "faucet-close-v1",
-        "push-back-v1",
-        "stick-pull-v1",
-        "handle-press-side-v1",
-        "push-v1",
-        "shelf-place-v1",
-        "window-close-v1",
-        "peg-unplug-side-v1",
+    "hammer-v2",
+    "push-wall-v2",
+    "faucet-close-v2",
+    "push-back-v2",
+    "stick-pull-v2",
+    "handle-press-side-v2",
+    "push-v2",
+    "shelf-place-v2",
+    "window-close-v2",
+    "peg-unplug-side-v2",
     ],
 }
 
diff --git a/continualworld/utils/utils.py b/continualworld/utils/utils.py
index 4236037..8dc3d73 100644
--- a/continualworld/utils/utils.py
+++ b/continualworld/utils/utils.py
@@ -4,7 +4,7 @@
 from datetime import datetime
 from typing import Callable, Dict, Optional, Type, Union
 
-import gym
+import gymnasium as gym
 import numpy as np
 import tensorflow as tf
 
diff --git a/continualworld/utils/wrappers.py b/continualworld/utils/wrappers.py
index ca57731..27bef17 100644
--- a/continualworld/utils/wrappers.py
+++ b/continualworld/utils/wrappers.py
@@ -1,10 +1,10 @@
 import random
 from typing import Any, Dict, List, Tuple
 
-import gym
+import gymnasium as gym
 import metaworld
 import numpy as np
-from gym.spaces import Box
+from gymnasium.spaces import Box
 
 
 class SuccessCounter(gym.Wrapper):
@@ -16,12 +16,13 @@ def __init__(self, env: gym.Env) -> None:
         self.current_success = False
 
     def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict]:
-        obs, reward, done, info = self.env.step(action)
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        done = np.logical_or(terminated,truncated)
         if info.get("success", False):
             self.current_success = True
         if done:
             self.successes.append(self.current_success)
-        return obs, reward, done, info
+        return obs, reward, terminated, truncated, info
 
     def pop_successes(self) -> List[bool]:
         res = self.successes
@@ -61,11 +62,12 @@ def _append_one_hot(self, obs: np.ndarray) -> np.ndarray:
         return np.concatenate([obs, self.to_append])
 
     def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict]:
-        obs, reward, done, info = self.env.step(action)
-        return self._append_one_hot(obs), reward, done, info
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        return self._append_one_hot(obs), reward, terminated, truncated, info
 
     def reset(self, **kwargs) -> np.ndarray:
-        return self._append_one_hot(self.env.reset(**kwargs))
+        obs, info = self.env.reset(**kwargs)
+        return self._append_one_hot(obs), info
 
 
 class RandomizationWrapper(gym.Wrapper):
diff --git a/input_args.py b/input_args.py
index 24b4a3f..66685c2 100644
--- a/input_args.py
+++ b/input_args.py
@@ -27,7 +27,7 @@ def cl_parse_args(args=None):
         type=str,
         nargs="+",
         choices=["neptune", "tensorboard", "tsv"],
-        default=["tsv"],
+        default=["tensorboard"], #tsv
         help="Types of logger used.",
     )
     parser.add_argument(
diff --git a/metaworld_play.ipynb b/metaworld_play.ipynb
new file mode 100644
index 0000000..4fd60d6
--- /dev/null
+++ b/metaworld_play.ipynb
@@ -0,0 +1,234 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# a notebook to check that metaworld is working, and a testbed for new RL algos \n",
+    "import metaworld\n",
+    "\n",
+    "SEED = 0  # some seed number here\n",
+    "benchmark = metaworld.ML1('pick-place-v2', seed=SEED)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['assembly-v2', 'basketball-v2', 'bin-picking-v2', 'box-close-v2', 'button-press-topdown-v2', 'button-press-topdown-wall-v2', 'button-press-v2', 'button-press-wall-v2', 'coffee-button-v2', 'coffee-pull-v2', 'coffee-push-v2', 'dial-turn-v2', 'disassemble-v2', 'door-close-v2', 'door-lock-v2', 'door-open-v2', 'door-unlock-v2', 'hand-insert-v2', 'drawer-close-v2', 'drawer-open-v2', 'faucet-open-v2', 'faucet-close-v2', 'hammer-v2', 'handle-press-side-v2', 'handle-press-v2', 'handle-pull-side-v2', 'handle-pull-v2', 'lever-pull-v2', 'peg-insert-side-v2', 'pick-place-wall-v2', 'pick-out-of-hole-v2', 'reach-v2', 'push-back-v2', 'push-v2', 'pick-place-v2', 'plate-slide-v2', 'plate-slide-side-v2', 'plate-slide-back-v2', 'plate-slide-back-side-v2', 'peg-unplug-side-v2', 'soccer-v2', 'stick-push-v2', 'stick-pull-v2', 'push-wall-v2', 'reach-wall-v2', 'shelf-place-v2', 'sweep-into-v2', 'sweep-v2', 'window-open-v2', 'window-close-v2']\n"
+     ]
+    }
+   ],
+   "source": [
+    "import random\n",
+    "\n",
+    "print(metaworld.ML1.ENV_NAMES)  # Check out the available environments\n",
+    "\n",
+    "ml1 = metaworld.ML1('pick-place-v2') # Construct the benchmark, sampling tasks\n",
+    "\n",
+    "env = ml1.train_classes['pick-place-v2']()  # Create an environment with task `pick_place`\n",
+    "task = random.choice(ml1.train_tasks)\n",
+    "env.set_task(task)  # Set task\n",
+    "\n",
+    "obs, info = env.reset()  # Reset environment\n",
+    "a = env.action_space.sample()  # Sample an action\n",
+    "obs, reward, terminal, truncated, info = env.step(a)  # Step the environment with the sampled random action\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ml10 = metaworld.ML10() # Construct the benchmark, sampling tasks\n",
+    "\n",
+    "training_envs = []\n",
+    "for name, env_cls in ml10.train_classes.items():\n",
+    "  env = env_cls()\n",
+    "  task = random.choice([task for task in ml10.train_tasks\n",
+    "                        if task.env_name == name])\n",
+    "  env.set_task(task)\n",
+    "  training_envs.append(env)\n",
+    "\n",
+    "for env in training_envs:\n",
+    "  obs, info = env.reset()  # Reset environment\n",
+    "  a = env.action_space.sample()  # Sample an action\n",
+    "  obs, reward, terminated, truncated, info = env.step(a)  # Step the environment with the sampled random action"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# now lets try continual_world\n",
+    "# run single\n",
+    "\n",
+    "from continualworld.envs import get_single_env, get_cl_env\n",
+    "from continualworld.utils.utils import get_activation_from_str\n",
+    "from continualworld.sac.models import MlpActor\n",
+    "from continualworld.utils.run_utils import get_sac_class\n",
+    "\n",
+    "steps_per_task = 1000000\n",
+    "activation=\"lrelu\"\n",
+    "num_tasks=1\n",
+    "hidden_sizes=[256, 256, 256, 256]\n",
+    "tasks=['basketball-v2']\n",
+    "# logger = EpochLogger(args[\"logger_output\"], config=args, group_id=args[\"group_id\"])\n",
+    "\n",
+    "logger = None\n",
+    "steps = steps_per_task * len(tasks)\n",
+    "num_heads = num_tasks\n",
+    "train_env = get_cl_env(tasks, steps_per_task)\n",
+    "test_envs = [\n",
+    "    get_single_env(task, one_hot_idx=i, one_hot_len=num_tasks) for i, task in enumerate(tasks)\n",
+    "]\n",
+    "\n",
+    "actor_kwargs = dict(\n",
+    "    hidden_sizes=hidden_sizes,\n",
+    "    activation=get_activation_from_str(activation),\n",
+    "    use_layer_norm=True,\n",
+    "    num_heads=num_heads,\n",
+    "    hide_task_id=True,\n",
+    ")\n",
+    "critic_kwargs = dict(\n",
+    "    hidden_sizes=hidden_sizes,\n",
+    "    activation=get_activation_from_str(activation),\n",
+    "    use_layer_norm=True,\n",
+    "    num_heads=num_heads,\n",
+    "    hide_task_id=True,\n",
+    ")\n",
+    "\n",
+    "actor_cl = MlpActor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from continualworld.utils.enums import BufferType\n",
+    "lr=0.0025\n",
+    "alpha='auto'\n",
+    "buffer_type=\"fifo\"\n",
+    "gamma=0.99\n",
+    "seed=13\n",
+    "clipnorm=None\n",
+    "target_output_std = 0.089\n",
+    "agent_policy_exploration = False\n",
+    "cl_reg_coef=0.0  # strength of MAS/EWC\n",
+    "regularize_critic = False\n",
+    "\n",
+    "vanilla_sac_kwargs = {\n",
+    "    \"env\": train_env,\n",
+    "    \"test_envs\": test_envs,\n",
+    "    \"logger\": logger,\n",
+    "    \"seed\": seed,\n",
+    "    \"steps\": steps,\n",
+    "    \"log_every\": 100,\n",
+    "    \"replay_size\": 10000,\n",
+    "    \"batch_size\": 10,\n",
+    "    \"actor_cl\": actor_cl,\n",
+    "    \"actor_kwargs\": actor_kwargs,\n",
+    "    \"critic_kwargs\": critic_kwargs,\n",
+    "    \"buffer_type\": BufferType(buffer_type),\n",
+    "    \"reset_buffer_on_task_change\": True,\n",
+    "    \"reset_optimizer_on_task_change\": True,\n",
+    "    \"lr\": lr,\n",
+    "    \"alpha\": alpha,\n",
+    "    \"reset_critic_on_task_change\": True,\n",
+    "    \"clipnorm\": clipnorm,\n",
+    "    \"gamma\": gamma,\n",
+    "    \"target_output_std\": target_output_std,\n",
+    "    \"agent_policy_exploration\": agent_policy_exploration,\n",
+    "}\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cl_method = None  # 'mas'\n",
+    "\n",
+    "sac_class = get_sac_class(cl_method)\n",
+    "if cl_method is None:\n",
+    "    sac = sac_class(**vanilla_sac_kwargs)\n",
+    "elif cl_method in [\"l2\", \"ewc\", \"mas\"]:\n",
+    "    sac = sac_class(\n",
+    "        **vanilla_sac_kwargs, \n",
+    "        cl_reg_coef=cl_reg_coef, \n",
+    "        regularize_critic=regularize_critic\n",
+    "    )\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "ValueError",
+     "evalue": "setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (2,) + inhomogeneous part.",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[22], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43msac\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/code/continual_world/continualworld/sac/sac.py:535\u001b[0m, in \u001b[0;36mSAC.run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    533\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"A method to run the SAC training, after the object has been created.\"\"\"\u001b[39;00m\n\u001b[1;32m    534\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstart_time \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[0;32m--> 535\u001b[0m obs, episode_return, episode_len \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreset\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m, \u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m    537\u001b[0m \u001b[38;5;66;03m# Main loop: collect experience in env and update/log each epoch\u001b[39;00m\n\u001b[1;32m    538\u001b[0m current_task_timestep \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n",
+      "File \u001b[0;32m~/code/continual_world/continualworld/envs.py:145\u001b[0m, in \u001b[0;36mContinualLearningEnv.reset\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    143\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mreset\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m np\u001b[38;5;241m.\u001b[39mndarray:\n\u001b[1;32m    144\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_steps_bound()\n\u001b[0;32m--> 145\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menvs\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcur_seq_idx\u001b[49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreset\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/code/continual_world/continualworld/utils/wrappers.py:33\u001b[0m, in \u001b[0;36mSuccessCounter.reset\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mreset\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m np\u001b[38;5;241m.\u001b[39mndarray:\n\u001b[1;32m     32\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcurrent_success \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[0;32m---> 33\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreset\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/gym/wrappers/time_limit.py:68\u001b[0m, in \u001b[0;36mTimeLimit.reset\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m     59\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Resets the environment with :param:`**kwargs` and sets the number of steps elapsed to zero.\u001b[39;00m\n\u001b[1;32m     60\u001b[0m \n\u001b[1;32m     61\u001b[0m \u001b[38;5;124;03mArgs:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     65\u001b[0m \u001b[38;5;124;03m    The reset environment\u001b[39;00m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     67\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_elapsed_steps \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[0;32m---> 68\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreset\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/code/continual_world/continualworld/utils/wrappers.py:68\u001b[0m, in \u001b[0;36mOneHotAdder.reset\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m     67\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mreset\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m np\u001b[38;5;241m.\u001b[39mndarray:\n\u001b[0;32m---> 68\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_append_one_hot\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreset\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/code/continual_world/continualworld/utils/wrappers.py:61\u001b[0m, in \u001b[0;36mOneHotAdder._append_one_hot\u001b[0;34m(self, obs)\u001b[0m\n\u001b[1;32m     59\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39morig_one_hot_dim \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m     60\u001b[0m     obs \u001b[38;5;241m=\u001b[39m obs[: \u001b[38;5;241m-\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39morig_one_hot_dim]\n\u001b[0;32m---> 61\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconcatenate\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[43mobs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto_append\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m<__array_function__ internals>:200\u001b[0m, in \u001b[0;36mconcatenate\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
+      "\u001b[0;31mValueError\u001b[0m: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (2,) + inhomogeneous part."
+     ]
+    }
+   ],
+   "source": [
+    "sac.run()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "continual_world",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/run_cl.py b/run_cl.py
index 801e9d3..609979a 100644
--- a/run_cl.py
+++ b/run_cl.py
@@ -43,7 +43,8 @@ def main(
     clipnorm: float,
     agent_policy_exploration: bool,
 ):
-    assert (tasks is None) != (task_list is None)
+    # assert (tasks is None) != (task_list is None)
+    tasks='CW10'  # TODO unhardcode this
     if tasks is not None:
         tasks = TASK_SEQS[tasks]
     else:
diff --git a/sb3_buffers.py b/sb3_buffers.py
new file mode 100644
index 0000000..e69de29