diff --git a/build.sh b/build.sh
index 454fda2..df67e56 100755
--- a/build.sh
+++ b/build.sh
@@ -1,4 +1,5 @@
+# auth in ~/.pypirc
 rm -rf ../python-interface/dist
 pip install --upgrade build twine
 python3 -m build
-python3 -m twine upload dist/* --password $PYPI_TOKEN
\ No newline at end of file
+python3 -m twine upload dist/* 
diff --git a/examples/custom_environment/my_pendulum/operations_generic.h b/examples/custom_environment/my_pendulum/operations_generic.h
index 273e6e8..1b446de 100644
--- a/examples/custom_environment/my_pendulum/operations_generic.h
+++ b/examples/custom_environment/my_pendulum/operations_generic.h
@@ -22,7 +22,7 @@ namespace rl_tools{
     template<typename DEVICE, typename SPEC>
     static void free(DEVICE& device, const MyPendulum<SPEC>& env){ }
     template<typename DEVICE, typename SPEC>
-    static void init(DEVICE& device, const MyPendulum<SPEC>& env, typename MyPendulum<SPEC>::Parameters& parameters){ }
+    static void init(DEVICE& device, const MyPendulum<SPEC>& env){ }
     template<typename DEVICE, typename SPEC>
     static void initial_parameters(DEVICE& device, const MyPendulum<SPEC>& env, typename MyPendulum<SPEC>::Parameters& parameters){ }
     template<typename DEVICE, typename SPEC, typename RNG>
diff --git a/examples/pendulum/ppo/train.py b/examples/pendulum/ppo/train.py
index 3e0d980..21e5a17 100644
--- a/examples/pendulum/ppo/train.py
+++ b/examples/pendulum/ppo/train.py
@@ -25,19 +25,19 @@ def env_factory(**kwargs):
         selected_env,
         interface_name=env_name,
         enable_evaluation=True,
-        evaluation_interval=10,
+        evaluation_interval=100,
         # force_blas=True,
         num_evaluation_episodes=10,
-        OPTIMIZER_ALPHA=1e-3,
+        OPTIMIZER_ALPHA=1e-4,
         ACTION_ENTROPY_COEFFICIENT = 0.0,
-        N_EPOCHS = 2,
+        N_EPOCHS = 1,
         GAMMA = 0.9,
-        BATCH_SIZE = 256,
-        ACTOR_HIDDEN_DIM = 64,
-        CRITIC_HIDDEN_DIM = 64,
-        ON_POLICY_RUNNER_STEPS_PER_ENV = 1024,
-        N_ENVIRONMENTS = 4,
-        STEP_LIMIT = 74, # 1024 * 4 * 74 ~ 300k steps
+        BATCH_SIZE = 128,
+        ACTOR_HIDDEN_DIM = 32,
+        CRITIC_HIDDEN_DIM = 32,
+        ON_POLICY_RUNNER_STEPS_PER_ENV = 128,
+        N_ENVIRONMENTS = 8,
+        STEP_LIMIT = 2000, # 1024 * 4 * 74 ~ 300k steps
     )
     state = ppo.State(seed)
 
diff --git a/pyproject.toml b/pyproject.toml
index baa2def..e1c2099 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "rltools"
-version = "1.2.1"
+version = "2.2.1"
 authors = [
   { name="Jonas Eschmann", email="jonas.eschmann@gmail.com" },
 ]
diff --git a/rltools/external/rl_tools b/rltools/external/rl_tools
index b1213d8..56fd71a 160000
--- a/rltools/external/rl_tools
+++ b/rltools/external/rl_tools
@@ -1 +1 @@
-Subproject commit b1213d8fd5af287cc0a4f2a7ff295247ef1b4eac
+Subproject commit 56fd71aeeedc56cb24f9db6cf40e5b9e49dd3bdd
diff --git a/rltools/interface/algorithms/ppo/template.h b/rltools/interface/algorithms/ppo/template.h
index c089833..3bf8206 100644
--- a/rltools/interface/algorithms/ppo/template.h
+++ b/rltools/interface/algorithms/ppo/template.h
@@ -7,8 +7,9 @@
 
 namespace rlt = rl_tools;
 
-template<typename T, typename TI, typename ENVIRONMENT>
-struct PPO_LOOP_CORE_PARAMETERS: rlt::rl::algorithms::ppo::loop::core::DefaultParameters<T, TI, ENVIRONMENT>{
+template<typename TYPE_POLICY, typename TI, typename ENVIRONMENT>
+struct PPO_LOOP_CORE_PARAMETERS: rlt::rl::algorithms::ppo::loop::core::DefaultParameters<TYPE_POLICY, TI, ENVIRONMENT>{
+    using T = typename TYPE_POLICY::DEFAULT;
     static constexpr TI STEP_LIMIT = $STEP_LIMIT;
     static constexpr TI ACTOR_HIDDEN_DIM = $ACTOR_HIDDEN_DIM;
     static constexpr TI ACTOR_NUM_LAYERS = $ACTOR_NUM_LAYERS;
@@ -22,7 +23,7 @@ struct PPO_LOOP_CORE_PARAMETERS: rlt::rl::algorithms::ppo::loop::core::DefaultPa
     static constexpr TI DATASET_SIZE = ON_POLICY_RUNNER_STEPS_PER_ENV * N_ENVIRONMENTS;
     static constexpr TI BATCH_SIZE = $BATCH_SIZE;
 
-    struct PPO_PARAMETERS: rlt::rl::algorithms::ppo::DefaultParameters<T, TI, BATCH_SIZE>{
+    struct PPO_PARAMETERS: rlt::rl::algorithms::ppo::DefaultParameters<TYPE_POLICY, TI, BATCH_SIZE>{
         static constexpr T GAMMA = $GAMMA;
         static constexpr T LAMBDA = $LAMBDA;
         static constexpr T EPSILON_CLIP = $EPSILON_CLIP;
@@ -46,7 +47,7 @@ struct PPO_LOOP_CORE_PARAMETERS: rlt::rl::algorithms::ppo::loop::core::DefaultPa
     };
 
 
-    struct OPTIMIZER_PARAMETERS: rlt::nn::optimizers::adam::DEFAULT_PARAMETERS_TENSORFLOW<T>{
+    struct OPTIMIZER_PARAMETERS: rlt::nn::optimizers::adam::DEFAULT_PARAMETERS_TENSORFLOW<TYPE_POLICY>{
         static constexpr T ALPHA = $OPTIMIZER_ALPHA;
         static constexpr T BETA_1 = $OPTIMIZER_BETA_1;
         static constexpr T BETA_2 = $OPTIMIZER_BETA_2;
@@ -56,5 +57,5 @@ struct PPO_LOOP_CORE_PARAMETERS: rlt::rl::algorithms::ppo::loop::core::DefaultPa
 };
 
 
-template <typename T, typename TI, typename RNG, typename ENVIRONMENT>
-using LOOP_CORE_CONFIG_FACTORY = rlt::rl::algorithms::ppo::loop::core::Config<T, TI, RNG, ENVIRONMENT, PPO_LOOP_CORE_PARAMETERS<T, TI, ENVIRONMENT>>;
\ No newline at end of file
+template <typename TYPE_POLICY, typename TI, typename RNG, typename ENVIRONMENT>
+using LOOP_CORE_CONFIG_FACTORY = rlt::rl::algorithms::ppo::loop::core::Config<TYPE_POLICY, TI, RNG, ENVIRONMENT, PPO_LOOP_CORE_PARAMETERS<TYPE_POLICY, TI, ENVIRONMENT>>;
\ No newline at end of file
diff --git a/rltools/interface/algorithms/sac/template.h b/rltools/interface/algorithms/sac/template.h
index 16a9f1e..bc6ac5f 100644
--- a/rltools/interface/algorithms/sac/template.h
+++ b/rltools/interface/algorithms/sac/template.h
@@ -10,9 +10,10 @@ namespace RL_TOOLS_MODULE_NAME{
     namespace rlt = rl_tools;
 
     // This should stay in sync with the parameters in https://github.com/rl-tools/rl-tools/blob/a98bc461f7ebfed0ba71c653216edec6d9334b78/include/rl_tools/rl/algorithms/sac/loop/core/config.h#L18
-    template <typename T, typename TI, typename ENVIRONMENT>
+    template <typename TYPE_POLICY, typename TI, typename ENVIRONMENT>
     // struct SAC_LOOP_CORE_PARAMETERS: rlt::rl::algorithms::sac::loop::core::DefaultParameters<T, TI, ENVIRONMENT>{
     struct SAC_LOOP_CORE_PARAMETERS{
+        using T = typename TYPE_POLICY::DEFAULT;
         struct SAC_PARAMETERS{
             static constexpr T GAMMA = $GAMMA;
             static constexpr TI ACTOR_BATCH_SIZE = $ACTOR_BATCH_SIZE;
@@ -52,11 +53,14 @@ namespace RL_TOOLS_MODULE_NAME{
         static constexpr bool COLLECT_EPISODE_STATS = $COLLECT_EPISODE_STATS;
         static constexpr TI EPISODE_STATS_BUFFER_SIZE = $EPISODE_STATS_BUFFER_SIZE;
         static constexpr bool SHARED_BATCH = $SHARED_BATCH;
+        static constexpr bool SAMPLE_ENVIRONMENT_PARAMETERS = true;
 
-        using INITIALIZER = rlt::nn::layers::dense::DefaultInitializer<T, TI>;
 
+        using INITIALIZER = rlt::nn::layers::dense::DefaultInitializer<TYPE_POLICY, TI>;
 
-        struct OPTIMIZER_PARAMETERS: rlt::nn::optimizers::adam::DEFAULT_PARAMETERS_TENSORFLOW<T>{
+        using BATCH_SAMPLING_PARAMETERS = rlt::rl::components::off_policy_runner::SequentialBatchParameters<TYPE_POLICY, TI, SAC_PARAMETERS::SEQUENCE_LENGTH, rlt::rl::components::off_policy_runner::SequentialBatchParametersDefault>;
+
+        struct OPTIMIZER_PARAMETERS: rlt::nn::optimizers::adam::DEFAULT_PARAMETERS_TENSORFLOW<TYPE_POLICY>{
             static constexpr T ALPHA = $OPTIMIZER_ALPHA;
             static constexpr T BETA_1 = $OPTIMIZER_BETA_1;
             static constexpr T BETA_2 = $OPTIMIZER_BETA_2;
@@ -70,6 +74,6 @@ namespace RL_TOOLS_MODULE_NAME{
 
     };
 
-    template <typename T, typename TI, typename RNG, typename ENVIRONMENT>
-    using LOOP_CORE_CONFIG_FACTORY = rlt::rl::algorithms::sac::loop::core::Config<T, TI, RNG, ENVIRONMENT, SAC_LOOP_CORE_PARAMETERS<T, TI, ENVIRONMENT>>;
+    template <typename TYPE_POLICY, typename TI, typename RNG, typename ENVIRONMENT>
+    using LOOP_CORE_CONFIG_FACTORY = rlt::rl::algorithms::sac::loop::core::Config<TYPE_POLICY, TI, RNG, ENVIRONMENT, SAC_LOOP_CORE_PARAMETERS<TYPE_POLICY, TI, ENVIRONMENT>>;
 }
\ No newline at end of file
diff --git a/rltools/interface/algorithms/td3/template.h b/rltools/interface/algorithms/td3/template.h
index 717b682..f6b235f 100644
--- a/rltools/interface/algorithms/td3/template.h
+++ b/rltools/interface/algorithms/td3/template.h
@@ -10,8 +10,9 @@ namespace RL_TOOLS_MODULE_NAME{
     namespace rlt = rl_tools;
 
     // This should stay in sync with the parameters in https://github.com/rl-tools/rl-tools/blob/a98bc461f7ebfed0ba71c653216edec6d9334b78/include/rl_tools/rl/algorithms/td3/loop/core/config.h#L18
-    template <typename T, typename TI, typename ENVIRONMENT>
-    struct TD3_LOOP_CORE_PARAMETERS: rlt::rl::algorithms::td3::loop::core::DefaultParameters<T, TI, ENVIRONMENT>{
+    template <typename TYPE_POLICY, typename TI, typename ENVIRONMENT>
+    struct TD3_LOOP_CORE_PARAMETERS: rlt::rl::algorithms::td3::loop::core::DefaultParameters<TYPE_POLICY, TI, ENVIRONMENT>{
+        using T = typename TYPE_POLICY::DEFAULT;
         struct TD3_PARAMETERS{
             static constexpr T GAMMA = $GAMMA;
             static constexpr TI ACTOR_BATCH_SIZE = $ACTOR_BATCH_SIZE;
@@ -45,8 +46,13 @@ namespace RL_TOOLS_MODULE_NAME{
         static constexpr TI EPISODE_STATS_BUFFER_SIZE = $EPISODE_STATS_BUFFER_SIZE;
         static constexpr T EXPLORATION_NOISE = $EXPLORATION_NOISE;
         static constexpr bool SHARED_BATCH = $SHARED_BATCH;
+        static constexpr bool SAMPLE_ENVIRONMENT_PARAMETERS = true;
 
-        struct OPTIMIZER_PARAMETERS: rlt::nn::optimizers::adam::DEFAULT_PARAMETERS_TENSORFLOW<T>{
+        using INITIALIZER = rlt::nn::layers::dense::DefaultInitializer<TYPE_POLICY, TI>;
+
+        using BATCH_SAMPLING_PARAMETERS = rlt::rl::components::off_policy_runner::SequentialBatchParameters<TYPE_POLICY, TI, TD3_PARAMETERS::SEQUENCE_LENGTH, rlt::rl::components::off_policy_runner::SequentialBatchParametersDefault>;
+
+        struct OPTIMIZER_PARAMETERS: rlt::nn::optimizers::adam::DEFAULT_PARAMETERS_TENSORFLOW<TYPE_POLICY>{
             static constexpr T ALPHA = $OPTIMIZER_ALPHA;
             static constexpr T BETA_1 = $OPTIMIZER_BETA_1;
             static constexpr T BETA_2 = $OPTIMIZER_BETA_2;
@@ -55,6 +61,6 @@ namespace RL_TOOLS_MODULE_NAME{
         };
     };
 
-    template <typename T, typename TI, typename RNG, typename ENVIRONMENT>
-    using LOOP_CORE_CONFIG_FACTORY = rlt::rl::algorithms::td3::loop::core::Config<T, TI, RNG, ENVIRONMENT, TD3_LOOP_CORE_PARAMETERS<T, TI, ENVIRONMENT>>;
+    template <typename TYPE_POLICY, typename TI, typename RNG, typename ENVIRONMENT>
+    using LOOP_CORE_CONFIG_FACTORY = rlt::rl::algorithms::td3::loop::core::Config<TYPE_POLICY, TI, RNG, ENVIRONMENT, TD3_LOOP_CORE_PARAMETERS<TYPE_POLICY, TI, ENVIRONMENT>>;
 }
\ No newline at end of file
diff --git a/rltools/interface/inference/inference.cpp b/rltools/interface/inference/inference.cpp
index 8a16447..33b93ad 100644
--- a/rltools/interface/inference/inference.cpp
+++ b/rltools/interface/inference/inference.cpp
@@ -4,6 +4,7 @@
 
 #include <rl_tools/nn/operations_generic.h>
 #include <rl_tools/nn/layers/dense/operations_generic.h>
+#include <rl_tools/nn/layers/td3_sampling/operations_generic.h>
 #include <rl_tools/nn/layers/standardize/operations_generic.h>
 #include <rl_tools/nn/layers/sample_and_squash/operations_generic.h>
 #include <rl_tools/nn_models/mlp_unconditional_stddev/operations_generic.h>
@@ -19,17 +20,21 @@ namespace rlt = rl_tools;
 
 
 using DEVICE = rlt::devices::DEVICE_FACTORY<>;
+using RNG = DEVICE::SPEC::RANDOM::ENGINE<>;
 
 DEVICE device;
-auto rng = rlt::random::default_engine(typename DEVICE::SPEC::RANDOM{});
+RNG rng;
 using MODEL_TYPE = decltype(policy::module);
-using T = typename MODEL_TYPE::T;
+using T = typename MODEL_TYPE::TYPE_POLICY::DEFAULT;
+using TYPE_POLICY = rlt::numeric_types::Policy<T>;
 using TI = typename DEVICE::index_t;
 typename MODEL_TYPE::template Buffer<1> buffer;
 bool initialized = false;
+constexpr TI seed = 0x1337;
 
 void init(){
     if(!initialized){
+        rlt::init(device, rng, seed);
         rlt::malloc(device, buffer);
     }
 }
diff --git a/rltools/interface/python_environment/operations_cpu.h b/rltools/interface/python_environment/operations_cpu.h
index 0824aef..aab6215 100644
--- a/rltools/interface/python_environment/operations_cpu.h
+++ b/rltools/interface/python_environment/operations_cpu.h
@@ -67,7 +67,7 @@ namespace rl_tools{
         // delete env.environment;
     }
     template<typename DEVICE, typename SPEC>
-    static void init(DEVICE& device, PythonEnvironment<SPEC>& env, const typename PythonEnvironment<SPEC>::Parameters& parameters){}
+    static void init(DEVICE& device, PythonEnvironment<SPEC>& env){}
     template<typename DEVICE, typename SPEC>
     static void initial_parameters(DEVICE& device, const PythonEnvironment<SPEC>& env, typename PythonEnvironment<SPEC>::Parameters& parameters){}
     template<typename DEVICE, typename SPEC, typename RNG>
diff --git a/rltools/interface/training/training.cpp b/rltools/interface/training/training.cpp
index 54b7951..ee409e6 100644
--- a/rltools/interface/training/training.cpp
+++ b/rltools/interface/training/training.cpp
@@ -24,6 +24,9 @@
 #include <rl_tools/nn_models/sequential/operations_generic.h>
 #include <rl_tools/nn/optimizers/adam/operations_generic.h>
 
+#include <rl_tools/containers/matrix/persist_code.h>
+#include <rl_tools/containers/tensor/persist_code.h>
+#include <rl_tools/numeric_types/persist_code.h>
 #include <rl_tools/nn/optimizers/adam/instance/persist_code.h>
 #include <rl_tools/nn/parameters/persist_code.h>
 #include <rl_tools/nn/layers/sample_and_squash/persist_code.h>
@@ -48,10 +51,11 @@ namespace RL_TOOLS_MODULE_NAME{
 #ifdef RL_TOOLS_FORCE_BLAS
     static_assert(DEVICE::DEVICE_ID == rlt::devices::DeviceId::CPU_MKL || DEVICE::DEVICE_ID == rlt::devices::DeviceId::CPU_ACCELERATE || DEVICE::DEVICE_ID == rlt::devices::DeviceId::CPU_OPENBLAS);
 #endif
-    using RNG = decltype(rlt::random::default_engine(typename DEVICE::SPEC::RANDOM{}));
+    using RNG = DEVICE::SPEC::RANDOM::ENGINE<>;
     using TI = typename DEVICE::index_t;
 
     using T = RL_TOOLS_DTYPE;
+    using TYPE_POLICY = rlt::numeric_types::Policy<T>;
 
 
     #ifdef RL_TOOLS_USE_PYTHON_ENVIRONMENT
@@ -71,11 +75,12 @@ namespace RL_TOOLS_MODULE_NAME{
 
 
 
-    using LOOP_CORE_CONFIG = LOOP_CORE_CONFIG_FACTORY<T, TI, RNG, ENVIRONMENT>;
+    using LOOP_CORE_CONFIG = LOOP_CORE_CONFIG_FACTORY<TYPE_POLICY, TI, RNG, ENVIRONMENT>;
 
 
 
     DEVICE device;
+    RNG rng;
 
     #ifdef RL_TOOLS_ENABLE_EVALUATION
     constexpr bool ENABLE_EVALUATION = true;
@@ -86,7 +91,7 @@ namespace RL_TOOLS_MODULE_NAME{
     constexpr TI PARAMETER_NUM_EVALUATION_EPISODES = RL_TOOLS_NUM_EVALUATION_EPISODES;
 
     template <typename NEXT>
-    struct LOOP_EVAL_PARAMETERS: rlt::rl::loop::steps::evaluation::Parameters<T, TI, NEXT>{
+    struct LOOP_EVAL_PARAMETERS: rlt::rl::loop::steps::evaluation::Parameters<TYPE_POLICY, TI, NEXT>{
         static constexpr TI EVALUATION_INTERVAL = PARAMETER_EVALUATION_INTERVAL;
         static constexpr TI NUM_EVALUATION_EPISODES = PARAMETER_NUM_EVALUATION_EPISODES;
         static constexpr TI N_EVALUATIONS = NEXT::CORE_PARAMETERS::STEP_LIMIT / EVALUATION_INTERVAL;
@@ -124,6 +129,7 @@ namespace RL_TOOLS_MODULE_NAME{
         EVALUATION_ACTOR_BUFFER evaluation_actor_buffer;
         bool evaluation_actor_synced = false;
         State(TI seed){
+            rlt::init(device, rng, seed);
             rlt::malloc(device, static_cast<LOOP_STATE&>(*this));
             rlt::init(device, static_cast<LOOP_STATE&>(*this), seed);
             rlt::malloc(device, evaluation_actor);
diff --git a/rltools/src/link_mkl.py b/rltools/src/link_mkl.py
index 6eb16ca..8230dbb 100644
--- a/rltools/src/link_mkl.py
+++ b/rltools/src/link_mkl.py
@@ -35,6 +35,8 @@ def link_mkl():
                 if not found:
                     print(f"MKL library {lib} not found.")
                     required_libraries_paths = None
+            if "MKL_NUM_THREADS" not in os.environ:
+                os.environ["MKL_NUM_THREADS"] = "1"
             required_libraries_paths_absolute = [str(p.locate()) for p in required_libraries_paths]
             required_libraries_search_paths = [str(os.path.dirname(p)) for p in required_libraries_paths_absolute]
             flags += [
diff --git a/tests/test_pytorch_import.py b/tests/test_pytorch_import.py
index f848656..b8ac6b3 100644
--- a/tests/test_pytorch_import.py
+++ b/tests/test_pytorch_import.py
@@ -1,18 +1,17 @@
 import os
+import pytest
 os.environ["RL_TOOLS_IGNORE_TORCH_WARNING"] = "1"
 
-if not "RL_TOOLS_DISABLE_PYTORCH_IMPORT_TEST" in os.environ:
+@pytest.mark.skip
+def test_pytorch_import():
+    if "RL_TOOLS_DISABLE_PYTORCH_IMPORT_TEST" in os.environ:
+        return
     import torch
     import torch.nn as nn
     from rltools import CACHE_PATH
     from rltools.onnx import load_mlp, evaluate
     from rltools.onnx import render
     from rltools import load_checkpoint_from_path
-
-@pytest.mark.skip
-def test_pytorch_import():
-    if "RL_TOOLS_DISABLE_PYTORCH_IMPORT_TEST" in os.environ:
-        return
     torch.manual_seed(0)
 
     n_input_features = 4