Repository: mlaux1/rl-blox
Pull request #101: MRG: TD3
Author: login: AlexanderFabisch
Created: 2025-05-08T13:34:36Z | Merged: 2025-05-16T11:25:16Z

Fix #6 

* Implements TD3 as an extension of DDPG (+ test, example, apidoc)
* Modifications to DDPG
    * Refactors DDPG: extracts `q_deterministic_bootstrap_estimate`
    * Extends documentation of DDPG
    * Removes policy delay from DDPG
    * Fixes control flow bug: take a training step, even when the episode finishes
    * Report losses at the correct environment step
    * Rename `update_target` to `soft_target_net_update`
* Modifications to SAC
    * Refactors DDPG: extracts `sac_q_target`
    * Fixes control flow bug: take a training step, even when the episode finishes
    * Report losses at the correct environment step
* Modifications to logger (required for comparison of DDPG and TD3 in the Hopper-v5 env)
    * Accept relative dirs for checkpointing with orbax by converting them to absolute dirs
    * Do not delete existing dir for checkpoints (makes more sense with multiple runs)
    * Fix concatenation of checkpoint path
    * Wait until checkpoint is saved before continuing (fixes an error with asynchronous checkpointing in the last step)

## Results on Hopper-v5 after 1e6 steps over 5 seeds with 10 evaluation episodes per run:

* TD3 return 2881.7565 +- 863.0822 (purple learning curve)
* DDPG return 810.2611 +- 716.9389 (green learning curve)
* DDPG return (with learning rates `3e-4` and exploration noise 0.2) 1446.0263 +- 425.4875 (blue learning curve)

![metrics-07_59_24-16-May-25](https://github.com/user-attachments/assets/382b1ee7-2901-4372-bbb2-16fd32cd4cb4)

*For reference,*

* stable-baselines3 [results on PyBullet envs](https://stable-baselines3.readthedocs.io/en/master/modules/ddpg.html#pybullet-environments)
    * DDPG: 1201 +/- 211
    * TD3: 2429 +/- 126
* results from the [TD3 paper](https://arxiv.org/pdf/1802.09477) (Hopper-v1, I used the same configuration as in "Our DDPG"):
    * TD3: 3564.07 ± 114.74
    * Our DDPG: 1860.02

![Screenshot from 2025-05-15 15-37-20](https://github.com/user-attachments/assets/fe0242ed-c9c3-419e-8ec4-9d1ed215f34e)


## Results on Pendulum-v1 after 12_000 steps over 20 seeds with 10 evaluation episodes per run:

* TD3 return -151.1679 +- 48.6921
* DDPG return -135.4195 +- 6.5115

## Example of an agent trained with TD3 on Hopper-v5

https://github.com/user-attachments/assets/36728be4-4f37-4826-bc7e-1f0eb54af652


Linked issue #6: Implement TD3
Papers:
- https://arxiv.org/pdf/1802.09477

Diff:
diff --git a/doc/source/api.rst b/doc/source/api.rst
index b030654d6..40c592cb9 100644
--- a/doc/source/api.rst
+++ b/doc/source/api.rst
@@ -23,6 +23,7 @@ Algorithm Interface
    reinforce.train_reinforce
    actor_critic.train_ac
    ddpg.train_ddpg
+   td3.train_td3
    sac.train_sac
    pets.train_pets
 
@@ -39,15 +40,21 @@ Functional Blox
    reinforce.policy_gradient_pseudo_loss
    reinforce.mse_value_loss
    actor_critic.actor_critic_policy_gradient
+   ddpg.sample_actions
    ddpg.mse_action_value_loss
    ddpg.deterministic_policy_value_loss
-   ddpg.update_target
+   ddpg.soft_target_net_update
    ddpg.ddpg_update_actor
    ddpg.ddpg_update_critic
+   ddpg.q_deterministic_bootstrap_estimate
+   td3.td3_update_critic
+   td3.double_q_deterministic_bootstrap_estimate
+   td3.sample_target_actions
    sac.sac_actor_loss
    sac.sac_exploration_loss
    sac.sac_update_actor
    sac.sac_update_critic
+   sac.sac_q_target
    double_q_learning.double_q_learning
    pets.mpc_action
    pets.ts_inf
diff --git a/examples/td3_example.py b/examples/td3_example.py
new file mode 100644
index 000000000..98c358ddc
--- /dev/null
+++ b/examples/td3_example.py
@@ -0,0 +1,97 @@
+import gymnasium as gym
+import jax.numpy as jnp
+import numpy as np
+
+from rl_blox.algorithm.td3 import create_td3_state, train_td3
+from rl_blox.logging.logger import AIMLogger, LoggerList, StandardLogger
+
+env_name = "Hopper-v5"
+env = gym.make(env_name)
+
+seed = 1
+verbose = 1
+env = gym.wrappers.RecordEpisodeStatistics(env)
+
+hparams_models = dict(
+    policy_hidden_nodes=[256, 256],
+    policy_learning_rate=1e-3,
+    q_hidden_nodes=[256, 256],
+    q_learning_rate=1e-3,
+    seed=seed,
+)
+hparams_algorithm = dict(
+    policy_delay=2,
+    exploration_noise=0.2,
+    noise_clip=0.5,
+    gradient_steps=1,
+    total_timesteps=1_000_000,
+    buffer_size=1_000_000,
+    learning_starts=25_000,
+    batch_size=256,
+    seed=seed,
+)
+
+if verbose:
+    print(
+        "This example uses the AIM logger. You will not see any output on "
+        "stdout. Run 'aim up' to analyze the progress."
+    )
+logger = LoggerList(
+    [AIMLogger(), StandardLogger(checkpoint_dir=".td3_example", verbose=0)]
+)
+logger.define_experiment(
+    env_name=env_name,
+    algorithm_name="TD3",
+    hparams=hparams_models | hparams_algorithm,
+)
+n_policy_epochs = (
+    hparams_algorithm["total_timesteps"] - hparams_algorithm["learning_starts"]
+) // hparams_algorithm["policy_delay"]
+logger.define_checkpoint_frequency("policy", n_policy_epochs)
+
+td3_state = create_td3_state(env, **hparams_models)
+
+td3_result = train_td3(
+    env,
+    td3_state.policy,
+    td3_state.policy_optimizer,
+    td3_state.q1,
+    td3_state.q1_optimizer,
+    td3_state.q2,
+    td3_state.q2_optimizer,
+    logger=logger,
+    **hparams_algorithm,
+)
+env.close()
+policy, _, _, q1, _, _, q2, _, _ = td3_result
+
+# Evaluation
+env = gym.make(env_name, render_mode="human")
+returns = []
+for i in range(10):
+    done = False
+    infos = {}
+    obs, _ = env.reset(seed=i)
+    accumulated_reward = 0.0
+    while not done:
+        action = np.asarray(policy(jnp.asarray(obs)))
+        next_obs, reward, termination, truncation, infos = env.step(action)
+        accumulated_reward += reward
+        done = termination or truncation
+        if verbose >= 2:
+            q1_value = float(
+                q1(jnp.concatenate((obs, action), axis=-1)).squeeze()
+            )
+            q2_value = float(
+                q2(jnp.concatenate((obs, action), axis=-1)).squeeze()
+            )
+            q_value = min(q1_value, q2_value)
+            print(f"{q_value=:.3f} {q1_value=:.3f} {q2_value=:.3f}")
+        obs = np.asarray(next_obs)
+    returns.append(accumulated_reward)
+env.close()
+print(f"{returns=}")
+print(
+    f"{np.mean(returns)} +- {np.std(returns)}, "
+    f"[min={min(returns)}, max={max(returns)}]"
+)
diff --git a/rl_blox/algorithm/ddpg.py b/rl_blox/algorithm/ddpg.py
index bd4abac1a..d7958aa64 100644
--- a/rl_blox/algorithm/ddpg.py
+++ b/rl_blox/algorithm/ddpg.py
@@ -207,11 +207,9 @@ def ddpg_update_critic(
     chex.assert_equal_shape_prefix((observations, terminations), prefix_len=1)
     chex.assert_equal_shape((rewards, terminations))
 
-    next_actions = policy_target(next_observations)
-    q_target_next = q_target(
-        jnp.concatenate((next_observations, next_actions), axis=-1)
-    ).squeeze()
-    q_bootstrap = rewards + (1 - terminations) * gamma * q_target_next
+    q_bootstrap = q_deterministic_bootstrap_estimate(
+        policy_target, rewards, terminations, gamma, q_target, next_observations
+    )
 
     q_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
         observations, actions, q_bootstrap, q
@@ -221,6 +219,52 @@ def ddpg_update_critic(
     return q_loss_value
 
 
+def q_deterministic_bootstrap_estimate(
+    policy: nnx.Module,
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+    gamma: float,
+    q: nnx.Module,
+    next_observations: jnp.ndarray,
+) -> jnp.ndarray:
+    r"""Bootstrap estimate of action-value function with deterministic policy.
+
+    .. math::
+
+        \mathbb{E}\left[R(o_t)\right]
+        \approx
+        r_{t+1} + \gamma Q(o_{t+1}, \pi(o_{t+1}))
+
+    Parameters
+    ----------
+    policy : nnx.Module
+        Deterministic policy for action selection.
+
+    rewards : array
+        Observed reward.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    gamma : float
+        Discount factor.
+
+    q : nnx.Module
+        Action-value function.
+
+    next_observations : array
+        Next observations.
+
+    Returns
+    -------
+    q_bootstrap : array
+        Bootstrap estimate of action-value function.
+    """
+    next_actions = policy(next_observations)
+    obs_act = jnp.concatenate((next_observations, next_actions), axis=-1)
+    return rewards + (1 - terminations) * gamma * q(obs_act).squeeze()
+
+
 @nnx.jit
 def ddpg_update_actor(
     policy: nnx.Module,
@@ -243,20 +287,34 @@ def ddpg_update_actor(
 
 
 @nnx.jit
-def update_target(net: nnx.Module, target_net: nnx.Module, tau: float) -> None:
-    """Update target network inplace with Polyak averaging.
+def soft_target_net_update(
+    net: nnx.Module, target_net: nnx.Module, tau: float
+) -> None:
+    r"""Inplace soft update for target network with Polyak averaging.
+
+    The soft update for the target network is supposed to be applied with about
+    the same frequency as the update for the live network.
+
+    Update formula:
+
+    .. math::
+
+        \theta' \leftarrow \tau \theta + (1 - \tau) \theta'
+
+    where :math:`\theta` are the weights of the live network and
+    :math:`\theta'` are the weights of the target network.
 
     Parameters
     ----------
     net : nnx.Module
-        Live network.
+        Live network with weights :math:`\theta`.
 
     target_net : nnx.Module
-        Target network.
+        Target network with weights :math:`\theta'`.
 
     tau : float
-        The step_size used to update the Polyak average, i.e., the coefficient
-        with which the live network's parameters will be multiplied.
+        The step size :math:`\tau`, i.e., the coefficient with which the live
+        network's parameters will be multiplied.
     """
     params = nnx.state(net)
     target_params = nnx.state(target_net)
@@ -344,73 +402,107 @@ def train_ddpg(
     gradient_steps: int = 1,
     exploration_noise: float = 0.1,
     learning_starts: int = 25_000,
-    policy_frequency: int = 2,
     policy_target: nnx.Optimizer | None = None,
     q_target: nnx.Optimizer | None = None,
     logger: LoggerBase | None = None,
 ) -> tuple[
     nnx.Module, nnx.Module, nnx.Optimizer, nnx.Module, nnx.Module, nnx.Optimizer
 ]:
-    """Deep Deterministic Policy Gradients (DDPG).
+    r"""Deep Deterministic Policy Gradients (DDPG).
 
     This is an off-policy actor-critic algorithm with a deterministic policy.
     The critic approximates the action-value function. The actor will maximize
     action values based on the approximation of the action-value function.
 
     DDPG [2]_ extends Deterministic Policy Gradients [1]_ to use neural
-    networks as function approximators with target networks.
+    networks as function approximators with target networks for the policy
+    :math:`\pi` and the action-value function :math:`Q`.
+
+    The deterministic target policy :math:`\mu(o) = a` that maps observations
+    to actions has to be passed as the parameter `policy`. The function
+    :func:`sample_actions` adds Gaussian noise to the target policy to create
+    the behavior policy that is used during training for exploration.
+    :func:`ddpg_update_actor` will update the policy with the
+    `policy_optimizer` to minimize the :func:`deterministic_policy_value_loss`.
+
+    The action-value function has to be passed as the parameter `q`. We use
+    :func:`ddpg_update_critic` to update the action-value function with the
+    `q_optimizer` and an `mse_action_value_loss` between the current
+    prediction of `q` and the target value generated by
+    :func:`q_deterministic_bootstrap_estimate` using the target networks of
+    `policy` and `q`. The target networks will be updated with Polyak
+    averaging in :func:`soft_target_net_update`.
 
     Parameters
     ----------
-    env: Vectorized Gymnasium environments.
-    policy: Deterministic policy network.
-    policy_optimizer: Optimizer for the policy network.
-    q: Q network.
-    q_optimizer: Optimizer for the Q network.
-    seed: Seed for random number generators in Jax and NumPy.
-    total_timesteps: Number of steps to execute in the environment.
-    actor_learning_rate: Learning rate of the actor.
-    q_learning_rate: Learning rate of the critic.
-    buffer_size: Size of the replay buffer.
-    gamma: Discount factor.
-    tau
+    env : gymnasium.Env
+        Gymnasium environment.
+
+    policy : nnx.Module
+        Deterministic policy network.
+
+    policy_optimizer : nnx.Optimizer
+        Optimizer for the policy network.
+
+    q : nnx.Module
+        Q network.
+
+    q_optimizer : nnx.Optimizer
+        Optimizer for the Q network.
+
+    seed : int
+        Seed for random number generators in Jax and NumPy.
+
+    total_timesteps : int
+        Number of steps to execute in the environment.
+
+    buffer_size : int
+        Size of the replay buffer.
+
+    gamma : float
+        Discount factor.
+
+    tau : float
         Learning rate for polyak averaging of target policy and value function.
-    batch_size
+
+    batch_size : int
         Size of a batch during gradient computation.
-    gradient_steps
+
+    gradient_steps : int
         Number of gradient steps during one training phase.
-    exploration_noise
+
+    exploration_noise : float
         Exploration noise in action space. Will be scaled by half of the range
         of the action space.
-    learning_starts
+
+    learning_starts : int
         Learning starts after this number of random steps was taken in the
         environment.
-    policy_frequency
-        The policy will only be updated after this number of steps. Target
-        policy and value function will be updated with the same frequency. The
-        value function will be updated after every step.
-    policy_target
+
+    policy_target : nnx.Module
         Target policy. Only has to be set if we want to continue training
         from an old state.
-    q_target
+
+    q_target : nnx.Module
         Target network. Only has to be set if we want to continue training
         from an old state.
+
     logger : LoggerBase, optional
         Experiment logger.
 
     Returns
     -------
-    policy
+    policy : nnx.Module
         Final policy.
-    policy_target
+    policy_target : nnx.Module
         Target policy.
-    policy_optimizer
+    policy_optimizer : nnx.Optimizer
         Policy optimizer.
-    q
+    q : nnx.Module
         Final state-action value function.
-    q_target
+    q_target : nnx.Module
         Target network.
-    q_optimizer
+    q_optimizer : nnx.Optimizer
         Optimizer for Q network.
 
     References
@@ -477,23 +569,7 @@ def train_ddpg(
             termination=termination,
         )
 
-        done = termination or truncated
-        if done:
-            if logger is not None:
-                if "episode" in info:
-                    logger.record_stat(
-                        "return", info["episode"]["r"], step=global_step
-                    )
-                logger.stop_episode(steps_per_episode)
-                logger.start_new_episode()
-
-            obs, _ = env.reset()
-            steps_per_episode = 0
-            continue
-
-        obs = next_obs
-
-        if global_step > learning_starts:
+        if global_step >= learning_starts:
             for _ in range(gradient_steps):
                 (
                     observations,
@@ -515,29 +591,40 @@ def train_ddpg(
                     rewards,
                     terminations,
                 )
+                actor_loss_value = ddpg_update_actor(
+                    policy, policy_optimizer, q, observations
+                )
+                soft_target_net_update(policy, policy_target, tau)
+                soft_target_net_update(q, q_target, tau)
+
                 if logger is not None:
-                    logger.record_stat("q loss", q_loss_value, step=global_step)
-                    logger.record_epoch("q", q, step=global_step)
+                    logger.record_stat(
+                        "q loss", q_loss_value, step=global_step + 1
+                    )
+                    logger.record_epoch("q", q, step=global_step + 1)
+                    logger.record_stat(
+                        "policy loss", actor_loss_value, step=global_step + 1
+                    )
+                    logger.record_epoch("policy", policy, step=global_step + 1)
+                    logger.record_epoch(
+                        "policy_target", policy_target, step=global_step + 1
+                    )
+                    logger.record_epoch(
+                        "q_target", q_target, step=global_step + 1
+                    )
 
-                if global_step % policy_frequency == 0:
-                    actor_loss_value = ddpg_update_actor(
-                        policy, policy_optimizer, q, observations
+        if termination or truncated:
+            if logger is not None:
+                if "episode" in info:
+                    logger.record_stat(
+                        "return", info["episode"]["r"], step=global_step + 1
                     )
+                logger.stop_episode(steps_per_episode)
+                logger.start_new_episode()
 
-                    update_target(policy, policy_target, tau)
-                    # TODO why is it updated less often than q?
-                    update_target(q, q_target, tau)
-
-                    if logger is not None:
-                        logger.record_stat(
-                            "policy loss", actor_loss_value, step=global_step
-                        )
-                        logger.record_epoch("policy", policy, step=global_step)
-                        logger.record_epoch(
-                            "policy_target", policy_target, step=global_step
-                        )
-                        logger.record_epoch(
-                            "q_target", q_target, step=global_step
-                        )
+            obs, _ = env.reset()
+            steps_per_episode = 0
+        else:
+            obs = next_obs
 
     return policy, policy_target, policy_optimizer, q, q_target, q_optimizer
diff --git a/rl_blox/algorithm/sac.py b/rl_blox/algorithm/sac.py
index 5d9d6915c..ffabc0ba5 100644
--- a/rl_blox/algorithm/sac.py
+++ b/rl_blox/algorithm/sac.py
@@ -15,7 +15,7 @@
     StochasticPolicyBase,
 )
 from ..logging.logger import LoggerBase
-from .ddpg import ReplayBuffer, mse_action_value_loss, update_target
+from .ddpg import ReplayBuffer, mse_action_value_loss, soft_target_net_update
 
 
 def sac_actor_loss(
@@ -66,10 +66,10 @@ def sac_actor_loss(
     actions = policy.sample(observations, action_key)
     log_prob = policy.log_probability(observations, actions)
     obs_act = jnp.concatenate((observations, actions), axis=-1)
-    qf1_pi = q1(obs_act).squeeze()
-    qf2_pi = q2(obs_act).squeeze()
-    min_qf_pi = jnp.minimum(qf1_pi, qf2_pi)
-    actor_loss = (alpha * log_prob - min_qf_pi).mean()
+    q1_value = q1(obs_act).squeeze()
+    q2_value = q2(obs_act).squeeze()
+    q_value = jnp.minimum(q1_value, q2_value)
+    actor_loss = (alpha * log_prob - q_value).mean()
     return actor_loss
 
 
@@ -426,16 +426,6 @@ def _sample_action(policy, obs, action_key):
         next_obs, reward, termination, truncation, info = env.step(action)
         steps_per_episode += 1
 
-        done = termination or truncation
-        if done:
-            if logger is not None:
-                logger.stop_episode(steps_per_episode)
-                logger.start_new_episode()
-                if "episode" in info:
-                    logger.record_stat("return", info["episode"]["r"])
-            obs, _ = env.reset()
-            steps_per_episode = 0
-
         rb.add_sample(
             observation=obs,
             action=action,
@@ -444,9 +434,7 @@ def _sample_action(policy, obs, action_key):
             termination=termination,
         )
 
-        obs = next_obs
-
-        if global_step > learning_starts:
+        if global_step >= learning_starts:
             observations, actions, rewards, next_observations, terminations = (
                 rb.sample_batch(batch_size, rng)
             )
@@ -470,10 +458,14 @@ def _sample_action(policy, obs, action_key):
                 entropy_control.alpha_,
             )
             if logger is not None:
-                logger.record_stat("q1 loss", q1_loss_value, step=global_step)
-                logger.record_epoch("q1", q1)
-                logger.record_stat("q2 loss", q2_loss_value, step=global_step)
-                logger.record_epoch("q2", q2)
+                logger.record_stat(
+                    "q1 loss", q1_loss_value, step=global_step + 1
+                )
+                logger.record_epoch("q1", q1, step=global_step + 1)
+                logger.record_stat(
+                    "q2 loss", q2_loss_value, step=global_step + 1
+                )
+                logger.record_epoch("q2", q2, step=global_step + 1)
 
             if global_step % policy_frequency == 0:
                 # compensate for delay by doing 'policy_frequency' updates
@@ -494,29 +486,48 @@ def _sample_action(policy, obs, action_key):
                     )
                     if logger is not None:
                         logger.record_stat(
-                            "policy loss", policy_loss_value, step=global_step
+                            "policy loss",
+                            policy_loss_value,
+                            step=global_step + 1,
+                        )
+                        logger.record_epoch(
+                            "policy", policy, step=global_step + 1
                         )
-                        logger.record_epoch("policy", policy, step=global_step)
                         logger.record_stat(
                             "alpha",
                             float(entropy_control.alpha_[0]),
-                            step=global_step,
+                            step=global_step + 1,
                         )
                         if autotune:
                             logger.record_stat(
                                 "alpha loss",
                                 exploration_loss_value,
-                                step=global_step,
+                                step=global_step + 1,
                             )
                             logger.record_epoch(
-                                "alpha", alpha, step=global_step
+                                "alpha", alpha, step=global_step + 1
                             )
 
             if global_step % target_network_frequency == 0:
-                update_target(q1, q1_target, tau)
-                logger.record_epoch("q1_target", q1_target, step=global_step)
-                update_target(q2, q2_target, tau)
-                logger.record_epoch("q2_target", q2_target, step=global_step)
+                soft_target_net_update(q1, q1_target, tau)
+                logger.record_epoch(
+                    "q1_target", q1_target, step=global_step + 1
+                )
+                soft_target_net_update(q2, q2_target, tau)
+                logger.record_epoch(
+                    "q2_target", q2_target, step=global_step + 1
+                )
+
+        if termination or truncation:
+            if logger is not None:
+                logger.stop_episode(steps_per_episode)
+                logger.start_new_episode()
+                if "episode" in info:
+                    logger.record_stat("return", info["episode"]["r"])
+            obs, _ = env.reset()
+            steps_per_episode = 0
+        else:
+            obs = next_obs
 
     return (
         policy,
@@ -647,16 +658,21 @@ def sac_update_critic(
     --------
     .ddpg.mse_action_value_loss
         The mean squared error loss.
+
+    sac_q_target
+        Generates target values for action-value functions.
     """
-    next_actions = policy.sample(next_observations, action_key)
-    next_log_pi = policy.log_probability(next_observations, next_actions)
-    next_obs_act = jnp.concatenate((next_observations, next_actions), axis=-1)
-    q1_next_target = q1_target(next_obs_act).squeeze()
-    q2_next_target = q2_target(next_obs_act).squeeze()
-    min_q_next_target = (
-        jnp.minimum(q1_next_target, q2_next_target) - alpha * next_log_pi
+    q_target_value = sac_q_target(
+        q1_target,
+        q2_target,
+        policy,
+        rewards,
+        next_observations,
+        terminations,
+        action_key,
+        alpha,
+        gamma,
     )
-    q_target_value = rewards + (1 - terminations) * gamma * min_q_next_target
 
     q1_loss_value, q1_grads = nnx.value_and_grad(
         mse_action_value_loss, argnums=3
@@ -668,3 +684,72 @@ def sac_update_critic(
     q2_optimizer.update(q2_grads)
 
     return q1_loss_value, q2_loss_value
+
+
+def sac_q_target(
+    q1_target: nnx.Module,
+    q2_target: nnx.Module,
+    policy: StochasticPolicyBase,
+    rewards: jnp.ndarray,
+    next_observations: jnp.ndarray,
+    terminations: jnp.ndarray,
+    action_key: jnp.ndarray,
+    alpha: float,
+    gamma: float,
+) -> jnp.ndarray:
+    r"""Target value for action-value functions in SAC.
+
+    Uses the bootstrap estimate
+
+    .. math::
+
+        r_{t+1} + \gamma
+        \left[\min(Q_1(o_{t+1}, a_{t+1}), Q_2(o_{t+1}, a_{t+1}))
+        - \alpha \log \pi(a_{t+1}|o_{t+1})\right]
+
+    based on the target networks of :math:`Q_1, Q_2` as a target value for the
+    Q network update with a mean squared error loss.
+
+    Parameters
+    ----------
+    q1_target : nnx.Module
+        Target network of q1.
+
+    q2_target : nnx.Module
+        Target network of q2.
+
+    policy : StochasticPolicyBase
+        Policy.
+
+    rewards : array
+        Rewards :math:`r_{t+1}`.
+
+    next_observations : array
+        Next observations :math:`o_{t+1}`.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    action_key : array
+        Random key for action sampling.
+
+    alpha : float
+        Entropy coefficient.
+
+    gamma : float
+        Discount factor of discounted infinite horizon return model.
+
+    Returns
+    -------
+    q_target_value : array
+        Target values for action-value functions.
+    """
+    next_actions = policy.sample(next_observations, action_key)
+    next_log_pi = policy.log_probability(next_observations, next_actions)
+    next_obs_act = jnp.concatenate((next_observations, next_actions), axis=-1)
+    q1_next_target = q1_target(next_obs_act).squeeze()
+    q2_next_target = q2_target(next_obs_act).squeeze()
+    min_q_next_target = (
+        jnp.minimum(q1_next_target, q2_next_target) - alpha * next_log_pi
+    )
+    return rewards + (1 - terminations) * gamma * min_q_next_target
diff --git a/rl_blox/algorithm/td3.py b/rl_blox/algorithm/td3.py
new file mode 100644
index 000000000..cfc8116a4
--- /dev/null
+++ b/rl_blox/algorithm/td3.py
@@ -0,0 +1,577 @@
+from collections import namedtuple
+from functools import partial
+
+import gymnasium as gym
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import tqdm
+from flax import nnx
+
+from ..blox.function_approximator.mlp import MLP
+from ..blox.function_approximator.policy_head import DeterministicTanhPolicy
+from ..logging.logger import LoggerBase
+from .ddpg import (
+    ReplayBuffer,
+    ddpg_update_actor,
+    mse_action_value_loss,
+    sample_actions,
+    soft_target_net_update,
+)
+
+
+@nnx.jit
+def td3_update_critic(
+    q1: nnx.Module,
+    q1_target: nnx.Module,
+    q1_optimizer: nnx.Optimizer,
+    q2: nnx.Module,
+    q2_target: nnx.Module,
+    q2_optimizer: nnx.Optimizer,
+    gamma: float,
+    observations: jnp.ndarray,
+    actions: jnp.ndarray,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+) -> tuple[float, float]:
+    r"""TD3 critic update.
+
+    Parameters
+    ----------
+    q1 : nnx.Module
+        Action-value function.
+
+    q1_target : nnx.Module
+        Target network of q.
+
+    q1_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    q2_target : nnx.Module
+        Target network of q.
+
+    q2_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    gamma : float
+        Discount factor of discounted infinite horizon return model.
+
+    observations : array
+        Observations :math:`o_t`.
+
+    actions : array
+        Actions :math:`a_t`.
+
+    next_observations : array
+        Next observations :math:`o_{t+1}`.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    rewards : array
+        Rewards :math:`r_{t+1}`.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    Returns
+    -------
+    q1_loss_value : float
+        Loss value.
+
+    q2_loss_value : float
+        Loss value.
+
+    See also
+    --------
+    mse_action_value_loss
+        The mean squared error loss.
+    """
+    q_bootstrap = double_q_deterministic_bootstrap_estimate(
+        rewards,
+        terminations,
+        gamma,
+        q1_target,
+        q2_target,
+        next_observations,
+        next_actions,
+    )
+
+    q1_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q1
+    )
+    q1_optimizer.update(grads)
+
+    q2_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q2
+    )
+    q2_optimizer.update(grads)
+
+    return q1_loss_value, q2_loss_value
+
+
+def double_q_deterministic_bootstrap_estimate(
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+    gamma: float,
+    q1: nnx.Module,
+    q2: nnx.Module,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+) -> jnp.ndarray:
+    r"""Bootstrap estimate of action-value function with deterministic policy.
+
+    .. math::
+
+        \mathbb{E}\left[R(o_t)\right]
+        \approx
+        r_{t+1}
+        + \gamma \min(Q_1(o_{t+1}, a_{t+1}), Q_2(o_{t+1}, a_{t+1}))
+
+    Parameters
+    ----------
+    rewards : array
+        Observed reward.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    gamma : float
+        Discount factor.
+
+    q1 : nnx.Module
+        Action-value function.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    next_observations : array
+        Next observations.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    Returns
+    -------
+    double_q_bootstrap : array
+        Double Q bootstrap estimate of action-value function.
+    """
+    obs_act = jnp.concatenate((next_observations, next_actions), axis=-1)
+    q_value = jnp.minimum(q1(obs_act).squeeze(), q2(obs_act).squeeze())
+    return rewards + (1 - terminations) * gamma * q_value
+
+
+def sample_target_actions(
+    action_low: jnp.ndarray,
+    action_high: jnp.ndarray,
+    action_scale: jnp.ndarray,
+    exploration_noise: float,
+    noise_clip: float,
+    policy: DeterministicTanhPolicy,
+    obs: jnp.ndarray,
+    key: jnp.ndarray,
+) -> jnp.ndarray:
+    """Sample target actions with truncated Gaussian noise.
+
+    Actions will be clipped to [action_low, action_high].
+    """
+    action = policy(obs)
+    noise = jax.random.multivariate_normal(
+        key,
+        jnp.zeros_like(action[0]),
+        jnp.diag(action_scale * exploration_noise),
+        shape=(action.shape[0],),
+    )
+    scaled_noise_clip = action_scale * noise_clip
+    clipped_noise = jnp.clip(noise, -scaled_noise_clip, scaled_noise_clip)
+    return jnp.clip(action + clipped_noise, action_low, action_high)
+
+
+def create_td3_state(
+    env: gym.Env[gym.spaces.Box, gym.spaces.Box],
+    policy_hidden_nodes: list[int] | tuple[int] = (256, 256),
+    policy_activation: str = "relu",
+    policy_learning_rate: float = 1e-3,
+    q_hidden_nodes: list[int] | tuple[int] = (256, 256),
+    q_activation: str = "relu",
+    q_learning_rate: float = 1e-3,
+    seed: int = 0,
+) -> namedtuple:
+    """Create components for SAC algorithm with default configuration."""
+    env.action_space.seed(seed)
+
+    policy_net = MLP(
+        env.observation_space.shape[0],
+        env.action_space.shape[0],
+        policy_hidden_nodes,
+        policy_activation,
+        nnx.Rngs(seed),
+    )
+    policy = DeterministicTanhPolicy(policy_net, env.action_space)
+    policy_optimizer = nnx.Optimizer(
+        policy, optax.adam(learning_rate=policy_learning_rate)
+    )
+
+    q1 = MLP(
+        env.observation_space.shape[0] + env.action_space.shape[0],
+        1,
+        q_hidden_nodes,
+        q_activation,
+        nnx.Rngs(seed),
+    )
+    q1_optimizer = nnx.Optimizer(q1, optax.adam(learning_rate=q_learning_rate))
+
+    q2 = MLP(
+        env.observation_space.shape[0] + env.action_space.shape[0],
+        1,
+        q_hidden_nodes,
+        q_activation,
+        nnx.Rngs(seed + 1),
+    )
+    q2_optimizer = nnx.Optimizer(q2, optax.adam(learning_rate=q_learning_rate))
+
+    return namedtuple(
+        "TD3State",
+        [
+            "policy",
+            "policy_optimizer",
+            "q1",
+            "q1_optimizer",
+            "q2",
+            "q2_optimizer",
+        ],
+    )(policy, policy_optimizer, q1, q1_optimizer, q2, q2_optimizer)
+
+
+def train_td3(
+    env: gym.Env[gym.spaces.Box, gym.spaces.Box],
+    policy: nnx.Module,
+    policy_optimizer: nnx.Optimizer,
+    q1: nnx.Module,
+    q1_optimizer: nnx.Optimizer,
+    q2: nnx.Module,
+    q2_optimizer: nnx.Optimizer,
+    seed: int = 1,
+    total_timesteps: int = 1_000_000,
+    buffer_size: int = 1_000_000,
+    gamma: float = 0.99,
+    tau: float = 0.005,
+    policy_delay: int = 2,
+    batch_size: int = 256,
+    gradient_steps: int = 1,
+    exploration_noise: float = 0.2,
+    noise_clip: float = 0.5,
+    learning_starts: int = 25_000,
+    policy_target: nnx.Optimizer | None = None,
+    q1_target: nnx.Optimizer | None = None,
+    q2_target: nnx.Optimizer | None = None,
+    logger: LoggerBase | None = None,
+) -> tuple[
+    nnx.Module,
+    nnx.Module,
+    nnx.Optimizer,
+    nnx.Module,
+    nnx.Module,
+    nnx.Optimizer,
+    nnx.Module,
+    nnx.Module,
+    nnx.Optimizer,
+]:
+    """Twin Delayed DDPG (TD3).
+
+    TD3 extends DDPG with three techniques to improve performance:
+
+    1. Clipped Double Q-Learning to mitigate overestimation bias of the value
+    2. Delayed policy updates, controlled by the parameter `policy_delay`
+    3. Target policy smoothing, i.e., sampling from the behavior policy with
+       clipped noise (parameter `noise_clip`) for the critic update.
+
+    As TD3 is based on DDPG, it uses :func:`.ddpg.ddpg_update_actor` to update
+    the policy and several details are the same.
+
+    To update the critic, we use :func:`td3_update_critic`. The target values
+    for the action value function are generated by
+    :func:`double_q_deterministic_bootstrap_estimate` based on the minimum of
+    the two target networks of the action-value functions `q1` and `q2`
+    (clipped double Q-learning) and next actions samples with
+    :func:`sample_target_actions` (target policy smoothing). The policy and
+    the target networks are only updated after `policy_delay` steps (delayed
+    policy updates).
+
+    Parameters
+    ----------
+    env : gymnasium.Env
+        Gymnasium environment.
+
+    policy : nnx.Module
+        Deterministic policy network.
+
+    policy_optimizer : nnx.Optimizer
+        Optimizer for the policy network.
+
+    q1 : nnx.Module
+        First Q network.
+
+    q1_optimizer: nnx.Optimizer
+        Optimizer for q1.
+
+    q2 : nnx.Module
+        Second Q network.
+
+    q2_optimizer : nnx.Optimizer
+        Optimizer for q2.
+
+    seed : int, optional
+        Seed for random number generators in Jax and NumPy.
+
+    total_timesteps : int, optional
+        Number of steps to execute in the environment.
+
+    buffer_size : int, optional
+        Size of the replay buffer.
+
+    gamma : float, optional
+        Discount factor.
+
+    tau : float, optional
+        Learning rate for polyak averaging of target policy and value function.
+
+    policy_delay : int, optional
+        Delayed policy updates. The policy is updated every `policy_delay`
+        steps.
+
+    batch_size : int, optional
+        Size of a batch during gradient computation.
+
+    gradient_steps : int, optional
+        Number of gradient steps during one training phase.
+
+    exploration_noise : float, optional
+        Exploration noise in action space. Will be scaled by half of the range
+        of the action space.
+
+    noise_clip : float, optional
+        Maximum absolute value of the exploration noise for sampling target
+        actions for the critic update. Will be scaled by half of the range
+        of the action space.
+
+    learning_starts : int, optional
+        Learning starts after this number of random steps was taken in the
+        environment.
+
+    policy_target : nnx.Module, optional
+        Target policy. Only has to be set if we want to continue training
+        from an old state.
+
+    q1_target : nnx.Module, optional
+        Target network. Only has to be set if we want to continue training
+        from an old state.
+
+    q2_target : nnx.Module, optional
+        Target network. Only has to be set if we want to continue training
+        from an old state.
+
+    logger : LoggerBase, optional
+        Experiment logger.
+
+    Returns
+    -------
+    policy : nnx.Module
+        Final policy.
+
+    policy_target : nnx.Module
+        Target policy.
+
+    policy_optimizer : nnx.Optimizer
+        Policy optimizer.
+
+    q1 : nnx.Module
+        Final state-action value function.
+
+    q1_target : nnx.Module
+        Target network.
+
+    q1_optimizer
+        Optimizer for Q network.
+
+    q2 : nnx.Module
+        Final state-action value function.
+
+    q2_target : nnx.Module
+        Target network.
+
+    q2_optimizer : nnx.Optimizer
+        Optimizer for Q network.
+
+    References
+    ----------
+    .. [1] Fujimoto, S., Hoof, H. &amp; Meger, D.. (2018). Addressing Function
+       Approximation Error in Actor-Critic Methods. Proceedings of the 35th
+       International Conference on Machine Learning, in Proceedings of Machine
+       Learning Research 80:1587-1596 Available from
+       https://proceedings.mlr.press/v80/fujimoto18a.html.
+
+    See Also
+    --------
+    .ddpg.train_ddpg
+        DDPG without the extensions of TD3.
+    """
+    rng = np.random.default_rng(seed)
+    key = jax.random.key(seed)
+
+    assert isinstance(
+        env.action_space, gym.spaces.Box
+    ), "only continuous action space is supported"
+
+    env.observation_space.dtype = np.float32
+    rb = ReplayBuffer(buffer_size)
+
+    action_scale = 0.5 * (env.action_space.high - env.action_space.low)
+    _sample_actions = nnx.jit(
+        partial(
+            sample_actions,
+            env.action_space.low,
+            env.action_space.high,
+            action_scale,
+            exploration_noise,
+        )
+    )
+    _sample_target_actions = nnx.jit(
+        partial(
+            sample_target_actions,
+            env.action_space.low,
+            env.action_space.high,
+            action_scale,
+            exploration_noise,
+            noise_clip,
+        )
+    )
+
+    if logger is not None:
+        logger.start_new_episode()
+    obs, _ = env.reset(seed=seed)
+    steps_per_episode = 0
+
+    if policy_target is None:
+        policy_target = nnx.clone(policy)
+    if q1_target is None:
+        q1_target = nnx.clone(q1)
+    if q2_target is None:
+        q2_target = nnx.clone(q2)
+
+    for global_step in tqdm.trange(total_timesteps):
+        if global_step < learning_starts:
+            action = env.action_space.sample()
+        else:
+            key, action_key = jax.random.split(key, 2)
+            action = np.asarray(
+                _sample_actions(policy, jnp.asarray(obs), action_key)
+            )
+
+        next_obs, reward, termination, truncated, info = env.step(action)
+        steps_per_episode += 1
+
+        rb.add_sample(
+            observation=obs,
+            action=action,
+            reward=reward,
+            next_observation=next_obs,
+            termination=termination,
+        )
+
+        if global_step >= learning_starts:
+            for _ in range(gradient_steps):
+                (
+                    observations,
+                    actions,
+                    rewards,
+                    next_observations,
+                    terminations,
+                ) = rb.sample_batch(batch_size, rng)
+
+                # policy smoothing: sample next actions from target policy
+                key, sampling_key = jax.random.split(key, 2)
+                next_actions = _sample_target_actions(
+                    policy_target, next_observations, sampling_key
+                )
+                q1_loss_value, q2_loss_value = td3_update_critic(
+                    q1,
+                    q1_target,
+                    q1_optimizer,
+                    q2,
+                    q2_target,
+                    q2_optimizer,
+                    gamma,
+                    observations,
+                    actions,
+                    next_observations,
+                    next_actions,
+                    rewards,
+                    terminations,
+                )
+
+                if logger is not None:
+                    logger.record_stat(
+                        "q1 loss", q1_loss_value, step=global_step + 1
+                    )
+                    logger.record_epoch("q1", q1, step=global_step + 1)
+                    logger.record_stat(
+                        "q2 loss", q2_loss_value, step=global_step + 1
+                    )
+                    logger.record_epoch("q2", q2, step=global_step + 1)
+
+                if global_step % policy_delay == 0:
+                    actor_loss_value = ddpg_update_actor(
+                        policy, policy_optimizer, q1, observations
+                    )
+                    soft_target_net_update(policy, policy_target, tau)
+                    soft_target_net_update(q1, q1_target, tau)
+                    soft_target_net_update(q2, q2_target, tau)
+                    if logger is not None:
+                        logger.record_stat(
+                            "policy loss",
+                            actor_loss_value,
+                            step=global_step + 1,
+                        )
+                    logger.record_epoch("policy", policy, step=global_step + 1)
+                    logger.record_epoch(
+                        "policy_target", policy_target, step=global_step + 1
+                    )
+                    logger.record_epoch(
+                        "q_target", q1_target, step=global_step + 1
+                    )
+                    logger.record_epoch(
+                        "q_target", q2_target, step=global_step + 1
+                    )
+
+        if termination or truncated:
+            if logger is not None:
+                if "episode" in info:
+                    logger.record_stat(
+                        "return", info["episode"]["r"], step=global_step + 1
+                    )
+                logger.stop_episode(steps_per_episode)
+                logger.start_new_episode()
+
+            obs, _ = env.reset()
+            steps_per_episode = 0
+        else:
+            obs = next_obs
+
+    return (
+        policy,
+        policy_target,
+        policy_optimizer,
+        q1,
+        q1_target,
+        q1_optimizer,
+        q2,
+        q2_target,
+        q2_optimizer,
+    )
diff --git a/rl_blox/logging/logger.py b/rl_blox/logging/logger.py
index 91b9532a6..f9bf97521 100644
--- a/rl_blox/logging/logger.py
+++ b/rl_blox/logging/logger.py
@@ -2,7 +2,6 @@
 import atexit
 import contextlib
 import os
-import shutil
 import time
 from typing import Any
 
@@ -152,7 +151,7 @@ class StandardLogger(LoggerBase):
 
         .. warning::
 
-            This directory will be deleted before the experiment starts!
+            This directory will be created if it does not exist.
 
     verbose : int, optional
         Verbosity level.
@@ -176,7 +175,7 @@ class StandardLogger(LoggerBase):
     checkpoint_path: dict[str, list[str]]
 
     def __init__(self, checkpoint_dir="/tmp/rl-blox/", verbose=0):
-        self.checkpoint_dir = checkpoint_dir
+        self.checkpoint_dir = os.path.abspath(checkpoint_dir)
         self.verbose = verbose
 
         self.env_name = None
@@ -341,9 +340,8 @@ def define_checkpoint_frequency(self, key: str, frequency: int):
 
     def _init_checkpointer(self):
         self.checkpointer = ocp.StandardCheckpointer()
-        if os.path.exists(self.checkpoint_dir):
-            shutil.rmtree(self.checkpoint_dir)
-        os.makedirs(self.checkpoint_dir)
+        if not os.path.exists(self.checkpoint_dir):
+            os.makedirs(self.checkpoint_dir)
 
     def record_epoch(
         self,
@@ -399,13 +397,14 @@ def record_epoch(
             self._save_checkpoint(key, value)
 
     def _save_checkpoint(self, key: str, value: Any):
-        checkpoint_path = (
-            f"{self.checkpoint_dir}"
+        checkpoint_path = os.path.join(
+            f"{self.checkpoint_dir}",
             f"{self.start_time}_{self.env_name}_{self.algorithm_name}_"
             f"{key}_{self.epoch[key]}/"
         )
         _, state = nnx.split(value)
         self.checkpointer.save(f"{checkpoint_path}", state)
+        self.checkpointer.wait_until_finished()
         self.checkpoint_path[key].append(checkpoint_path)
         if self.verbose:
             tqdm.tqdm.write(
diff --git a/tests/test_td3.py b/tests/test_td3.py
new file mode 100644
index 000000000..1effa0c4f
--- /dev/null
+++ b/tests/test_td3.py
@@ -0,0 +1,29 @@
+import gymnasium as gym
+
+from rl_blox.algorithm.td3 import create_td3_state, train_td3
+
+
+def test_td3():
+    env = gym.make("InvertedPendulum-v5")
+
+    td3_state = create_td3_state(
+        env,
+        policy_hidden_nodes=[32, 32],
+        policy_learning_rate=3e-4,
+        q_hidden_nodes=[128, 128],
+        q_learning_rate=1e-2,
+        seed=42,
+    )
+
+    train_td3(
+        env,
+        td3_state.policy,
+        td3_state.policy_optimizer,
+        td3_state.q1,
+        td3_state.q1_optimizer,
+        td3_state.q2,
+        td3_state.q2_optimizer,
+        total_timesteps=10,
+    )
+
+    env.close()


Reviews:
user: login: AlexanderFabisch
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T20:44:10Z
user: login: AlexanderFabisch
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T20:46:13Z
user: login: AlexanderFabisch
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T20:46:37Z
user: login: AlexanderFabisch
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T20:49:15Z
user: login: AlexanderFabisch
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T20:50:47Z
user: login: AlexanderFabisch
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T20:54:09Z
user: login: mlaux1
body: @AlexanderFabisch, this is still marked as WIP, so I won't merge it yet. Let me know when it's ready.
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-16T11:16:21Z
user: login: AlexanderFabisch
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-16T11:21:37Z
user: login: AlexanderFabisch
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-16T11:23:10Z

Review comments:
diff_hunk: @@ -0,0 +1,80 @@
+import gymnasium as gym
+import jax.numpy as jnp
+import numpy as np
+
+from rl_blox.algorithm.td3 import create_td3_state, train_td3
+from rl_blox.logging.logger import AIMLogger
+
+env_name = "Hopper-v5"
+env = gym.make(env_name)
+
+seed = 1
+verbose = 2
+env = gym.wrappers.RecordEpisodeStatistics(env)
+
+hparams_models = dict(
+    policy_hidden_nodes=[256, 256],
+    policy_learning_rate=3e-4,
+    q_hidden_nodes=[256, 256],
+    q_learning_rate=3e-4,
+    seed=seed,
+)
+hparams_algorithm = dict(
+    policy_delay=2,
+    exploration_noise=0.2,
+    noise_clip=0.5,
+    gradient_steps=1,
+    total_timesteps=1_000_000,
+    buffer_size=1_000_000,
+    learning_starts=5_000,
path: examples/td3_example.py
user: login: AlexanderFabisch
body: 25e3
created_at: 2025-05-09T20:44:10Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,80 @@
+import gymnasium as gym
+import jax.numpy as jnp
+import numpy as np
+
+from rl_blox.algorithm.td3 import create_td3_state, train_td3
+from rl_blox.logging.logger import AIMLogger
+
+env_name = "Hopper-v5"
+env = gym.make(env_name)
+
+seed = 1
+verbose = 2
+env = gym.wrappers.RecordEpisodeStatistics(env)
+
+hparams_models = dict(
+    policy_hidden_nodes=[256, 256],
+    policy_learning_rate=3e-4,
+    q_hidden_nodes=[256, 256],
+    q_learning_rate=3e-4,
path: examples/td3_example.py
user: login: AlexanderFabisch
body: 1e-3
created_at: 2025-05-09T20:46:13Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,80 @@
+import gymnasium as gym
+import jax.numpy as jnp
+import numpy as np
+
+from rl_blox.algorithm.td3 import create_td3_state, train_td3
+from rl_blox.logging.logger import AIMLogger
+
+env_name = "Hopper-v5"
+env = gym.make(env_name)
+
+seed = 1
+verbose = 2
+env = gym.wrappers.RecordEpisodeStatistics(env)
+
+hparams_models = dict(
+    policy_hidden_nodes=[256, 256],
+    policy_learning_rate=3e-4,
path: examples/td3_example.py
user: login: AlexanderFabisch
body: 1e-3
created_at: 2025-05-09T20:46:37Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,488 @@
+import gymnasium as gym
+from flax import nnx
+import numpy as np
+import jax
+import jax.numpy as jnp
+from functools import partial
+from collections import namedtuple
+import optax
+import tqdm
+
+from ..logging.logger import LoggerBase
+from .ddpg import ReplayBuffer, sample_actions, ddpg_update_actor, update_target, mse_action_value_loss
+from ..blox.function_approximator.mlp import MLP
+from ..blox.function_approximator.policy_head import DeterministicTanhPolicy
+
+
+@nnx.jit
+def td3_update_critic(
+    q1: nnx.Module,
+    q1_target: nnx.Module,
+    q1_optimizer: nnx.Optimizer,
+    q2: nnx.Module,
+    q2_target: nnx.Module,
+    q2_optimizer: nnx.Optimizer,
+    gamma: float,
+    observations: jnp.ndarray,
+    actions: jnp.ndarray,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+) -> tuple[float, float]:
+    r"""TD3 critic update.
+
+    Parameters
+    ----------
+    policy_target : nnx.Module
+        Target network of policy.
+
+    q1 : nnx.Module
+        Action-value function.
+
+    q1_target : nnx.Module
+        Target network of q.
+
+    q1_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    q2_target : nnx.Module
+        Target network of q.
+
+    q2_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    gamma : float
+        Discount factor of discounted infinite horizon return model.
+
+    observations : array
+        Observations :math:`o_t`.
+
+    actions : array
+        Actions :math:`a_t`.
+
+    next_observations : array
+        Next observations :math:`o_{t+1}`.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    rewards : array
+        Rewards :math:`r_{t+1}`.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    Returns
+    -------
+    q1_loss_value : float
+        Loss value.
+
+    q2_loss_value : float
+        Loss value.
+
+    See also
+    --------
+    mse_action_value_loss
+        The mean squared error loss.
+    """
+    q_bootstrap = double_q_deterministic_bootstrap_estimate(
+        rewards, terminations, gamma, q1_target, q2_target, next_observations, next_actions
+    )
+
+    q1_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q1
+    )
+    q1_optimizer.update(grads)
+
+    q2_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q2
+    )
+    q2_optimizer.update(grads)
+
+    return q1_loss_value, q2_loss_value
+
+
+def double_q_deterministic_bootstrap_estimate(
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+    gamma: float,
+    q1: nnx.Module,
+    q2: nnx.Module,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+) -> jnp.ndarray:
+    r"""Bootstrap estimate of action-value function with deterministic policy.
+
+    .. math::
+
+        \mathbb{E}\left[R(o_t)\right]
+        \approx
+        r_{t+1}
+        + \gamma \min(Q_1(o_{t+1}, a_{t+1}), Q_2(o_{t+1}, a_{t+1}))
+
+    Parameters
+    ----------
+    policy : nnx.Module
+        Deterministic policy for action selection.
+
+    rewards : array
+        Observed reward.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    gamma : float
+        Discount factor.
+
+    q1 : nnx.Module
+        Action-value function.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    next_observations : array
+        Next observations.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    Returns
+    -------
+    double_q_bootstrap : array
+        Double Q bootstrap estimate of action-value function.
+    """
+    obs_act = jnp.concatenate((next_observations, next_actions), axis=-1)
+    q_value = jnp.minimum(q1(obs_act).squeeze(), q2(obs_act).squeeze())
+    return rewards + (1 - terminations) * gamma * q_value
+
+
+def sample_target_actions(
+    action_low: jnp.ndarray,
+    action_high: jnp.ndarray,
+    action_scale: jnp.ndarray,
+    exploration_noise: float,
+    noise_clip: float,
+    policy: DeterministicTanhPolicy,
+    obs: jnp.ndarray,
+    key: jnp.ndarray,
+) -> jnp.ndarray:
+    """Sample actions with deterministic policy and Gaussian action noise.
+
+    Actions will be clipped to [action_low, action_high].
+    """
+    action = policy(obs)
+    noise = jax.random.multivariate_normal(
+        key,
+        jnp.zeros_like(action[0]),
+        jnp.diag(action_scale * exploration_noise),
+        shape=(action.shape[0],),
+    )
+    clipped_noise = jnp.clip(noise, -noise_clip, noise_clip)
+    return jnp.clip(action + clipped_noise, action_low, action_high)
+
+
+def create_td3_state(
+    env: gym.Env[gym.spaces.Box, gym.spaces.Box],
+    policy_hidden_nodes: list[int] | tuple[int] = (256, 256),
+    policy_activation: str = "swish",
+    policy_learning_rate: float = 3e-4,
path: rl_blox/algorithm/td3.py
user: login: AlexanderFabisch
body: 1e-3
created_at: 2025-05-09T20:49:15Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,488 @@
+import gymnasium as gym
+from flax import nnx
+import numpy as np
+import jax
+import jax.numpy as jnp
+from functools import partial
+from collections import namedtuple
+import optax
+import tqdm
+
+from ..logging.logger import LoggerBase
+from .ddpg import ReplayBuffer, sample_actions, ddpg_update_actor, update_target, mse_action_value_loss
+from ..blox.function_approximator.mlp import MLP
+from ..blox.function_approximator.policy_head import DeterministicTanhPolicy
+
+
+@nnx.jit
+def td3_update_critic(
+    q1: nnx.Module,
+    q1_target: nnx.Module,
+    q1_optimizer: nnx.Optimizer,
+    q2: nnx.Module,
+    q2_target: nnx.Module,
+    q2_optimizer: nnx.Optimizer,
+    gamma: float,
+    observations: jnp.ndarray,
+    actions: jnp.ndarray,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+) -> tuple[float, float]:
+    r"""TD3 critic update.
+
+    Parameters
+    ----------
+    policy_target : nnx.Module
+        Target network of policy.
+
+    q1 : nnx.Module
+        Action-value function.
+
+    q1_target : nnx.Module
+        Target network of q.
+
+    q1_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    q2_target : nnx.Module
+        Target network of q.
+
+    q2_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    gamma : float
+        Discount factor of discounted infinite horizon return model.
+
+    observations : array
+        Observations :math:`o_t`.
+
+    actions : array
+        Actions :math:`a_t`.
+
+    next_observations : array
+        Next observations :math:`o_{t+1}`.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    rewards : array
+        Rewards :math:`r_{t+1}`.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    Returns
+    -------
+    q1_loss_value : float
+        Loss value.
+
+    q2_loss_value : float
+        Loss value.
+
+    See also
+    --------
+    mse_action_value_loss
+        The mean squared error loss.
+    """
+    q_bootstrap = double_q_deterministic_bootstrap_estimate(
+        rewards, terminations, gamma, q1_target, q2_target, next_observations, next_actions
+    )
+
+    q1_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q1
+    )
+    q1_optimizer.update(grads)
+
+    q2_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q2
+    )
+    q2_optimizer.update(grads)
+
+    return q1_loss_value, q2_loss_value
+
+
+def double_q_deterministic_bootstrap_estimate(
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+    gamma: float,
+    q1: nnx.Module,
+    q2: nnx.Module,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+) -> jnp.ndarray:
+    r"""Bootstrap estimate of action-value function with deterministic policy.
+
+    .. math::
+
+        \mathbb{E}\left[R(o_t)\right]
+        \approx
+        r_{t+1}
+        + \gamma \min(Q_1(o_{t+1}, a_{t+1}), Q_2(o_{t+1}, a_{t+1}))
+
+    Parameters
+    ----------
+    policy : nnx.Module
+        Deterministic policy for action selection.
+
+    rewards : array
+        Observed reward.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    gamma : float
+        Discount factor.
+
+    q1 : nnx.Module
+        Action-value function.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    next_observations : array
+        Next observations.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    Returns
+    -------
+    double_q_bootstrap : array
+        Double Q bootstrap estimate of action-value function.
+    """
+    obs_act = jnp.concatenate((next_observations, next_actions), axis=-1)
+    q_value = jnp.minimum(q1(obs_act).squeeze(), q2(obs_act).squeeze())
+    return rewards + (1 - terminations) * gamma * q_value
+
+
+def sample_target_actions(
+    action_low: jnp.ndarray,
+    action_high: jnp.ndarray,
+    action_scale: jnp.ndarray,
+    exploration_noise: float,
+    noise_clip: float,
+    policy: DeterministicTanhPolicy,
+    obs: jnp.ndarray,
+    key: jnp.ndarray,
+) -> jnp.ndarray:
+    """Sample actions with deterministic policy and Gaussian action noise.
+
+    Actions will be clipped to [action_low, action_high].
+    """
+    action = policy(obs)
+    noise = jax.random.multivariate_normal(
+        key,
+        jnp.zeros_like(action[0]),
+        jnp.diag(action_scale * exploration_noise),
+        shape=(action.shape[0],),
+    )
+    clipped_noise = jnp.clip(noise, -noise_clip, noise_clip)
+    return jnp.clip(action + clipped_noise, action_low, action_high)
+
+
+def create_td3_state(
+    env: gym.Env[gym.spaces.Box, gym.spaces.Box],
+    policy_hidden_nodes: list[int] | tuple[int] = (256, 256),
+    policy_activation: str = "swish",
+    policy_learning_rate: float = 3e-4,
+    q_hidden_nodes: list[int] | tuple[int] = (256, 256),
+    q_activation: str = "relu",
+    q_learning_rate: float = 1e-3,
+    seed: int = 0,
+) -> namedtuple:
+    """Create components for SAC algorithm with default configuration."""
+    env.action_space.seed(seed)
+
+    policy_net = MLP(
+        env.observation_space.shape[0],
+        env.action_space.shape[0],
+        policy_hidden_nodes,
+        policy_activation,
+        nnx.Rngs(seed),
+    )
+    policy = DeterministicTanhPolicy(policy_net, env.action_space)
+    policy_optimizer = nnx.Optimizer(
+        policy, optax.adam(learning_rate=policy_learning_rate)
+    )
+
+    q1 = MLP(
+        env.observation_space.shape[0] + env.action_space.shape[0],
+        1,
+        q_hidden_nodes,
+        q_activation,
+        nnx.Rngs(seed),
+    )
+    q1_optimizer = nnx.Optimizer(q1, optax.adam(learning_rate=q_learning_rate))
+
+    q2 = MLP(
+        env.observation_space.shape[0] + env.action_space.shape[0],
+        1,
+        q_hidden_nodes,
+        q_activation,
+        nnx.Rngs(seed + 1),
+    )
+    q2_optimizer = nnx.Optimizer(q2, optax.adam(learning_rate=q_learning_rate))
+
+    return namedtuple(
+        "TD3State",
+        [
+            "policy",
+            "policy_optimizer",
+            "q1",
+            "q1_optimizer",
+            "q2",
+            "q2_optimizer",
+        ],
+    )(policy, policy_optimizer, q1, q1_optimizer, q2, q2_optimizer)
+
+
+def train_td3(
+    env: gym.Env[gym.spaces.Box, gym.spaces.Box],
+    policy: nnx.Module,
+    policy_optimizer: nnx.Optimizer,
+    q1: nnx.Module,
+    q1_optimizer: nnx.Optimizer,
+    q2: nnx.Module,
+    q2_optimizer: nnx.Optimizer,
+    seed: int = 1,
+    total_timesteps: int = 1_000_000,
+    buffer_size: int = 1_000_000,
+    gamma: float = 0.99,
+    tau: float = 0.005,
+    policy_delay: int = 2,
+    batch_size: int = 256,
+    gradient_steps: int = 1,
+    exploration_noise: float = 0.2,
path: rl_blox/algorithm/td3.py
user: login: AlexanderFabisch
body: Scale by action space
created_at: 2025-05-09T20:50:47Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,488 @@
+import gymnasium as gym
+from flax import nnx
+import numpy as np
+import jax
+import jax.numpy as jnp
+from functools import partial
+from collections import namedtuple
+import optax
+import tqdm
+
+from ..logging.logger import LoggerBase
+from .ddpg import ReplayBuffer, sample_actions, ddpg_update_actor, update_target, mse_action_value_loss
+from ..blox.function_approximator.mlp import MLP
+from ..blox.function_approximator.policy_head import DeterministicTanhPolicy
+
+
+@nnx.jit
+def td3_update_critic(
+    q1: nnx.Module,
+    q1_target: nnx.Module,
+    q1_optimizer: nnx.Optimizer,
+    q2: nnx.Module,
+    q2_target: nnx.Module,
+    q2_optimizer: nnx.Optimizer,
+    gamma: float,
+    observations: jnp.ndarray,
+    actions: jnp.ndarray,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+) -> tuple[float, float]:
+    r"""TD3 critic update.
+
+    Parameters
+    ----------
+    policy_target : nnx.Module
+        Target network of policy.
+
+    q1 : nnx.Module
+        Action-value function.
+
+    q1_target : nnx.Module
+        Target network of q.
+
+    q1_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    q2_target : nnx.Module
+        Target network of q.
+
+    q2_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    gamma : float
+        Discount factor of discounted infinite horizon return model.
+
+    observations : array
+        Observations :math:`o_t`.
+
+    actions : array
+        Actions :math:`a_t`.
+
+    next_observations : array
+        Next observations :math:`o_{t+1}`.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    rewards : array
+        Rewards :math:`r_{t+1}`.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    Returns
+    -------
+    q1_loss_value : float
+        Loss value.
+
+    q2_loss_value : float
+        Loss value.
+
+    See also
+    --------
+    mse_action_value_loss
+        The mean squared error loss.
+    """
+    q_bootstrap = double_q_deterministic_bootstrap_estimate(
+        rewards, terminations, gamma, q1_target, q2_target, next_observations, next_actions
+    )
+
+    q1_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q1
+    )
+    q1_optimizer.update(grads)
+
+    q2_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q2
+    )
+    q2_optimizer.update(grads)
+
+    return q1_loss_value, q2_loss_value
+
+
+def double_q_deterministic_bootstrap_estimate(
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+    gamma: float,
+    q1: nnx.Module,
+    q2: nnx.Module,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+) -> jnp.ndarray:
+    r"""Bootstrap estimate of action-value function with deterministic policy.
+
+    .. math::
+
+        \mathbb{E}\left[R(o_t)\right]
+        \approx
+        r_{t+1}
+        + \gamma \min(Q_1(o_{t+1}, a_{t+1}), Q_2(o_{t+1}, a_{t+1}))
+
+    Parameters
+    ----------
+    policy : nnx.Module
+        Deterministic policy for action selection.
+
+    rewards : array
+        Observed reward.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    gamma : float
+        Discount factor.
+
+    q1 : nnx.Module
+        Action-value function.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    next_observations : array
+        Next observations.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    Returns
+    -------
+    double_q_bootstrap : array
+        Double Q bootstrap estimate of action-value function.
+    """
+    obs_act = jnp.concatenate((next_observations, next_actions), axis=-1)
+    q_value = jnp.minimum(q1(obs_act).squeeze(), q2(obs_act).squeeze())
+    return rewards + (1 - terminations) * gamma * q_value
+
+
+def sample_target_actions(
+    action_low: jnp.ndarray,
+    action_high: jnp.ndarray,
+    action_scale: jnp.ndarray,
+    exploration_noise: float,
+    noise_clip: float,
+    policy: DeterministicTanhPolicy,
+    obs: jnp.ndarray,
+    key: jnp.ndarray,
+) -> jnp.ndarray:
+    """Sample actions with deterministic policy and Gaussian action noise.
+
+    Actions will be clipped to [action_low, action_high].
+    """
+    action = policy(obs)
+    noise = jax.random.multivariate_normal(
+        key,
+        jnp.zeros_like(action[0]),
+        jnp.diag(action_scale * exploration_noise),
+        shape=(action.shape[0],),
+    )
+    clipped_noise = jnp.clip(noise, -noise_clip, noise_clip)
+    return jnp.clip(action + clipped_noise, action_low, action_high)
+
+
+def create_td3_state(
+    env: gym.Env[gym.spaces.Box, gym.spaces.Box],
+    policy_hidden_nodes: list[int] | tuple[int] = (256, 256),
+    policy_activation: str = "swish",
path: rl_blox/algorithm/td3.py
user: login: AlexanderFabisch
body: Relu
created_at: 2025-05-09T20:54:09Z
author_association: COLLABORATOR
diff_hunk: @@ -15,7 +15,7 @@
     StochasticPolicyBase,
 )
 from ..logging.logger import LoggerBase
-from .ddpg import ReplayBuffer, mse_action_value_loss, update_target
+from .ddpg import ReplayBuffer, mse_action_value_loss, soft_target_net_update
path: rl_blox/algorithm/sac.py
user: login: mlaux1
body: these parts should eventually be part of the blox interface
created_at: 2025-05-16T11:08:17Z
author_association: OWNER
diff_hunk: @@ -0,0 +1,577 @@
+from collections import namedtuple
+from functools import partial
+
+import gymnasium as gym
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import tqdm
+from flax import nnx
+
+from ..blox.function_approximator.mlp import MLP
+from ..blox.function_approximator.policy_head import DeterministicTanhPolicy
+from ..logging.logger import LoggerBase
+from .ddpg import (
+    ReplayBuffer,
+    ddpg_update_actor,
+    mse_action_value_loss,
+    sample_actions,
+    soft_target_net_update,
+)
+
+
+@nnx.jit
+def td3_update_critic(
+    q1: nnx.Module,
+    q1_target: nnx.Module,
+    q1_optimizer: nnx.Optimizer,
+    q2: nnx.Module,
+    q2_target: nnx.Module,
+    q2_optimizer: nnx.Optimizer,
+    gamma: float,
+    observations: jnp.ndarray,
+    actions: jnp.ndarray,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+) -> tuple[float, float]:
+    r"""TD3 critic update.
+
+    Parameters
+    ----------
+    q1 : nnx.Module
+        Action-value function.
+
+    q1_target : nnx.Module
+        Target network of q.
+
+    q1_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    q2_target : nnx.Module
+        Target network of q.
+
+    q2_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    gamma : float
+        Discount factor of discounted infinite horizon return model.
+
+    observations : array
+        Observations :math:`o_t`.
+
+    actions : array
+        Actions :math:`a_t`.
+
+    next_observations : array
+        Next observations :math:`o_{t+1}`.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    rewards : array
+        Rewards :math:`r_{t+1}`.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    Returns
+    -------
+    q1_loss_value : float
+        Loss value.
+
+    q2_loss_value : float
+        Loss value.
+
+    See also
+    --------
+    mse_action_value_loss
+        The mean squared error loss.
+    """
+    q_bootstrap = double_q_deterministic_bootstrap_estimate(
+        rewards,
+        terminations,
+        gamma,
+        q1_target,
+        q2_target,
+        next_observations,
+        next_actions,
+    )
+
+    q1_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q1
+    )
+    q1_optimizer.update(grads)
+
+    q2_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q2
+    )
+    q2_optimizer.update(grads)
+
+    return q1_loss_value, q2_loss_value
+
+
+def double_q_deterministic_bootstrap_estimate(
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+    gamma: float,
+    q1: nnx.Module,
+    q2: nnx.Module,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+) -> jnp.ndarray:
+    r"""Bootstrap estimate of action-value function with deterministic policy.
+
+    .. math::
+
+        \mathbb{E}\left[R(o_t)\right]
+        \approx
+        r_{t+1}
+        + \gamma \min(Q_1(o_{t+1}, a_{t+1}), Q_2(o_{t+1}, a_{t+1}))
+
+    Parameters
+    ----------
+    rewards : array
+        Observed reward.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    gamma : float
+        Discount factor.
+
+    q1 : nnx.Module
+        Action-value function.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    next_observations : array
+        Next observations.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    Returns
+    -------
+    double_q_bootstrap : array
+        Double Q bootstrap estimate of action-value function.
+    """
+    obs_act = jnp.concatenate((next_observations, next_actions), axis=-1)
+    q_value = jnp.minimum(q1(obs_act).squeeze(), q2(obs_act).squeeze())
+    return rewards + (1 - terminations) * gamma * q_value
+
+
+def sample_target_actions(
+    action_low: jnp.ndarray,
+    action_high: jnp.ndarray,
+    action_scale: jnp.ndarray,
+    exploration_noise: float,
+    noise_clip: float,
+    policy: DeterministicTanhPolicy,
+    obs: jnp.ndarray,
+    key: jnp.ndarray,
+) -> jnp.ndarray:
+    """Sample target actions with truncated Gaussian noise.
+
+    Actions will be clipped to [action_low, action_high].
+    """
+    action = policy(obs)
+    noise = jax.random.multivariate_normal(
+        key,
+        jnp.zeros_like(action[0]),
+        jnp.diag(action_scale * exploration_noise),
+        shape=(action.shape[0],),
+    )
+    scaled_noise_clip = action_scale * noise_clip
+    clipped_noise = jnp.clip(noise, -scaled_noise_clip, scaled_noise_clip)
+    return jnp.clip(action + clipped_noise, action_low, action_high)
+
+
+def create_td3_state(
path: rl_blox/algorithm/td3.py
user: login: mlaux1
body: I think these create_state functions are ok for now, but I think not the most intuitive. Some of these values I'd expect to be able to configure from the algorithm API directly, e.g. learning rates.
created_at: 2025-05-16T11:13:20Z
author_association: OWNER
diff_hunk: @@ -15,7 +15,7 @@
     StochasticPolicyBase,
 )
 from ..logging.logger import LoggerBase
-from .ddpg import ReplayBuffer, mse_action_value_loss, update_target
+from .ddpg import ReplayBuffer, mse_action_value_loss, soft_target_net_update
path: rl_blox/algorithm/sac.py
user: login: AlexanderFabisch
body: Definitely, but not in this PR as it is already quite large.
created_at: 2025-05-16T11:21:36Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,577 @@
+from collections import namedtuple
+from functools import partial
+
+import gymnasium as gym
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import tqdm
+from flax import nnx
+
+from ..blox.function_approximator.mlp import MLP
+from ..blox.function_approximator.policy_head import DeterministicTanhPolicy
+from ..logging.logger import LoggerBase
+from .ddpg import (
+    ReplayBuffer,
+    ddpg_update_actor,
+    mse_action_value_loss,
+    sample_actions,
+    soft_target_net_update,
+)
+
+
+@nnx.jit
+def td3_update_critic(
+    q1: nnx.Module,
+    q1_target: nnx.Module,
+    q1_optimizer: nnx.Optimizer,
+    q2: nnx.Module,
+    q2_target: nnx.Module,
+    q2_optimizer: nnx.Optimizer,
+    gamma: float,
+    observations: jnp.ndarray,
+    actions: jnp.ndarray,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+) -> tuple[float, float]:
+    r"""TD3 critic update.
+
+    Parameters
+    ----------
+    q1 : nnx.Module
+        Action-value function.
+
+    q1_target : nnx.Module
+        Target network of q.
+
+    q1_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    q2_target : nnx.Module
+        Target network of q.
+
+    q2_optimizer : nnx.Optimizer
+        Optimizer of q.
+
+    gamma : float
+        Discount factor of discounted infinite horizon return model.
+
+    observations : array
+        Observations :math:`o_t`.
+
+    actions : array
+        Actions :math:`a_t`.
+
+    next_observations : array
+        Next observations :math:`o_{t+1}`.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    rewards : array
+        Rewards :math:`r_{t+1}`.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    Returns
+    -------
+    q1_loss_value : float
+        Loss value.
+
+    q2_loss_value : float
+        Loss value.
+
+    See also
+    --------
+    mse_action_value_loss
+        The mean squared error loss.
+    """
+    q_bootstrap = double_q_deterministic_bootstrap_estimate(
+        rewards,
+        terminations,
+        gamma,
+        q1_target,
+        q2_target,
+        next_observations,
+        next_actions,
+    )
+
+    q1_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q1
+    )
+    q1_optimizer.update(grads)
+
+    q2_loss_value, grads = nnx.value_and_grad(mse_action_value_loss, argnums=3)(
+        observations, actions, q_bootstrap, q2
+    )
+    q2_optimizer.update(grads)
+
+    return q1_loss_value, q2_loss_value
+
+
+def double_q_deterministic_bootstrap_estimate(
+    rewards: jnp.ndarray,
+    terminations: jnp.ndarray,
+    gamma: float,
+    q1: nnx.Module,
+    q2: nnx.Module,
+    next_observations: jnp.ndarray,
+    next_actions: jnp.ndarray,
+) -> jnp.ndarray:
+    r"""Bootstrap estimate of action-value function with deterministic policy.
+
+    .. math::
+
+        \mathbb{E}\left[R(o_t)\right]
+        \approx
+        r_{t+1}
+        + \gamma \min(Q_1(o_{t+1}, a_{t+1}), Q_2(o_{t+1}, a_{t+1}))
+
+    Parameters
+    ----------
+    rewards : array
+        Observed reward.
+
+    terminations : array
+        Indicates if a terminal state was reached in this step.
+
+    gamma : float
+        Discount factor.
+
+    q1 : nnx.Module
+        Action-value function.
+
+    q2 : nnx.Module
+        Action-value function.
+
+    next_observations : array
+        Next observations.
+
+    next_actions : array
+        Sampled target actions :math:`a_{t+1}`.
+
+    Returns
+    -------
+    double_q_bootstrap : array
+        Double Q bootstrap estimate of action-value function.
+    """
+    obs_act = jnp.concatenate((next_observations, next_actions), axis=-1)
+    q_value = jnp.minimum(q1(obs_act).squeeze(), q2(obs_act).squeeze())
+    return rewards + (1 - terminations) * gamma * q_value
+
+
+def sample_target_actions(
+    action_low: jnp.ndarray,
+    action_high: jnp.ndarray,
+    action_scale: jnp.ndarray,
+    exploration_noise: float,
+    noise_clip: float,
+    policy: DeterministicTanhPolicy,
+    obs: jnp.ndarray,
+    key: jnp.ndarray,
+) -> jnp.ndarray:
+    """Sample target actions with truncated Gaussian noise.
+
+    Actions will be clipped to [action_low, action_high].
+    """
+    action = policy(obs)
+    noise = jax.random.multivariate_normal(
+        key,
+        jnp.zeros_like(action[0]),
+        jnp.diag(action_scale * exploration_noise),
+        shape=(action.shape[0],),
+    )
+    scaled_noise_clip = action_scale * noise_clip
+    clipped_noise = jnp.clip(noise, -scaled_noise_clip, scaled_noise_clip)
+    return jnp.clip(action + clipped_noise, action_low, action_high)
+
+
+def create_td3_state(
path: rl_blox/algorithm/td3.py
user: login: AlexanderFabisch
body: My Idea was that this is a convenient way to generate a good default, but you wouldn't normally use it. For example, if you want to use Simba+SAC, you would replace this by the initialization of a Simba-style network.
created_at: 2025-05-16T11:23:10Z
author_association: COLLABORATOR

PR comments:
user: login: AlexanderFabisch
created_at: 2025-05-16T06:05:29Z
body: @mlaux1 Finished the comparison of TD3 and DDPG on the hopper environment. Results look good and in line with the paper and stable baselines although they used slightly different environments. TD3 is much better than DDPG.
author_association: COLLABORATOR
user: login: AlexanderFabisch
created_at: 2025-05-16T11:24:06Z
body: > @AlexanderFabisch, this is still marked as WIP, so I won't merge it yet. Let me know when it's ready.

I Forgot about that. You can merge it.

author_association: COLLABORATOR