helloRL 1.0__py3-none-any.whl

helloRL/__init__.py

@@ -0,0 +1,7 @@
+from .modules import *
+from .utils import *
+
+from . import trainer
+from . import modal_training
+
+__version__ = "1.0.0"

helloRL/modal_training.py

@@ -0,0 +1,98 @@
+import modal
+import time
+from functools import partial
+
+from . import trainer
+from .utils import plot
+
+def gather_modal_results_for_calls(calls, n_timesteps, n_sessions, progress_dict):
+    # Monitor progress by polling the Dict
+    total_timesteps = n_sessions * n_timesteps
+
+    # this file is imported to modal, but progress bar shouldn't be used there
+    from .utils.progress import RemoteProgressBar
+    
+    with RemoteProgressBar("Training", n_steps=total_timesteps, n_sessions=n_sessions) as bar:
+        while True:
+            time.sleep(0.2)  # Poll every 0.2 seconds
+            
+            # Sum up completed timesteps across all sessions
+            completed_timesteps = 0
+            completed_sessions = 0
+            
+            for i in range(n_sessions):
+                value = progress_dict.get(i)
+
+                if value is None:
+                    continue
+
+                completed_timesteps += value
+
+                if value == n_timesteps:
+                    completed_sessions += 1
+
+            bar.update_completed_sessions(completed_sessions)
+            bar.update_value(completed_timesteps)
+
+            if completed_sessions >= n_sessions:
+                break
+    
+    # Gather results
+    results = modal.FunctionCall.gather(*calls)
+
+    return results
+
+def train_session_on_modal_with_func(train_func, session_id, n_timesteps, progress_dict):
+    def progress_callback(current_timestep):
+        progress_dict[session_id] = current_timestep
+
+    train_func_return = train_func(progress_callback=progress_callback)
+    progress_callback(n_timesteps)
+
+    return train_func_return
+
+def create_modal_train_function(app, image, timeout=3600):
+    @app.function(image=image, timeout=timeout, serialized=True)
+    def _modal_train(n_timesteps, setup_func, session_id=None, progress_dict=None):
+        agent, env_name, continuous, params = setup_func()
+
+        training_func = partial(trainer.train, agent, env_name, continuous=continuous, 
+                               n_timesteps=n_timesteps, should_print=False)
+        
+        train_results = train_session_on_modal_with_func(training_func, session_id, 
+                                                          n_timesteps, progress_dict)
+        
+        return (*train_results, agent, env_name, continuous, params)
+    
+    return _modal_train
+
+def train(n_sessions, n_timesteps, setup_func, app, image, timeout=3600):
+    # Create the modal function before entering app.run() so it gets registered
+    modal_train = create_modal_train_function(app, image, timeout)
+
+    with app.run():
+        progress_dict = modal.Dict.from_name("training-progress", create_if_missing=True)
+        progress_dict.clear()
+
+        calls = [modal_train.spawn(n_timesteps, setup_func=setup_func, session_id=i, progress_dict=progress_dict
+        ) for i in range(n_sessions)]
+
+        results = gather_modal_results_for_calls(calls, n_timesteps, n_sessions, progress_dict)
+
+    return results
+
+def plot_results(results, title, n_timesteps, nb_name=None, save_dir=None):
+    plot_results = [(returns, lengths) for returns, lengths, _, _, _, _ in results]
+    _, _, agent, env_name, continuous, params = results[0]
+
+    plot.plot_sessions(
+        plot_results,
+        title,
+        env_name=env_name,
+        continuous=continuous,
+        params=params,
+        agent=agent,
+        n_timesteps=n_timesteps,
+        nb_name=nb_name,
+        save_dir=save_dir
+    )

helloRL/modules/__init__.py

@@ -0,0 +1,17 @@
+from .actors import *
+from .critics import *
+from .agents import *
+from .params import *
+from .foundation import *
+
+from .a2c import *
+from .critic_loss_clipped import *
+from .critic_loss_q import *
+from .epochs import *
+from .gae import *
+from .grad_norm import *
+from .lr_anneal import *
+from .monte_carlo import *
+from .po_clipped import *
+from .replay import *
+from .rollout_data import *

helloRL/modules/a2c.py

@@ -0,0 +1,102 @@
+from dataclasses import dataclass
+import torch
+
+from .foundation import *
+    
+@dataclass
+class AdvantageTransformNormalize(AdvantageTransform):
+    # mean or sum normalization
+    method: str = "mean"
+
+    def transform(self, raw_advantages: torch.Tensor) -> torch.Tensor:
+        if self.method == "sum":
+            return raw_advantages / (raw_advantages.sum() + 1e-9)
+        else:  # default to mean normalization
+            if len(raw_advantages) < 2:
+                return raw_advantages
+
+            return (raw_advantages - raw_advantages.mean()) / (raw_advantages.std() + 1e-9)
+
+@dataclass
+class RolloutMethodA2C(RolloutMethod):
+    n_steps: int = 16
+    n_envs: int = 4
+
+    def collect_rollout_data(
+        self, envs: gym.vector.VectorEnv, initial_states: torch.Tensor, agent: AgentProtocol, tracker: SessionTracker
+        ) -> tuple[RolloutData, list[np.ndarray]]:
+            n_envs = envs.num_envs
+            state_space = envs.single_observation_space.shape[0]
+            action_dim = envs.single_action_space.shape[0] if isinstance(envs.single_action_space, gym.spaces.Box) else 1
+
+            rollout_states_t = torch.zeros(n_envs, self.n_steps, state_space) # (n_envs, n_steps, state_space)
+            rollout_actions_t = torch.zeros(n_envs, self.n_steps, action_dim) # (n_envs, n_steps, action_dim)
+            rollout_next_states_t = torch.zeros(n_envs, self.n_steps, state_space) # (n_envs, n_steps, state_space)
+            rollout_rewards_t = torch.zeros(n_envs, self.n_steps, 1) # (n_envs, n_steps, 1 value)
+            rollout_terminateds_t = torch.zeros(n_envs, self.n_steps, 1) # (n_envs, n_steps, 1 value)
+            rollout_truncateds_t = torch.zeros(n_envs, self.n_steps, 1) # (n_envs, n_steps, 1 value)
+            rollout_dones_t = torch.zeros(n_envs, self.n_steps, 1) # (n_envs, n_steps, 1 value)
+            rollout_critic_values_t = torch.zeros(n_envs, self.n_steps, 1) # (n_envs, n_steps, 1 value)
+            rollout_log_probs_t = torch.zeros(n_envs, self.n_steps, 1) # (n_envs, n_steps, 1 value)
+
+            states = initial_states
+
+            for step in range(self.n_steps):
+                tracker.increment_timestep(n=n_envs)
+
+                states_t = torch.tensor(states).float() # (n_envs, state_space)
+
+                with torch.no_grad():
+                    actions_t, _ = agent.actor.output(states_t)
+                    # apply exploration noise etc to the actions
+                    actions_t = agent.actor.exploration(actions_t)
+                    log_probs_t, _ = agent.actor.get_log_prob_and_entropy(states_t, actions_t)
+                    critic_values_t = agent.get_critic_value(states_t, actions_t)
+                actions_np = actions_t.squeeze(-1).numpy() # (n_envs)
+                next_states, rewards, terminateds, truncateds, infos = envs.step(actions_np)
+
+                dones = terminateds | truncateds
+                dones_t = torch.tensor(dones).float().reshape(n_envs, 1)  # (n_envs, 1 value)
+                terminateds_t = torch.tensor(terminateds).float().reshape(n_envs, 1)  # (n_envs, 1 value)
+                truncateds_t = torch.tensor(truncateds).float().reshape(n_envs, 1)  # (n_envs, 1 value)
+
+                rewards_t = torch.tensor(rewards).float().reshape(n_envs, 1)  # (n_envs, 1 value)
+                
+                rollout_states_t[:, step] = states_t  # (n_envs, n_steps, state_space)
+                rollout_actions_t[:, step] = actions_t  # (n_envs, n_steps, 1 value)
+                rollout_next_states_t[:, step] = torch.tensor(next_states).float()  # (n_envs, n_steps, state_space)
+                rollout_rewards_t[:, step] = rewards_t  # (n_envs, n_steps, 1 value)
+                rollout_terminateds_t[:, step] = terminateds_t  # (n_envs, n_steps, 1 value)
+                rollout_truncateds_t[:, step] = truncateds_t  # (n_envs, n_steps, 1 value)
+                rollout_dones_t[:, step] = dones_t  # (n_envs, n_steps, 1 value)
+                rollout_critic_values_t[:, step] = critic_values_t  # (n_envs, n_steps, 1 value)
+                rollout_log_probs_t[:, step] = log_probs_t  # (n_envs, n_steps, 1 value)
+
+                if dones.any():
+                    episode_returns = infos['episode']['r'][dones]
+                    episode_lengths = infos['episode']['l'][dones]
+
+                    tracker.finish_episodes(episode_returns, episode_lengths)
+                    
+                    next_states, _ = envs.reset(options={'reset_mask': dones})
+
+                states = next_states
+            
+            rollout_returns_t = torch.zeros_like(rollout_rewards_t)  # (n_envs, n_steps, 1)
+            rollout_advantages_t = torch.zeros_like(rollout_rewards_t)  # (n_envs, n_steps, 1)
+
+            rollout_data = RolloutData(
+                states=rollout_states_t,
+                actions=rollout_actions_t,
+                next_states=rollout_next_states_t,
+                rewards=rollout_rewards_t,
+                terminateds=rollout_terminateds_t,
+                truncateds=rollout_truncateds_t,
+                dones=rollout_dones_t,
+                critic_values=rollout_critic_values_t,
+                log_probs=rollout_log_probs_t,
+                returns=rollout_returns_t,
+                advantages=rollout_advantages_t
+            )
+
+            return rollout_data, next_states

helloRL/modules/actors.py

@@ -0,0 +1,256 @@
+from collections.abc import Callable
+import torch
+import torch.nn as nn
+from torch.distributions import Categorical, Normal
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+
+HIDDEN_SIZES_DEFAULT = [64, 64]
+
+class ActorProtocol(ABC):
+    @abstractmethod
+    def output(self, state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Given a state, output an action and its log probability.
+        
+        Args:
+            state: The input state (torch.Tensor)
+
+        Returns:
+            action: The action to take (torch.Tensor)
+            log_prob: The log probability of the action (torch.Tensor)
+        """
+        pass
+
+    @abstractmethod
+    def get_loss(self, data, critic_value_func: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]) -> torch.Tensor:
+        """Compute the actor loss given the data and critic.
+        Advantages should be precomputed and stored in data.
+
+        Args:
+            data: The rollout data (RolloutData)
+            critic_value_func: A function that takes state and action tensors and returns critic values
+            (Callable[[torch.Tensor, torch.Tensor], torch.Tensor])
+        """
+        pass
+
+    @abstractmethod
+    def exploration(self, action) -> torch.Tensor:
+        """Apply exploration noise to the given action.
+
+        Args:
+            action: The action to apply exploration to (torch.Tensor)
+
+        Returns:
+            explored_action: The action after applying exploration (torch.Tensor)
+        """
+        pass
+
+    @abstractmethod
+    def get_log_prob_and_entropy(self, state, action) -> tuple[torch.Tensor, torch.Tensor]:
+        """Given a state and action, output its log probability and entropy.
+        
+        Args:
+            state: The input state (torch.Tensor)
+            action: The input action (torch.Tensor)
+
+        Returns:
+            log_prob: The log probability of the action (torch.Tensor)
+            entropy: The entropy of the action distribution (torch.Tensor)
+        """
+        pass
+
+@dataclass
+class ActorParams:
+    pass
+
+@dataclass
+class PolicyObjectiveMethod(ABC):
+    @abstractmethod
+    def compute_policy_objective(self, ratio: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
+        """Compute the policy objective given the ratio of new to old policy probabilities and the advantages.
+        
+        Args:
+            ratio: Ratio of new to old policy probabilities (torch.Tensor)
+            advantages: Advantages computed from the advantage method (torch.Tensor)
+            
+        Returns:
+            Policy objective (torch.Tensor)
+        """
+        pass
+
+class PolicyObjectiveMethodStandard(PolicyObjectiveMethod):
+    def compute_policy_objective(self, ratio: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
+        return ratio * advantages
+
+@dataclass
+class DistributionalActorParams(ActorParams):
+    policy_objective_method: PolicyObjectiveMethod = field(
+        default_factory=PolicyObjectiveMethodStandard
+    )
+    entropy_coef: float = 0.0
+
+class DiscreteActorNetwork(nn.Module):
+    def __init__(self, state_dim, action_dim, hidden_sizes=HIDDEN_SIZES_DEFAULT):
+        super(DiscreteActorNetwork, self).__init__()
+
+        self.hidden_layers = nn.ModuleList()
+
+        prev_size = state_dim
+
+        for size in hidden_sizes:
+            layer = nn.Linear(prev_size, size, dtype=torch.float32)
+            self.hidden_layers.append(layer)
+            prev_size = size
+
+        self.head = nn.Linear(prev_size, action_dim, dtype=torch.float32)
+
+    def forward(self, state): # logits: (batch_size, action_space)
+        x = state
+
+        for layer in self.hidden_layers:
+            x = torch.relu(layer(x))
+
+        logits = self.head(x)
+        return logits
+
+class DistributionalActor(ActorProtocol, nn.Module):
+    def __init__(self, params=DistributionalActorParams()):
+        super(DistributionalActor, self).__init__()
+        self.params = params
+
+    def get_log_prob_and_entropy(self, state, action):
+        pass
+
+    def forward(self, state):
+        return self.network(state)
+
+    def get_loss(self, data, critic_value_func):
+        old_log_probs = data.log_probs.detach()  # stored at rollout time
+
+        new_log_probs, entropies = self.get_log_prob_and_entropy(data.states, data.actions)
+        log_ratio = new_log_probs - old_log_probs
+        ratio = torch.exp(log_ratio)
+
+        policy_objective = self.params.policy_objective_method.compute_policy_objective(ratio, data.advantages)
+        actor_loss = -(policy_objective.sum())
+        actor_loss -= (self.params.entropy_coef * entropies.sum())
+        
+        return actor_loss
+    
+    def exploration(self, action) -> torch.Tensor:
+        return action
+
+class DiscreteActor(DistributionalActor):
+    def __init__(self, state_dim, action_dim, hidden_sizes=HIDDEN_SIZES_DEFAULT, params=DistributionalActorParams()):
+        super(DiscreteActor, self).__init__(params=params)
+
+        self.network = DiscreteActorNetwork(state_dim, action_dim, hidden_sizes=hidden_sizes)
+
+    def output(self, state): # action: (batch_size, 1), log_prob: (batch_size, 1)
+        actor_logits = self.network.forward(state)
+        action_distribution = Categorical(logits=actor_logits)
+        action = action_distribution.sample().unsqueeze(-1) # shape: (batch_size, 1)
+        log_prob = action_distribution.log_prob(action.squeeze()).unsqueeze(-1) # shape: (batch_size, 1)
+        
+        return action, log_prob
+
+    def get_log_prob_and_entropy(self, state, action):
+        actor_logits = self.network.forward(state)
+        action_distribution = Categorical(logits=actor_logits)
+        log_prob = action_distribution.log_prob(action.squeeze(-1)).unsqueeze(-1)
+        entropy = action_distribution.entropy().unsqueeze(-1)  # (batch_size, 1)
+        return log_prob, entropy
+    
+class ContinuousActorNetwork(nn.Module):
+    def __init__(self, state_dim, action_dim, action_range, hidden_sizes=HIDDEN_SIZES_DEFAULT):
+        super(ContinuousActorNetwork, self).__init__()
+
+        self.hidden_layers = nn.ModuleList()
+
+        prev_size = state_dim
+
+        for size in hidden_sizes:
+            layer = nn.Linear(prev_size, size, dtype=torch.float32)
+            self.hidden_layers.append(layer)
+            prev_size = size
+
+        self.head = nn.Linear(prev_size, action_dim, dtype=torch.float32)
+        self.action_range = action_range
+
+    def forward(self, state): # logits: (batch_size, action_space)
+        x = state
+
+        for layer in self.hidden_layers:
+            x = torch.relu(layer(x))
+
+        x = torch.tanh(self.head(x))  # Bounds to -1 to 1
+
+        mins = self.action_range.min(dim=0).values
+        maxs = self.action_range.max(dim=0).values
+        length = maxs - mins
+        x = (x * (length / 2)) + ((mins + maxs) / 2)
+
+        return x
+    
+class StochasticActor(DistributionalActor):
+    def __init__(self, state_dim, action_dim, action_range, hidden_sizes=HIDDEN_SIZES_DEFAULT, params=DistributionalActorParams()):
+        super(StochasticActor, self).__init__(params=params)
+        self.network = ContinuousActorNetwork(state_dim, action_dim, action_range, hidden_sizes=hidden_sizes)
+        self.log_std = nn.Parameter(torch.zeros(action_dim))
+
+    def output(self, state): # action: (batch_size, 1), log_prob: (batch_size, 1)
+        # mean, also referred to as 'mu'
+        mean = self.network(state)
+        std = torch.exp(self.log_std)  # Exp to make positive
+        dist = Normal(mean, std)
+        action = dist.sample()  # Shape: (batch_size, action_dim)
+        log_prob = dist.log_prob(action).sum(dim=-1, keepdim=True)  # Sum for multi-dim, shape: (batch_size, 1)
+        
+        return action, log_prob
+
+    def get_log_prob_and_entropy(self, state, action):
+        mean = self.network(state)
+        std = torch.exp(self.log_std)  # Exp to make positive
+        dist = Normal(mean, std)
+        log_prob = dist.log_prob(action).sum(dim=-1, keepdim=True)  # (batch_size, 1)
+        entropy = dist.entropy().sum(dim=-1, keepdim=True)  # (batch_size, 1)
+        return log_prob, entropy
+
+@dataclass
+class DeterministicActorParams(ActorParams):
+    exploration_std: float = 0.1
+    
+class DeterministicActor(ActorProtocol, nn.Module):
+    def __init__(self, state_dim, action_dim, action_range, hidden_sizes=HIDDEN_SIZES_DEFAULT, params=DeterministicActorParams()):
+        super(DeterministicActor, self).__init__()
+        self.network = ContinuousActorNetwork(state_dim, action_dim, action_range, hidden_sizes=hidden_sizes)
+        self.params = params
+
+    def forward(self, state): # action: (batch_size, 1)
+        return self.network(state)
+    
+    def output(self, state): # action: (batch_size, 1), log_prob: (batch_size, 1)
+        # return zeroes instead of logprobs
+        return self.network(state), torch.zeros((state.shape[0], 1), dtype=torch.float32)
+    
+    def get_loss(self, data, critic_value_func):
+        actions_pi = self.network(data.states)
+        actor_loss = -(critic_value_func(data.states, actions_pi).mean())
+        
+        return actor_loss
+    
+    def exploration(self, action) -> torch.Tensor:
+        action = action + (torch.randn_like(action) * self.params.exploration_std)
+
+        mins = self.network.action_range.min(dim=0).values
+        maxs = self.network.action_range.max(dim=0).values
+        action = torch.clamp(action, mins, maxs)
+
+        return action
+    
+    def get_log_prob_and_entropy(self, state, action):
+        # return zeroes instead of logprobs and entropies
+        batch_size = state.shape[0]
+        log_prob = torch.zeros((batch_size, 1), dtype=torch.float32)
+        entropy = torch.zeros((batch_size, 1), dtype=torch.float32)
+        return log_prob, entropy

helloRL/modules/agents.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from copy import deepcopy
+
+from .actors import ActorProtocol
+from .critics import CriticProtocol
+
+class AgentProtocol(ABC):
+    actor: ActorProtocol
+    critics: list[CriticProtocol]
+
+    @abstractmethod
+    def get_action(self, state):
+        """Get the action from the actor for the given state.
+
+        Args:
+            state: The input state (torch.Tensor)
+        """
+        pass
+
+    def get_critic_value(self, state, action):
+        """Get the critic value for the given state and action.
+
+        Args:
+            state: The input state (torch.Tensor)
+            action: The input action (torch.Tensor)
+        """
+        pass
+
+    @abstractmethod
+    def get_target_action(self, state):
+        """Get the target actor's output for the given state.
+        A simple implementation could use the main actor if no target networks are used.
+
+        Args:
+            state: The input state (torch.Tensor)
+        """
+        pass
+
+    @abstractmethod
+    def get_target_critic_value(self, state, action):
+        """Get the target critic's output for the given state and action.
+        A simple implementation could use the main critics if no target networks are used.
+
+        Args:
+            state: The input state (torch.Tensor)
+            action: The input action (torch.Tensor)
+        """
+        pass
+
+    @abstractmethod
+    def update_targets(self):
+        """Update the target networks, if applicable."""
+        pass
+
+@dataclass
+class AgentParams:
+    pass
+
+class Agent(AgentProtocol):
+    def __init__(self, actor: ActorProtocol, critics: list[CriticProtocol], params: AgentParams=AgentParams()):
+
+        super(Agent, self).__init__()
+
+        self.actor = actor
+        self.critics = critics
+        self.params = params
+
+    def get_action(self, state):
+        action, _ = self.actor.output(state)
+
+        return action
+    
+    def get_critic_value(self, state, action):
+        # Return the minimum value across all critics
+        values = [critic.output(state, action) for critic in self.critics]
+        return torch.min(torch.stack(values), dim=0).values
+    
+    def get_target_action(self, state):
+        return self.get_action(state)
+    
+    def get_target_critic_value(self, state, action):
+        return self.get_critic_value(state, action)
+    
+    def update_targets(self):
+        # No target networks to update
+        pass
+
+@dataclass
+class AgentWithTargetsParams(AgentParams):
+    tau: float = 0.1
+
+class AgentWithTargets(Agent):
+    def __init__(self, actor: ActorProtocol, critics: list[CriticProtocol], params: AgentWithTargetsParams):
+        super(AgentWithTargets, self).__init__(actor, critics, params)
+
+        self.target_actor_network = deepcopy(self.actor.network)
+        self.target_critic_networks = [deepcopy(critic.network) for critic in self.critics]
+
+    def get_target_action(self, state):
+        return self.target_actor_network(state)
+    
+    def get_target_critic_value(self, state, action):
+        # Return the minimum value across all target critics
+        values = [target_critic_network(state, action) for target_critic_network in self.target_critic_networks]
+        return torch.min(torch.stack(values), dim=0).values
+    
+    def update_targets(self):
+        # Soft update targets
+        for target_param, param in zip(self.target_actor_network.parameters(), self.actor.network.parameters()):
+            target_param.data.copy_((self.params.tau * param.data) + ((1 - self.params.tau) * target_param.data))
+
+        for i, critic in enumerate(self.critics):
+            for target_param, param in zip(self.target_critic_networks[i].parameters(), critic.network.parameters()):
+                target_param.data.copy_((self.params.tau * param.data) + ((1 - self.params.tau) * target_param.data))

helloRL/modules/critic_loss_clipped.py

@@ -0,0 +1,20 @@
+import torch
+
+from .foundation import *
+from .critics import CriticLossMethod
+
+class CriticLossMethodClipped(CriticLossMethod):
+    def __init__(self, clip_range_vf: float):
+        self.clip_range_vf = clip_range_vf
+
+    def compute_critic_loss(self, reference_actor, reference_critic, new_values, data, gamma):
+        old_values = data.critic_values.detach()
+
+        value_delta = (new_values - old_values)
+        value_delta_clipped = torch.clamp(value_delta, -self.clip_range_vf, self.clip_range_vf)
+        values_clipped = old_values + value_delta_clipped
+
+        loss_unclipped = (new_values - data.returns).pow(2)
+        loss_clipped = (values_clipped - data.returns).pow(2)
+
+        return torch.maximum(loss_unclipped, loss_clipped).mean()

helloRL/modules/critic_loss_q.py

@@ -0,0 +1,18 @@
+import torch
+
+from .foundation import *
+from .critics import CriticLossMethod
+
+class CriticLossMethodQ(CriticLossMethod):
+    def compute_critic_loss(self, target_action_func, target_critic_value_func, new_values, data, gamma):
+        with torch.no_grad():
+            # this should be the target actor
+            a_next = target_action_func(data.next_states)
+            # this should be the target critic
+            q_next = target_critic_value_func(data.next_states, a_next)
+            y = data.rewards + ((gamma * (1.0 - data.terminateds)) * q_next)
+        
+        # new values come from the current critic, passed into this function
+        critic_loss = torch.nn.MSELoss()(new_values, y)
+
+        return critic_loss