From 6313e5b0705e655242459cf257f077b89b9c74cc Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 28 Jun 2026 08:19:29 +0100
Subject: [PATCH 01/19] Add construction of functional for constructing
 log-probabilities

---
 src/tfrlrl/policies/base.py                   |  23 +++-
 src/tfrlrl/policies/dense_neural_network.py   |  37 +++++-
 src/tfrlrl/policies/linear_soft_max.py        |  30 ++++-
 .../policies/test_dense_neural_network.py     | 110 ++++++++++++++++++
 tests/tfrlrl/policies/test_linear_soft_max.py |  74 ++++++++++++
 5 files changed, 271 insertions(+), 3 deletions(-)

diff --git a/src/tfrlrl/policies/base.py b/src/tfrlrl/policies/base.py
index 97d1fc4..df80bca 100644
--- a/src/tfrlrl/policies/base.py
+++ b/src/tfrlrl/policies/base.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Iterator, Union
+from typing import Any, Callable, Dict, Iterator, Tuple, Union
 
 import gymnasium as gym
 import numpy as np
@@ -187,3 +187,24 @@ def calculate_log_probabilities(self, observations: np.ndarray, actions: np.ndar
 
         """
         ...
+
+    @abstractmethod
+    def make_log_prob_fn(self, observations: np.ndarray, actions: np.ndarray) -> Tuple[Callable[[Dict], Tensor], Dict]:
+        """
+        Construct a PyTorch functional to calculate the log-probabilities of the given state-action pairs.
+
+        Construct a PyTorch functional that takes the policy parameters as inputs and returns the log-probabilites
+        of the given state-action pairs. This functional can be used in various functionality, such as in the
+        construction of the Jacboian of the policy.
+
+        Args:
+            observations: A NumPy array of the observations for which to calculate the log-probabilities.
+            actions: A NumPy array of the actions for which to calculate the log-probabilities (of the corresponding
+            observations).
+
+        Returns:
+            A tuple consisting of a function that takes the policy parameters as inputs and returns the
+            log-probabilities of the given state-action pairs and a dictionary of policy parameters.
+
+        """
+        ...
diff --git a/src/tfrlrl/policies/dense_neural_network.py b/src/tfrlrl/policies/dense_neural_network.py
index 5c4ed0e..bdccdd5 100644
--- a/src/tfrlrl/policies/dense_neural_network.py
+++ b/src/tfrlrl/policies/dense_neural_network.py
@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from typing import List, Union
+from typing import Callable, Dict, List, Tuple, Union
 
 import gymnasium as gym
 import numpy as np
@@ -11,6 +11,7 @@
     tensor,
 )
 from torch.distributions.normal import Normal
+from torch.func import functional_call
 
 from tfrlrl.policies.base import BasePyTorchPolicy, PolicyException
 
@@ -138,3 +139,37 @@ def calculate_log_probabilities(self, observations: np.ndarray, actions: np.ndar
             .log_prob(tensor(actions))
             .T
         )
+
+    def make_log_prob_fn(self, observations: np.ndarray, actions: np.ndarray) -> Tuple[Callable[[Dict], Tensor], Dict]:
+        """
+        Construct a PyTorch functional to calculate the log-probabilities of the given state-action pairs.
+
+        Construct a PyTorch functional that takes the policy parameters as inputs and returns the log-probabilites
+        of the given state-action pairs. This functional can be used in various functionality, such as in the
+        construction of the Jacboian of the policy.
+
+        Args:
+            observations: A NumPy array of the observations for which to calculate the log-probabilities.
+            actions: A NumPy array of the actions for which to calculate the log-probabilities (of the corresponding
+            observations).
+
+        Returns:
+            A tuple consisting of a function that takes the policy parameters as inputs and returns the
+            log-probabilities of the given state-action pairs and a dictionary of policy parameters.
+
+        """
+        params = dict(self.network.named_parameters())
+
+        def log_prob_fn(params):
+            action_means, action_stddevs = functional_call(self.network, params, (tensor(observations).T,))
+
+            return (
+                Normal(
+                    action_means[:, 0] + self.eps,
+                    action_stddevs[:, 0] + self.eps,
+                )
+                .log_prob(tensor(actions))
+                .T
+            )
+
+        return log_prob_fn, params
diff --git a/src/tfrlrl/policies/linear_soft_max.py b/src/tfrlrl/policies/linear_soft_max.py
index 020924c..0525eae 100644
--- a/src/tfrlrl/policies/linear_soft_max.py
+++ b/src/tfrlrl/policies/linear_soft_max.py
@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from typing import Union
+from typing import Callable, Dict, Tuple, Union
 
 import gymnasium as gym
 import numpy as np
@@ -9,6 +9,7 @@
     tensor,
 )
 from torch.distributions.multinomial import Categorical
+from torch.func import functional_call
 
 from tfrlrl.features.base import FeatureFunction
 from tfrlrl.policies.base import (
@@ -146,3 +147,30 @@ def calculate_log_probabilities(self, observations: np.ndarray, actions: np.ndar
 
         """
         return self.calculate_action_distribution(observations).log_prob(tensor(actions))
+
+    def make_log_prob_fn(self, observations: np.ndarray, actions: np.ndarray) -> Tuple[Callable[[Dict], Tensor], Dict]:
+        """
+        Construct a PyTorch functional to calculate the log-probabilities of the given state-action pairs.
+
+        Construct a PyTorch functional that takes the policy parameters as inputs and returns the log-probabilites
+        of the given state-action pairs. This functional can be used in various functionality, such as in the
+        construction of the Jacboian of the policy.
+
+        Args:
+            observations: A NumPy array of the observations for which to calculate the log-probabilities.
+            actions: A NumPy array of the actions for which to calculate the log-probabilities (of the corresponding
+            observations).
+
+        Returns:
+            A tuple consisting of a function that takes the policy parameters as inputs and returns the
+            log-probabilities of the given state-action pairs and a dictionary of policy parameters.
+
+        """
+        params = dict(self.network.named_parameters())
+
+        def log_prob_fn(params):
+            network_input = self.construct_network_input(observations)
+            probs = functional_call(self.network, params, (network_input,)).squeeze()
+            return Categorical(probs=probs).log_prob(tensor(actions))
+
+        return log_prob_fn, params
diff --git a/tests/tfrlrl/policies/test_dense_neural_network.py b/tests/tfrlrl/policies/test_dense_neural_network.py
index 1bb94f9..e1a707c 100644
--- a/tests/tfrlrl/policies/test_dense_neural_network.py
+++ b/tests/tfrlrl/policies/test_dense_neural_network.py
@@ -172,3 +172,113 @@ def test_calculate_log_probabilities_from_dense_network_policy_multiple_observat
         assert log_probabilities.shape == (n_observations, 1)
     else:
         assert log_probabilities.shape == (n_observations,)
+
+
+@pytest.mark.parametrize('env_id', ['InvertedPendulum-v5'])
+def test_calculate_log_probabilities_from_functional_single_observation(env_id):
+    """
+    Test make_log_prob_fn with a single observation with a DenseNetworkPolicy.
+
+    Args:
+        env_id: The environment I.D. from which to sample episodes.
+
+    """
+    env = gym.make(env_id)
+
+    hidden_space_dims = [16, 32]
+    policy = DenseNetworkPolicy(
+        env_id=env_id,
+        hidden_space_dims=hidden_space_dims,
+    )
+
+    observation = env.observation_space.sample()[..., np.newaxis]
+    action = env.action_space.sample()
+    log_probability = (
+        policy.calculate_log_probabilities(
+            observation,
+            action,
+        )
+        .detach()
+        .numpy()
+    )
+    log_prob_fn, params = policy.make_log_prob_fn(
+        observation,
+        action,
+    )
+    log_probability_from_functional = log_prob_fn(params).detach().numpy()
+    np.testing.assert_allclose(
+        log_probability,
+        log_probability_from_functional,
+        rtol=1e-6,
+        atol=1e-9,
+    )
+
+
+@pytest.mark.parametrize(
+    'env_id, extend_actions',
+    [
+        (
+            'InvertedPendulum-v5',
+            True,
+        ),
+        (
+            'InvertedPendulum-v5',
+            False,
+        ),
+    ],
+)
+@given(
+    n_observations=st.integers(min_value=1, max_value=20),
+)
+@settings(deadline=None)
+def test_calculate_log_probabilities_from_functional_multiple_observations(env_id, n_observations, extend_actions):
+    """
+    Test make_log_prob_fn with multiple observations with a DenseNetworkPolicy.
+
+    Args:
+        env_id: The environment I.D. from which to sample episodes.
+        n_observations: The number of observations to sample.
+        extend_actions: A Boolean indicating whether to extend the diemsions of the actions.
+
+    """
+    env = gym.make(env_id)
+
+    hidden_space_dims = [16, 32]
+    policy = DenseNetworkPolicy(
+        env_id=env_id,
+        hidden_space_dims=hidden_space_dims,
+    )
+
+    observations = np.concatenate(
+        [env.observation_space.sample()[..., np.newaxis] for _ in range(n_observations)],
+        axis=1,
+    )
+    if extend_actions:
+        actions = np.concatenate(
+            [env.action_space.sample()[..., np.newaxis] for _ in range(n_observations)],
+            axis=1,
+        )
+    else:
+        actions = np.concatenate(
+            [env.action_space.sample() for _ in range(n_observations)],
+        )
+
+    log_probabilities = (
+        policy.calculate_log_probabilities(
+            observations,
+            actions,
+        )
+        .detach()
+        .numpy()
+    )
+    log_prob_fn, params = policy.make_log_prob_fn(
+        observations,
+        actions,
+    )
+    log_probability_from_functional = log_prob_fn(params).detach().numpy()
+    np.testing.assert_allclose(
+        log_probabilities,
+        log_probability_from_functional,
+        rtol=1e-6,
+        atol=1e-9,
+    )
diff --git a/tests/tfrlrl/policies/test_linear_soft_max.py b/tests/tfrlrl/policies/test_linear_soft_max.py
index daa6b2f..fc41f75 100644
--- a/tests/tfrlrl/policies/test_linear_soft_max.py
+++ b/tests/tfrlrl/policies/test_linear_soft_max.py
@@ -247,6 +247,80 @@ def test_linear_softmax_policy_log_probabilities(env_id: str, n_observations: in
         )
 
 
+@pytest.mark.parametrize('env_id', ['CliffWalking-v1'])
+@given(
+    observation=st.integers(min_value=0, max_value=47),
+    action=st.integers(min_value=0, max_value=3),
+    seed=st.integers(min_value=0, max_value=10000),
+)
+@settings(deadline=None)
+def test_make_log_prob_fn_single_observation(env_id: str, observation: int, action: int, seed: int):
+    """
+    Test make_log_prob_fn returns function that can be used to calculate log-probability of given state-action pair.
+
+    Args:
+        env_id: The Gymnasium environment ID with a discrete action space.
+        observation: A valid observation (state) from the environment.
+        action: A valid action from the environment.
+        seed: Random seed for generating softmax parameters.
+
+    """
+    env = gym.make(env_id)
+    np.random.seed(seed)
+    feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+    policy = LinearSoftMax(env_id, feature_fn)
+
+    action = np.array([action])
+    observation = np.array([observation])
+
+    log_probability = policy.calculate_log_probabilities(observation, action).detach().numpy()
+    log_prob_fn, params = policy.make_log_prob_fn(observation, action)
+    log_probability_from_functional = log_prob_fn(params).detach().numpy()
+    np.testing.assert_allclose(
+        log_probability,
+        log_probability_from_functional,
+        rtol=1e-6,
+        atol=1e-9,
+    )
+
+
+@pytest.mark.parametrize('env_id', ['CliffWalking-v1'])
+@given(
+    n_observations=st.integers(min_value=2, max_value=100),
+    seed=st.integers(min_value=0, max_value=10000),
+)
+@settings(deadline=None)
+def test_make_log_prob_fn_multiple_observations(env_id: str, n_observations: int, seed: int):
+    """
+    Test make_log_prob_fn returns function that can be used to calculate log-probabilities of given state-action pairs.
+
+    Args:
+        env_id: The Gymnasium environment ID with a discrete action space.
+        n_observations: The number of observations for which to calculate the log-probabilities.
+        seed: Random seed for generating softmax parameters.
+
+    """
+    env = gym.make(env_id)
+    np.random.seed(seed)
+    feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+    policy = LinearSoftMax(env_id, feature_fn)
+
+    observations = np.random.randint(low=0, high=47, size=(1, n_observations))
+    actions = np.random.randint(low=0, high=3, size=(1, n_observations))
+
+    log_probabilities = policy.calculate_log_probabilities(observations, actions).detach().numpy()
+    log_prob_fn, params = policy.make_log_prob_fn(observations, actions)
+    log_probability_from_functional = log_prob_fn(params).detach().numpy()
+
+    for n in range(n_observations):
+        np.testing.assert_allclose(
+            log_probabilities[0, n],
+            log_probability_from_functional[0, n],
+            rtol=1e-6,
+            atol=1e-9,
+        )
+
+
 @pytest.mark.parametrize('env_id', ['CliffWalking-v1'])
 @given(
     observation=st.integers(min_value=0, max_value=47),

From 2599eee457ee4b1af55813993ad1befff5c0d058 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 28 Jun 2026 18:26:03 +0100
Subject: [PATCH 02/19] Add calculation of Jacobian

---
 src/tfrlrl/policies/base.py                   |  21 ++++
 .../policies/test_dense_neural_network.py     | 119 ++++++++++++++++++
 tests/tfrlrl/policies/test_linear_soft_max.py | 114 +++++++++++++++++
 3 files changed, 254 insertions(+)

diff --git a/src/tfrlrl/policies/base.py b/src/tfrlrl/policies/base.py
index df80bca..727a8ef 100644
--- a/src/tfrlrl/policies/base.py
+++ b/src/tfrlrl/policies/base.py
@@ -7,6 +7,7 @@
     Tensor,
     nn,
 )
+from torch.func import jacfwd
 
 
 class PolicyException(Exception):
@@ -208,3 +209,23 @@ def make_log_prob_fn(self, observations: np.ndarray, actions: np.ndarray) -> Tup
 
         """
         ...
+
+    def calculate_jacobian(self, observations: np.ndarray, actions: np.ndarray) -> Dict[str, Tensor]:
+        """
+        Calculate the Jacobian of the log-probabilites at the given state-action pairs.
+
+        This function calculates the Jacobian of the log-probabilities of the policy for the given
+        state-action pairs. The Jacobian is given in the form of a dictionary, with the keys corresponding
+        to the different parameters of the policy (network).
+
+        Args:
+            observations: A NumPy array of the observations for which to calculate the log-probabilities.
+            actions: A NumPy array of the actions for which to calculate the log-probabilities (of the corresponding
+            observations).
+
+        Returns:
+            A dictionary mapping the parameter name to the (rows of the) Jacobian corresponding to that parameter.
+
+        """
+        log_prob_fn, params = self.make_log_prob_fn(observations, actions)
+        return jacfwd(log_prob_fn)(params)
diff --git a/tests/tfrlrl/policies/test_dense_neural_network.py b/tests/tfrlrl/policies/test_dense_neural_network.py
index e1a707c..c87ae8f 100644
--- a/tests/tfrlrl/policies/test_dense_neural_network.py
+++ b/tests/tfrlrl/policies/test_dense_neural_network.py
@@ -1,3 +1,5 @@
+import copy
+
 import gymnasium as gym
 import numpy as np
 import pytest
@@ -282,3 +284,120 @@ def test_calculate_log_probabilities_from_functional_multiple_observations(env_i
         rtol=1e-6,
         atol=1e-9,
     )
+
+
+@pytest.mark.parametrize('env_id', ['InvertedPendulum-v5'])
+@given(
+    seed=st.integers(min_value=0, max_value=10000),
+)
+@settings(deadline=None)
+def test_dense_neural_network_calculate_jacobian_single_observation(env_id: str, seed: int):
+    """
+    Test calculate_jacobian for a single state-action pair in the DenseNetworkPolicy.
+
+    Args:
+        env_id: The Gymnasium environment ID with a continuous action space.
+        seed: Random seed for sampling observations, actions and generating network parameters.
+
+    """
+    env = gym.make(env_id)
+    np.random.seed(seed)
+
+    hidden_space_dims = [4, 4]
+    policy = DenseNetworkPolicy(env_id=env_id, hidden_space_dims=hidden_space_dims)
+
+    observation = env.observation_space.sample()[..., np.newaxis]
+    action = env.action_space.sample()
+
+    jacobian = policy.calculate_jacobian(observation, action)
+
+    eps = 0.001
+    policy_dict = copy.deepcopy(policy.get_state())
+
+    for parameter_name, parameter_value in policy_dict.items():
+        param_shape = tuple(parameter_value.shape)
+        df_finite_diffs = np.zeros(param_shape)
+
+        for idx in np.ndindex(param_shape):
+            new_policy_dict_plus = copy.deepcopy(policy_dict)
+            new_policy_dict_minus = copy.deepcopy(policy_dict)
+
+            new_policy_dict_plus[parameter_name][idx] += eps
+            policy.set_state(new_policy_dict_plus)
+            log_prob_plus = policy.calculate_log_probabilities(observation, action).item()
+
+            new_policy_dict_minus[parameter_name][idx] -= eps
+            policy.set_state(new_policy_dict_minus)
+            log_prob_minus = policy.calculate_log_probabilities(observation, action).item()
+
+            df_finite_diffs[idx] = 0.5 * (log_prob_plus - log_prob_minus) / eps
+
+        np.testing.assert_almost_equal(
+            jacobian[parameter_name].squeeze().detach().numpy(),
+            df_finite_diffs.squeeze(),
+            decimal=2,
+        )
+
+
+@pytest.mark.parametrize('env_id', ['InvertedPendulum-v5'])
+@given(
+    n_observations=st.integers(min_value=2, max_value=10),
+    seed=st.integers(min_value=0, max_value=10000),
+)
+@settings(deadline=None)
+def test_dense_neural_network_calculate_jacobian_multiple_observations(env_id: str, n_observations: int, seed: int):
+    """
+    Test calculate_jacobian for multiple state-action pairs in the DenseNetworkPolicy.
+
+    Args:
+        env_id: The Gymnasium environment ID with a continuous action space.
+        n_observations: The number of observations to sample from the environment.
+        seed: Random seed for sampling observations, actions and generating network parameters.
+
+    """
+    env = gym.make(env_id)
+    np.random.seed(seed)
+
+    hidden_space_dims = [4, 4]
+    policy = DenseNetworkPolicy(env_id=env_id, hidden_space_dims=hidden_space_dims)
+
+    observations = np.concatenate(
+        [env.observation_space.sample()[..., np.newaxis] for _ in range(n_observations)],
+        axis=1,
+    )
+    actions = np.concatenate(
+        [env.action_space.sample() for _ in range(n_observations)],
+    )
+
+    jacobian = policy.calculate_jacobian(observations, actions)
+
+    eps = 0.01
+    policy_dict = copy.deepcopy(policy.get_state())
+
+    for parameter_name, parameter_value in policy_dict.items():
+        param_shape = tuple(parameter_value.shape)
+        df_finite_diffs = np.zeros((n_observations, *param_shape))
+
+        for i in range(n_observations):
+            obs_i = observations[:, i][..., np.newaxis]
+            act_i = actions[i]
+
+            for idx in np.ndindex(param_shape):
+                new_policy_dict_plus = copy.deepcopy(policy_dict)
+                new_policy_dict_minus = copy.deepcopy(policy_dict)
+
+                new_policy_dict_plus[parameter_name][idx] += eps
+                policy.set_state(new_policy_dict_plus)
+                log_prob_plus = policy.calculate_log_probabilities(obs_i, act_i).item()
+
+                new_policy_dict_minus[parameter_name][idx] -= eps
+                policy.set_state(new_policy_dict_minus)
+                log_prob_minus = policy.calculate_log_probabilities(obs_i, act_i).item()
+
+                df_finite_diffs[(i, *idx)] = 0.5 * (log_prob_plus - log_prob_minus) / eps
+
+        np.testing.assert_almost_equal(
+            jacobian[parameter_name].squeeze().detach().numpy(),
+            df_finite_diffs.squeeze(),
+            decimal=2,
+        )
diff --git a/tests/tfrlrl/policies/test_linear_soft_max.py b/tests/tfrlrl/policies/test_linear_soft_max.py
index fc41f75..16c2dd0 100644
--- a/tests/tfrlrl/policies/test_linear_soft_max.py
+++ b/tests/tfrlrl/policies/test_linear_soft_max.py
@@ -461,3 +461,117 @@ def test_linear_softmax_policy_log_probabilities_derivatives_multiple_observatio
         df_log_probs_finite_diffs,
         decimal=2,
     )
+
+
+@pytest.mark.parametrize('env_id', ['CliffWalking-v1'])
+@given(
+    observation=st.integers(min_value=0, max_value=47),
+    action=st.integers(min_value=0, max_value=3),
+    seed=st.integers(min_value=0, max_value=10000),
+)
+@settings(deadline=None)
+def test_linear_softmax_calculate_jacobian_single_observation(env_id: str, observation: int, action: int, seed: int):
+    """
+    Test calculate_jacobian for a single state-action pair in the linear soft-max policy.
+
+    Args:
+        env_id: The Gymnasium environment ID with a discrete action space.
+        observation: A valid observation (state) from the environment.
+        action: A valid action from the environment.
+        seed: Random seed for generating softmax parameters.
+
+    """
+    env = gym.make(env_id)
+    np.random.seed(seed)
+    feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+    policy = LinearSoftMax(env_id, feature_fn)
+
+    action = np.array([action])
+    observation = np.array([observation])
+
+    jacobian = policy.calculate_jacobian(observation, action)
+
+    eps = 0.001
+    policy_dict = copy.deepcopy(policy.get_state())
+
+    for parameter_name, parameter_value in policy_dict.items():
+        param_shape = tuple(parameter_value.shape)
+        df_finite_diffs = np.zeros(param_shape)
+
+        for idx in np.ndindex(param_shape):
+            new_policy_dict_plus = copy.deepcopy(policy_dict)
+            new_policy_dict_minus = copy.deepcopy(policy_dict)
+
+            new_policy_dict_plus[parameter_name][idx] += eps
+            policy.set_state(new_policy_dict_plus)
+            log_probability_plus = policy.calculate_log_probabilities(observation, action).item()
+
+            new_policy_dict_minus[parameter_name][idx] -= eps
+            policy.set_state(new_policy_dict_minus)
+            log_probability_minus = policy.calculate_log_probabilities(observation, action).item()
+
+            df_finite_diffs[idx] = 0.5 * (log_probability_plus - log_probability_minus) / eps
+
+        np.testing.assert_almost_equal(
+            jacobian[parameter_name].squeeze().detach().numpy(),
+            df_finite_diffs.squeeze(),
+            decimal=2,
+        )
+
+
+@pytest.mark.parametrize('env_id', ['CliffWalking-v1'])
+@given(
+    n_observations=st.integers(min_value=2, max_value=10),
+    seed=st.integers(min_value=0, max_value=10000),
+)
+@settings(deadline=None)
+def test_linear_softmax_calculate_jacobian_multiple_observations(env_id: str, n_observations: int, seed: int):
+    """
+    Test calculate_jacobian for a multiple state-action pairs in the linear soft-max policy.
+
+    Args:
+        env_id: The Gymnasium environment ID with a discrete action space.
+        n_observations: The number of observations to sample from the environment.
+        seed: Random seed for generating softmax parameters.
+
+    """
+    env = gym.make(env_id)
+    np.random.seed(seed)
+    feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+    policy = LinearSoftMax(env_id, feature_fn)
+
+    observations = np.random.randint(low=0, high=47, size=(1, n_observations))
+    actions = np.random.randint(low=0, high=3, size=(1, n_observations))
+
+    jacobian = policy.calculate_jacobian(observations, actions)
+
+    eps = 0.01
+    policy_dict = copy.deepcopy(policy.get_state())
+
+    for parameter_name, parameter_value in policy_dict.items():
+        param_shape = tuple(parameter_value.shape)
+        df_finite_diffs = np.zeros((n_observations, *param_shape))
+
+        for i in range(n_observations):
+            obs_i = observations[:, i]
+            act_i = actions[:, i]
+
+            for idx in np.ndindex(param_shape):
+                new_policy_dict_plus = copy.deepcopy(policy_dict)
+                new_policy_dict_minus = copy.deepcopy(policy_dict)
+
+                new_policy_dict_plus[parameter_name][idx] += eps
+                policy.set_state(new_policy_dict_plus)
+                log_probability_plus = policy.calculate_log_probabilities(obs_i, act_i).item()
+
+                new_policy_dict_minus[parameter_name][idx] -= eps
+                policy.set_state(new_policy_dict_minus)
+                log_probability_minus = policy.calculate_log_probabilities(obs_i, act_i).item()
+
+                df_finite_diffs[(i, *idx)] = 0.5 * (log_probability_plus - log_probability_minus) / eps
+
+        np.testing.assert_almost_equal(
+            jacobian[parameter_name].squeeze().detach().numpy(),
+            df_finite_diffs.squeeze(),
+            decimal=2,
+        )

From 3ee0679d88a5f37cc5b6ebf1041f90e5f2947f2f Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 28 Jun 2026 19:00:11 +0100
Subject: [PATCH 03/19] Add conjugate gradient algorithm

---
 .../optimisation/conjugate_gradients.py       | 54 +++++++++++++++++++
 1 file changed, 54 insertions(+)
 create mode 100644 src/tfrlrl/optimisation/conjugate_gradients.py

diff --git a/src/tfrlrl/optimisation/conjugate_gradients.py b/src/tfrlrl/optimisation/conjugate_gradients.py
new file mode 100644
index 0000000..4176224
--- /dev/null
+++ b/src/tfrlrl/optimisation/conjugate_gradients.py
@@ -0,0 +1,54 @@
+from typing import Callable, Optional
+
+import numpy as np
+
+
+def calculate_conjugate_gradient(
+    mat_v_mult_fn: Callable[[np.ndarray], np.ndarray], b: np.ndarray, n_iters: Optional[int] = None, tol: float = 1e-10
+):
+    """
+    Perform the conjugate-gradient algorithm to find the (approximate) solution, x, of Ax = b.
+
+    Args:
+        mat_v_mult_fn: A callable that takes an input vector, x, and returns the matrix multiplication of A
+        with x. The purpose of making this a callable is to allow the user to perform conjugate--gradients without
+        ever having to materialise the matrix, A.
+        b: The target solution vector of the matrix-vector multiplication in the conjugate-gradient algorithm.
+        n_iters: The maximum number of iterations to perform in conjugate-gradients.
+        tol: The toleration below which the conjugate-gradient algorithm will terminate.
+
+    Returns:
+        A NumPy array, x, that forms the approximate solution of the equation, Ax = b.
+
+    Raises:
+        If the b not a vector, i.e. either a one-dimensional NumPy array or a two dimensional array with a single
+        element in the last dimension, then a ValueError will be thrown.
+
+    """
+    if b.ndim > 2 or (b.ndim == 2 and b.shape[1] > 1):
+        raise ValueError('Input vector to conjugate gradients needs to be a 1-dimensional array: %s', b.shape)
+
+    if n_iters is None:
+        n_iters = b.flatten().shape[0]
+
+    p = b.copy()
+    r = b.copy()
+    x = np.zeros_like(b)
+    rdotr = np.inner(r, r)
+
+    for _ in range(n_iters):
+        z = mat_v_mult_fn(p)
+
+        v = rdotr / p.T.dot(z)
+        x += v * p
+        r -= v * z
+
+        newrdotr = r.dot(r)
+        mu = newrdotr / rdotr
+        p = r + mu * p
+
+        rdotr = newrdotr
+        if rdotr < tol:
+            break
+
+    return x

From ee2a11bb647b61962add38befb33ce644cd8804c Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 28 Jun 2026 19:34:06 +0100
Subject: [PATCH 04/19] Add tests for cg algorithm

---
 CLAUDE.md                                     |  14 +-
 tests/tfrlrl/optimisation/__init__.py         |   0
 .../optimisation/test_conjugate_gradients.py  | 127 ++++++++++++++++++
 3 files changed, 139 insertions(+), 2 deletions(-)
 create mode 100644 tests/tfrlrl/optimisation/__init__.py
 create mode 100644 tests/tfrlrl/optimisation/test_conjugate_gradients.py

diff --git a/CLAUDE.md b/CLAUDE.md
index d299bbf..f5643b3 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -118,13 +118,23 @@ Command-line interface tools for common tasks:
   - Inherits logging configuration from centralized settings (no --verbose flag)
 - Entry point registered as `tfrlrl-sample` console script in pyproject.toml
 
-## Code Style Requirements
+## Code Style \& Conventions 
+
+### Requirements
 
 - Maximum line length: 120 characters
 - Maximum complexity: 8 (flake8)
 - isort configuration: line_length=120, multi_line_output=3, include_trailing_comma=true
 - Docstrings required for classes and functions (excluding D100, D104)
-- Double quotes for strings (enforced by flake8-quotes)
+- Single quotes for strings (enforced by ruff)
+- Type checking enforced through MyPy
+
+### Conventions
+
+- The test-suite must follow the same structure as the main project. 
+- Tests should not be collected in classes, but written as individual functions. 
+- There is a strong preference to using PyTest parametrize to avoid re-writing duplicate test code to handle different test cases. 
+- The use of the notation, A @ v, should be avoided. An explicit np.matmul or np.dot should be used instead.  
 
 ## Testing Requirements
 
diff --git a/tests/tfrlrl/optimisation/__init__.py b/tests/tfrlrl/optimisation/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/tfrlrl/optimisation/test_conjugate_gradients.py b/tests/tfrlrl/optimisation/test_conjugate_gradients.py
new file mode 100644
index 0000000..b167354
--- /dev/null
+++ b/tests/tfrlrl/optimisation/test_conjugate_gradients.py
@@ -0,0 +1,127 @@
+import numpy as np
+import pytest
+from hypothesis import given, settings
+from hypothesis import strategies as st
+
+from tfrlrl.optimisation.conjugate_gradients import calculate_conjugate_gradient
+
+
+def _make_spd_matrix(n: int) -> np.ndarray:
+    """Construct a random symmetric positive definite matrix of size (n, n)."""
+    A = np.random.randn(n, n)
+    return np.matmul(A.T, A) + n * np.eye(n)
+
+
+@given(seed=st.integers(min_value=0, max_value=10000))
+@settings(deadline=None)
+def test_exact_solution_with_n_iters_none(seed: int):
+    """
+    Test that calculate_conjugate_gradient returns the exact solution when n_iters is None.
+
+    When n_iters is None the algorithm defaults to n iterations (the dimension of b), which is sufficient
+    for exact convergence on a symmetric positive definite system.
+
+    Args:
+        seed: Random seed for generating the matrix and target vector.
+
+    """
+    np.random.seed(seed)
+    n = 5
+    A = _make_spd_matrix(n)
+    b = np.random.randn(n)
+
+    def mat_v_mult_fn(v):
+        return np.matmul(A, v)
+
+    x = calculate_conjugate_gradient(mat_v_mult_fn, b)
+
+    np.testing.assert_allclose(np.matmul(A, x), b, rtol=1e-5, atol=1e-5)
+
+
+@given(seed=st.integers(min_value=0, max_value=10000))
+@settings(deadline=None)
+def test_exact_solution_with_n_iters_equal_to_dimension(seed: int):
+    """
+    Test that calculate_conjugate_gradient returns the exact solution when n_iters equals the dimension.
+
+    This verifies that an explicit n_iters equal to the problem dimension is equivalent to n_iters=None.
+
+    Args:
+        seed: Random seed for generating the matrix and target vector.
+
+    """
+    np.random.seed(seed)
+    n = 5
+    A = _make_spd_matrix(n)
+    b = np.random.randn(n)
+
+    def mat_v_mult_fn(v):
+        return np.matmul(A, v)
+
+    x = calculate_conjugate_gradient(mat_v_mult_fn, b, n_iters=n)
+
+    np.testing.assert_allclose(np.matmul(A, x), b, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize('n_iters', [1, 2, 3])
+@given(seed=st.integers(min_value=0, max_value=10000))
+@settings(deadline=None)
+def test_approximate_solution_with_n_iters_less_than_dimension(n_iters: int, seed: int):
+    """
+    Test that calculate_conjugate_gradient reduces the A-norm error when n_iters is less than the dimension.
+
+    With fewer iterations than the problem dimension, the algorithm returns an approximate solution. This
+    test verifies that the output has the correct shape and that the A-norm error is reduced relative to
+    the initial A-norm error (starting from x=0), which is guaranteed by the monotone convergence property
+    of conjugate gradients on symmetric positive definite systems.
+
+    Args:
+        n_iters: The number of conjugate gradient iterations to perform (less than the problem dimension).
+        seed: Random seed for generating the matrix and target vector.
+
+    """
+    np.random.seed(seed)
+    n = 5
+    A = _make_spd_matrix(n)
+    b = np.random.randn(n)
+
+    def mat_v_mult_fn(v):
+        return np.matmul(A, v)
+
+    x = calculate_conjugate_gradient(mat_v_mult_fn, b, n_iters=n_iters)
+    x_exact = np.linalg.solve(A, b)
+
+    assert x.shape == b.shape
+
+    e_approx = x - x_exact
+    e_init = -x_exact
+    assert np.dot(e_approx, np.matmul(A, e_approx)) <= np.dot(e_init, np.matmul(A, e_init))
+
+
+@pytest.mark.parametrize(
+    'b_shape',
+    [
+        (5, 2),
+        (5, 3),
+        (5, 1, 1),
+        (2, 3, 5),
+    ],
+)
+def test_raises_value_error_for_invalid_b_shape(b_shape: tuple):
+    """
+    Test that calculate_conjugate_gradient raises a ValueError for invalid b shapes.
+
+    The function requires b to be a 1-D array or a 2-D column vector. Any other shape — including 2-D
+    arrays with more than one column and arrays with more than two dimensions — should raise a ValueError.
+
+    Args:
+        b_shape: The invalid shape for the input vector b.
+
+    """
+    b = np.random.randn(*b_shape)
+
+    def mat_v_mult_fn(v):
+        return v
+
+    with pytest.raises(ValueError):
+        calculate_conjugate_gradient(mat_v_mult_fn, b)

From 33b11835d50bf2378728efc0a3523cbf2d496769 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Mon, 29 Jun 2026 07:18:53 +0100
Subject: [PATCH 05/19] fix style

---
 src/tfrlrl/policies/utils.py           |  22 +++
 src/tfrlrl/training_algorithms/tnpg.py | 197 +++++++++++++++++++++++++
 2 files changed, 219 insertions(+)
 create mode 100644 src/tfrlrl/policies/utils.py
 create mode 100644 src/tfrlrl/training_algorithms/tnpg.py

diff --git a/src/tfrlrl/policies/utils.py b/src/tfrlrl/policies/utils.py
new file mode 100644
index 0000000..2f7d237
--- /dev/null
+++ b/src/tfrlrl/policies/utils.py
@@ -0,0 +1,22 @@
+from typing import Dict
+
+from torch import (
+    Tensor,
+    cat,
+)
+
+
+def flatten_tensor_dict(x: Dict[str, Tensor]) -> Tensor:
+    """
+    Flatten the given dictionary of PyTorch tensors into a signle PyTorch tensor.
+
+    This function takes a dictionary of PyTorch tensors and flattens them into a single tensor.
+
+    Args:
+        x: The dictionary of tensors to be flattened.
+
+    Returns:
+        The flattened tensor.
+
+    """
+    return cat([j.flatten(start_dim=1) for j in x.values()], dim=1)
diff --git a/src/tfrlrl/training_algorithms/tnpg.py b/src/tfrlrl/training_algorithms/tnpg.py
new file mode 100644
index 0000000..dacbb8b
--- /dev/null
+++ b/src/tfrlrl/training_algorithms/tnpg.py
@@ -0,0 +1,197 @@
+"""Stochastic Gradient Descent training algorithm."""
+
+import logging
+from typing import Callable, Optional, Union
+
+import numpy as np
+import ray
+from torch import (
+    no_grad,
+    sum,
+    tensor,
+)
+from torch.optim import (
+    Optimizer,
+)
+
+from tfrlrl import settings
+from tfrlrl.baselines.linear import Baseline
+from tfrlrl.data_models.reward_models import AverageEpisodicReward, DiscountedReward
+from tfrlrl.data_models.statistics import BaseStatistics
+from tfrlrl.optimisation.conjugate_gradients import calculate_conjugate_gradient
+from tfrlrl.policies.base import BasePyTorchPolicy
+from tfrlrl.policies.utils import flatten_tensor_dict
+from tfrlrl.sampling.episodic_sampler import (
+    EpisodicSampler,
+    RayEpisodicSampler,
+)
+from tfrlrl.sampling.statistics_collection import EpisocidPolicyGradientStatisticsCollector
+
+logger = logging.getLogger(__name__)
+
+
+def calculate_steepest_gradient_direction(
+    policy: BasePyTorchPolicy, statistics: BaseStatistics, optimizer: Optimizer
+) -> np.ndarray:
+    """
+    Calculate the direction of steepest gradient ascent of the policy.
+
+    This function calculates the steepest gradient ascent direction, i.e., standard policy gradients, and returns
+    it in the form of a vector (i.e. NumPy array).
+
+    Args:
+        policy: The policy for which the Fisher Infromation matrix is to be calculated.
+        statistics: The statistics over which the expectation within the Fisher Information matrix calculation
+        is to be performed.
+        optimizer: A PyTorch optimizer class that will be used to calculate the policy gradient.
+
+    Returns:
+        A NumPy array containing the policy gradient.
+
+    """
+    optimizer.zero_grad()
+    log_probabilities = policy.calculate_log_probabilities(
+        observations=statistics.observations,
+        actions=statistics.actions,
+    )
+    loss = -sum(log_probabilities * tensor(statistics.total_expected_rewards))
+    loss.backward()
+
+    # TODO: Convert sgd to a NumPy array
+    sgd = flatten_tensor_dict(
+        {name: param.grad for name, param in policy.network.named_parameters()},
+    )
+    return sgd
+
+
+def construct_fim_vector_product_fn(
+    policy: BasePyTorchPolicy, statistics: BaseStatistics
+) -> Callable[[np.ndarray], np.ndarray]:
+    """
+    Construct function for calculating the product of the Fisher Information matrix with a vector.
+
+    This function constructs a function that when called with a vector (NumPy array) will return the
+    product of the Fisher Information matrix with that vector. The Fisher Information matrix is calculated
+    w.r.t. the given policy, while the expetation in Fisher Information matrix is calculated over the sample
+      state-action pairs in the given statistics.
+
+    Args:
+        policy: The policy for which the Fisher Infromation matrix is to be calculated.
+        statistics: The statistics over which the expectation within the Fisher Information matrix calculation
+        is to be performed.
+
+    Returns:
+        A callable that takes an input vector as an argument and returns the product of the Fisher Information
+        matrix with that vector.
+
+    """
+    jacobian = policy.calculate_jacobian(
+        observations=statistics.observations,
+        actions=statistics.actions,
+    )
+    # TODO: Convert Jacobian to a NumPy array
+    jacobian = flatten_tensor_dict(jacobian)
+
+    def calculate_fim_vector_product(v: np.ndarray):
+        return np.matmul(jacobian.T, np.matmul(jacobian, v))
+
+    return calculate_fim_vector_product
+
+
+def train_policy_gradient(
+    env_id: str,
+    policy: BasePyTorchPolicy,
+    n_iterations: int,
+    n_episodes: int,
+    optimizer: Optimizer,
+    n_samplers: int = 1,
+    baseline: Optional[Baseline] = None,
+    reward_model: Optional[Union[AverageEpisodicReward, DiscountedReward]] = None,
+    n_iteration_logging: int = 10,
+    **kwargs,
+) -> BasePyTorchPolicy:
+    """
+    Train a policy using stochastic gradient ascent on the policy gradient.
+
+    Args:
+        env_id: Gymnasium environment ID (e.g., CartPole-v1, MountainCar-v0).
+        policy: The policy to train. Must have get_parameters() and set_parameters() methods.
+        n_iterations: The number of policy updates to perform.
+        n_episodes: The number of episodes to sample during each policy update.
+        optimizer: An instance of a PyTorch optimizer class that will be used to optimise the policy.
+        n_samplers: The number of samplers to used to sample from the environment.
+        baseline: An instance of a baseline class, if one is given.
+        reward_model: The reward model to use when computing total expected rewards. Defaults to
+        AverageEpisodicReward if not specified.
+        n_iteration_logging: The number of algorithm iterations between logging algorithm performance.
+        kwargs: Additional keyword arguments to pass to the EpisodicSampler (e.g., is_slippery).
+
+    Returns:
+        The trained policy.
+
+    """
+    statistics_collector = EpisocidPolicyGradientStatisticsCollector(
+        env_id,
+        baseline=baseline,
+        reward_model=reward_model,
+    )
+
+    if n_samplers > 1:
+        if not ray.is_initialized():
+            ray.init(
+                num_cpus=settings.ray_cpu,
+                ignore_reinit_error=True,
+            )
+            logger.info('Ray initialized for %s CPUs', settings.ray_cpu)
+        sampler: Union[RayEpisodicSampler, EpisodicSampler] = RayEpisodicSampler(
+            n_samplers=n_samplers,
+            env_id=env_id,
+            n_episodes=n_episodes,
+            policy=policy,
+            statistics_collector=statistics_collector,
+            **kwargs,
+        )
+    else:
+        sampler = EpisodicSampler(
+            env_id=env_id,
+            n_episodes=n_episodes,
+            policy=policy,
+            statistics_collector=statistics_collector,
+            **kwargs,
+        )
+
+    for n in range(n_iterations):
+        with no_grad():
+            statistics = sampler.sample()
+
+        sgd = calculate_steepest_gradient_direction(
+            policy=policy,
+            statistics=statistics,
+            optimizer=optimizer,
+        )
+
+        calculate_conjugate_gradient(
+            mat_v_mult_fn=construct_fim_vector_product_fn(
+                policy=policy,
+                statistics=statistics,
+            ),
+            b=sgd,
+            n_iters=10,
+        )
+
+        # optimizer.step()
+
+        if n % n_iteration_logging == 0:
+            logger.info('Policy update: %s', n)
+            logger.info('Average total episodic reward: %s', np.average(statistics.total_reward))
+
+        if baseline:
+            baseline.fit(statistics.baseline_features, statistics.baseline_targets)
+
+        sampler.reset()
+        sampler.update(
+            policy_state_dict=policy.get_state(),
+            baseline_state_dict=baseline.get_state() if baseline else None,
+        )
+
+    return policy

From 69588bfc6ec992c0391f2a0beba9b1f055c95726 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sat, 4 Jul 2026 11:09:15 +0100
Subject: [PATCH 06/19] Add tests for calculate_steepest_gradient_direction and
 fix flatten_tensor_dict for vector outputs

---
 src/tfrlrl/policies/base.py                   |  3 ++-
 src/tfrlrl/policies/dense_neural_network.py   | 27 +++++++------------
 src/tfrlrl/policies/linear_soft_max.py        |  3 ++-
 src/tfrlrl/policies/utils.py                  |  5 ++--
 src/tfrlrl/training_algorithms/tnpg.py        | 14 +++-------
 .../policies/test_dense_neural_network.py     |  2 +-
 6 files changed, 22 insertions(+), 32 deletions(-)

diff --git a/src/tfrlrl/policies/base.py b/src/tfrlrl/policies/base.py
index 727a8ef..8b2d290 100644
--- a/src/tfrlrl/policies/base.py
+++ b/src/tfrlrl/policies/base.py
@@ -184,7 +184,8 @@ def calculate_log_probabilities(self, observations: np.ndarray, actions: np.ndar
             observations).
 
         Returns:
-            A PyTorch Tensor containing the log-probabilities of the given (observation, action) pairs.
+            A PyTorch Tensor containing the log-probabilities of the given (observation, action) pairs. The shape of
+            the output is expected to be either (n_observation) or (1, n_observations).
 
         """
         ...
diff --git a/src/tfrlrl/policies/dense_neural_network.py b/src/tfrlrl/policies/dense_neural_network.py
index bdccdd5..e9f2250 100644
--- a/src/tfrlrl/policies/dense_neural_network.py
+++ b/src/tfrlrl/policies/dense_neural_network.py
@@ -127,18 +127,15 @@ def calculate_log_probabilities(self, observations: np.ndarray, actions: np.ndar
             actions: The actions taken in the given observation state.
 
         Returns:
-            A Tensor containing the log-probabilities of the policy for the given observation-action pairs.
+            A Tensor containing the log-probabilities of the policy for the given observation-action pairs. The shape
+            of the output is expected to be either (n_observation) or (1, n_observations).
 
         """
         action_means, action_stddevs = self.network(tensor(observations).T)
-        return (
-            Normal(
-                action_means[:, 0] + self.eps,
-                action_stddevs[:, 0] + self.eps,
-            )
-            .log_prob(tensor(actions))
-            .T
-        )
+        return Normal(
+            action_means[:, 0] + self.eps,
+            action_stddevs[:, 0] + self.eps,
+        ).log_prob(tensor(actions))
 
     def make_log_prob_fn(self, observations: np.ndarray, actions: np.ndarray) -> Tuple[Callable[[Dict], Tensor], Dict]:
         """
@@ -163,13 +160,9 @@ def make_log_prob_fn(self, observations: np.ndarray, actions: np.ndarray) -> Tup
         def log_prob_fn(params):
             action_means, action_stddevs = functional_call(self.network, params, (tensor(observations).T,))
 
-            return (
-                Normal(
-                    action_means[:, 0] + self.eps,
-                    action_stddevs[:, 0] + self.eps,
-                )
-                .log_prob(tensor(actions))
-                .T
-            )
+            return Normal(
+                action_means[:, 0] + self.eps,
+                action_stddevs[:, 0] + self.eps,
+            ).log_prob(tensor(actions))
 
         return log_prob_fn, params
diff --git a/src/tfrlrl/policies/linear_soft_max.py b/src/tfrlrl/policies/linear_soft_max.py
index 0525eae..4d1f0a5 100644
--- a/src/tfrlrl/policies/linear_soft_max.py
+++ b/src/tfrlrl/policies/linear_soft_max.py
@@ -143,7 +143,8 @@ def calculate_log_probabilities(self, observations: np.ndarray, actions: np.ndar
             actions: The actions taken in the given observation state.
 
         Returns:
-            return: The log-probabilities of the policy for the given observation-action pairs.
+            return: The log-probabilities of the policy for the given observation-action pairs. The shape of
+            the output is expected to be either (n_observation) or (1, n_observations).
 
         """
         return self.calculate_action_distribution(observations).log_prob(tensor(actions))
diff --git a/src/tfrlrl/policies/utils.py b/src/tfrlrl/policies/utils.py
index 2f7d237..8082cef 100644
--- a/src/tfrlrl/policies/utils.py
+++ b/src/tfrlrl/policies/utils.py
@@ -16,7 +16,8 @@ def flatten_tensor_dict(x: Dict[str, Tensor]) -> Tensor:
         x: The dictionary of tensors to be flattened.
 
     Returns:
-        The flattened tensor.
+        The flattened tensor. The flattened tensor will have a shape of (1, n_elems) in which n_elems
+        is the number of elements across the different tensors in the input dictionary.
 
     """
-    return cat([j.flatten(start_dim=1) for j in x.values()], dim=1)
+    return cat([j.flatten() for j in x.values()])[None, :]
diff --git a/src/tfrlrl/training_algorithms/tnpg.py b/src/tfrlrl/training_algorithms/tnpg.py
index dacbb8b..0ab3b7f 100644
--- a/src/tfrlrl/training_algorithms/tnpg.py
+++ b/src/tfrlrl/training_algorithms/tnpg.py
@@ -5,11 +5,7 @@
 
 import numpy as np
 import ray
-from torch import (
-    no_grad,
-    sum,
-    tensor,
-)
+from torch import Tensor, no_grad, sum, tensor
 from torch.optim import (
     Optimizer,
 )
@@ -32,7 +28,7 @@
 
 def calculate_steepest_gradient_direction(
     policy: BasePyTorchPolicy, statistics: BaseStatistics, optimizer: Optimizer
-) -> np.ndarray:
+) -> Tensor:
     """
     Calculate the direction of steepest gradient ascent of the policy.
 
@@ -46,7 +42,7 @@ def calculate_steepest_gradient_direction(
         optimizer: A PyTorch optimizer class that will be used to calculate the policy gradient.
 
     Returns:
-        A NumPy array containing the policy gradient.
+        A PyTorch tensor containing the policy gradient.
 
     """
     optimizer.zero_grad()
@@ -57,11 +53,9 @@ def calculate_steepest_gradient_direction(
     loss = -sum(log_probabilities * tensor(statistics.total_expected_rewards))
     loss.backward()
 
-    # TODO: Convert sgd to a NumPy array
-    sgd = flatten_tensor_dict(
+    return flatten_tensor_dict(
         {name: param.grad for name, param in policy.network.named_parameters()},
     )
-    return sgd
 
 
 def construct_fim_vector_product_fn(
diff --git a/tests/tfrlrl/policies/test_dense_neural_network.py b/tests/tfrlrl/policies/test_dense_neural_network.py
index c87ae8f..272b385 100644
--- a/tests/tfrlrl/policies/test_dense_neural_network.py
+++ b/tests/tfrlrl/policies/test_dense_neural_network.py
@@ -171,7 +171,7 @@ def test_calculate_log_probabilities_from_dense_network_policy_multiple_observat
     )
 
     if extend_actions:
-        assert log_probabilities.shape == (n_observations, 1)
+        assert log_probabilities.shape == (1, n_observations)
     else:
         assert log_probabilities.shape == (n_observations,)
 

From 56b0bd55f10708f24283d990591990d1265ff9af Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 5 Jul 2026 07:59:45 +0100
Subject: [PATCH 07/19] extend flatten_tensor_dict to support dimension
 argument

---
 src/tfrlrl/policies/utils.py           | 18 ++++++++++++++----
 src/tfrlrl/training_algorithms/tnpg.py | 14 ++++++++------
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/src/tfrlrl/policies/utils.py b/src/tfrlrl/policies/utils.py
index 8082cef..4c0d7ce 100644
--- a/src/tfrlrl/policies/utils.py
+++ b/src/tfrlrl/policies/utils.py
@@ -6,7 +6,7 @@
 )
 
 
-def flatten_tensor_dict(x: Dict[str, Tensor]) -> Tensor:
+def flatten_tensor_dict(x: Dict[str, Tensor], dim: int = 0) -> Tensor:
     """
     Flatten the given dictionary of PyTorch tensors into a signle PyTorch tensor.
 
@@ -14,10 +14,20 @@ def flatten_tensor_dict(x: Dict[str, Tensor]) -> Tensor:
 
     Args:
         x: The dictionary of tensors to be flattened.
+        dim: The starting dimension on which to flatten and concatenate the tensors.
 
     Returns:
-        The flattened tensor. The flattened tensor will have a shape of (1, n_elems) in which n_elems
-        is the number of elements across the different tensors in the input dictionary.
+        The flattened tensor. The flattened tensor will have the following shape:
+                    (start_dims, n_elems / n_start_dim_elems).
+        start_dims is dimensions of the dimensions precceding, n_start_dim_elems is the number of
+        dimensions in the starting dimensions (or 1 when dim is zero) and  n_elems is the number of
+        elements across the different tensors in the input dictionary.
+
+    Raises:
+        RuntimeError: When the starting dimensions of the tensors in x, i.e. the dimensions less than
+        dim, then a RuntimeError will be thrown by the call to cat.
 
     """
-    return cat([j.flatten() for j in x.values()])[None, :]
+    if dim == 0:
+        return cat([j.flatten(start_dim=dim) for j in x.values()], dim=dim)[None, :]
+    return cat([j.flatten(start_dim=dim) for j in x.values()], dim=dim)
diff --git a/src/tfrlrl/training_algorithms/tnpg.py b/src/tfrlrl/training_algorithms/tnpg.py
index 0ab3b7f..03ee208 100644
--- a/src/tfrlrl/training_algorithms/tnpg.py
+++ b/src/tfrlrl/training_algorithms/tnpg.py
@@ -13,7 +13,6 @@
 from tfrlrl import settings
 from tfrlrl.baselines.linear import Baseline
 from tfrlrl.data_models.reward_models import AverageEpisodicReward, DiscountedReward
-from tfrlrl.data_models.statistics import BaseStatistics
 from tfrlrl.optimisation.conjugate_gradients import calculate_conjugate_gradient
 from tfrlrl.policies.base import BasePyTorchPolicy
 from tfrlrl.policies.utils import flatten_tensor_dict
@@ -21,13 +20,16 @@
     EpisodicSampler,
     RayEpisodicSampler,
 )
-from tfrlrl.sampling.statistics_collection import EpisocidPolicyGradientStatisticsCollector
+from tfrlrl.sampling.statistics_collection import (
+    EpisocidPolicyGradientStatisticsCollector,
+    EpisodePolicyGradientStatistics,
+)
 
 logger = logging.getLogger(__name__)
 
 
 def calculate_steepest_gradient_direction(
-    policy: BasePyTorchPolicy, statistics: BaseStatistics, optimizer: Optimizer
+    policy: BasePyTorchPolicy, statistics: EpisodePolicyGradientStatistics, optimizer: Optimizer
 ) -> Tensor:
     """
     Calculate the direction of steepest gradient ascent of the policy.
@@ -59,7 +61,7 @@ def calculate_steepest_gradient_direction(
 
 
 def construct_fim_vector_product_fn(
-    policy: BasePyTorchPolicy, statistics: BaseStatistics
+    policy: BasePyTorchPolicy, statistics: EpisodePolicyGradientStatistics
 ) -> Callable[[np.ndarray], np.ndarray]:
     """
     Construct function for calculating the product of the Fisher Information matrix with a vector.
@@ -84,10 +86,10 @@ def construct_fim_vector_product_fn(
         actions=statistics.actions,
     )
     # TODO: Convert Jacobian to a NumPy array
-    jacobian = flatten_tensor_dict(jacobian)
+    jacobian_matrix = flatten_tensor_dict(jacobian)
 
     def calculate_fim_vector_product(v: np.ndarray):
-        return np.matmul(jacobian.T, np.matmul(jacobian, v))
+        return np.matmul(jacobian_matrix.T, np.matmul(jacobian_matrix, v))
 
     return calculate_fim_vector_product
 

From 0092659aa59fe8915b3d8d662273604c9820589f Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 5 Jul 2026 08:01:20 +0100
Subject: [PATCH 08/19] extend flatten_tensor_dict to support dimension
 argument

---
 tests/tfrlrl/policies/test_utils.py           | 132 ++++++++++++++++++
 tests/tfrlrl/training_algorithms/test_tnpg.py | 124 ++++++++++++++++
 2 files changed, 256 insertions(+)
 create mode 100644 tests/tfrlrl/policies/test_utils.py
 create mode 100644 tests/tfrlrl/training_algorithms/test_tnpg.py

diff --git a/tests/tfrlrl/policies/test_utils.py b/tests/tfrlrl/policies/test_utils.py
new file mode 100644
index 0000000..8c5862f
--- /dev/null
+++ b/tests/tfrlrl/policies/test_utils.py
@@ -0,0 +1,132 @@
+import gymnasium as gym
+import numpy as np
+import pytest
+from hypothesis import given, settings
+from hypothesis import strategies as st
+
+from tfrlrl.features.onehot import OneHotFeatureFunction
+from tfrlrl.policies.dense_neural_network import DenseNetworkPolicy
+from tfrlrl.policies.linear_soft_max import LinearSoftMax
+from tfrlrl.policies.utils import flatten_tensor_dict
+
+
+@pytest.mark.parametrize('env_id', ['FrozenLake-v1', 'InvertedPendulum-v5'])
+def test_flatten_tensor_dict(env_id: str):
+    """
+    Test flatten_tensor_dict function on policy parameters.
+
+    This function tests that flatten_tensor_dict flattens policy parameters and maintains
+    the parameter ordering.
+
+    Args:
+        env_id: The environment I.D. on which to run the test. This is used to determine the policy on which
+        to perform the test.
+
+    """
+    if env_id == 'FrozenLake-v1':
+        env = gym.make(env_id)
+        feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+        policy = LinearSoftMax(env_id, feature_fn)
+    elif env_id == 'InvertedPendulum-v5':
+        policy = DenseNetworkPolicy(
+            env_id=env_id,
+            hidden_space_dims=[16, 32],
+        )
+    else:
+        raise ValueError('Unexpected environment: %s', env_id)
+
+    flattend_tensor = flatten_tensor_dict(
+        {name: param for name, param in policy.network.named_parameters()},
+    )
+
+    indx = 0
+    for _, param in policy.network.named_parameters():
+        for p_indx in np.ndindex(param.shape):
+            assert param[p_indx] == flattend_tensor[0, indx]
+            indx += 1
+
+    assert indx == flattend_tensor.shape[1]
+
+
+@pytest.mark.parametrize(
+    'env_id, extend_actions',
+    [
+        (
+            'CliffWalking-v1',
+            False,
+        ),
+        (
+            'InvertedPendulum-v5',
+            True,
+        ),
+        (
+            'InvertedPendulum-v5',
+            False,
+        ),
+    ],
+)
+@given(
+    n_observations=st.integers(min_value=5, max_value=10),
+)
+@settings(deadline=None)
+def test_flatten_tensor_dict_jacobian(env_id: str, extend_actions, n_observations: int):
+    """
+    Test flatten_tensor_dict function on the policy Jacobian.
+
+    This function tests that flatten_tensor_dict flattens policy Jacobian and maintains
+    the parameter ordering.
+
+    Args:
+        env_id: The environment I.D. on which to run the test. This is used to determine the policy on which
+        to perform the test.
+        extend_actions: A Boolean indicating whether to extend the action dimensions. This is used to capture
+        the case of a 1-dimensional conintuous action space, in which the actions could be stored as either
+        a (n_observations) or an (1, n_observations) array.
+        n_observations: The number of observations.
+
+    """
+    env = gym.make(env_id)
+
+    if env_id == 'CliffWalking-v1':
+        feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+        policy = LinearSoftMax(env_id, feature_fn)
+        observations = np.random.randint(low=0, high=47, size=(1, n_observations))
+        actions = np.random.randint(low=0, high=3, size=(1, n_observations))
+    elif env_id == 'InvertedPendulum-v5':
+        policy = DenseNetworkPolicy(
+            env_id=env_id,
+            hidden_space_dims=[16, 32],
+        )
+        observations = np.concatenate(
+            [env.observation_space.sample()[..., np.newaxis] for _ in range(n_observations)],
+            axis=1,
+        )
+        if extend_actions:
+            actions = np.concatenate(
+                [env.action_space.sample()[..., np.newaxis] for _ in range(n_observations)],
+                axis=1,
+            )
+        else:
+            actions = np.concatenate(
+                [env.action_space.sample() for _ in range(n_observations)],
+            )
+    else:
+        raise ValueError('Unexpected environment: %s', env_id)
+
+    jacobian = policy.calculate_jacobian(observations, actions)
+    flattend_jacobian = flatten_tensor_dict(
+        jacobian,
+        dim=actions.ndim,
+    )
+
+    indx = 0
+    for _, jac_param in jacobian.items():
+        param_shape = tuple(jac_param.shape[actions.ndim :])
+        for p_indx in np.ndindex(param_shape):
+            for i in range(n_observations):
+                if actions.ndim == 1:
+                    assert jac_param[(i,) + p_indx] == flattend_jacobian[i, indx]
+                else:
+                    assert jac_param[(0, i) + p_indx] == flattend_jacobian[0, i, indx]
+            indx += 1
+    assert indx == flattend_jacobian.shape[actions.ndim]
diff --git a/tests/tfrlrl/training_algorithms/test_tnpg.py b/tests/tfrlrl/training_algorithms/test_tnpg.py
new file mode 100644
index 0000000..e66c1c7
--- /dev/null
+++ b/tests/tfrlrl/training_algorithms/test_tnpg.py
@@ -0,0 +1,124 @@
+import copy
+from typing import Tuple
+
+import gymnasium as gym
+import numpy as np
+import pytest
+from hypothesis import given, settings
+from hypothesis import strategies as st
+from torch import (
+    Tensor,
+    tensor,
+)
+from torch import (
+    sum as torch_sum,
+)
+from torch.optim import (
+    SGD,
+)
+
+from tfrlrl.features.onehot import OneHotFeatureFunction
+from tfrlrl.policies.dense_neural_network import DenseNetworkPolicy
+from tfrlrl.policies.linear_soft_max import LinearSoftMax
+from tfrlrl.sampling.episodic_sampler import EpisodicSampler
+from tfrlrl.sampling.statistics_collection import EpisocidPolicyGradientStatisticsCollector
+from tfrlrl.training_algorithms.tnpg import calculate_steepest_gradient_direction
+
+
+@pytest.mark.parametrize(
+    'env_id, expected_shape',
+    [
+        ('FrozenLake-v1', (1, 49)),
+        ('InvertedPendulum-v5', (1, 690)),
+    ],
+)
+@given(
+    n_episodes=st.integers(min_value=10, max_value=100),
+)
+@settings(deadline=None)
+def test_calculate_steepest_gradient_direction(env_id: str, expected_shape: Tuple[int], n_episodes: int):
+    """
+    Test that calculate_steepest_gradient_direction returns the correct policy gradient.
+
+    The steepest gradient direction is verified against central finite differences of the loss function
+    L(θ) = -sum_t log π_θ(a_t | s_t) * R_t with respect to each policy parameter element.
+
+    Args:
+        env_id: The Gymnasium environment ID with a discrete action space.
+        expected_shape: The expected shape of the gradient.
+        n_episodes: The number of episodes to sample when estimating the gradient.
+
+    """
+    env = gym.make(env_id)
+    statistics_collector = EpisocidPolicyGradientStatisticsCollector(env_id)
+
+    if env_id == 'FrozenLake-v1':
+        feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+        policy = LinearSoftMax(env_id, feature_fn)
+        sampler = EpisodicSampler(
+            env_id=env_id,
+            n_episodes=n_episodes,
+            policy=policy,
+            statistics_collector=statistics_collector,
+            is_slippery=False,
+        )
+    elif env_id == 'InvertedPendulum-v5':
+        policy = DenseNetworkPolicy(
+            env_id=env_id,
+            hidden_space_dims=[16, 32],
+        )
+        sampler = EpisodicSampler(
+            env_id=env_id,
+            n_episodes=n_episodes,
+            policy=policy,
+            statistics_collector=statistics_collector,
+        )
+    else:
+        raise ValueError('Unexpected environment: %s', env_id)
+
+    alpha = 0.1
+    optimizer = SGD(policy.get_parameters(), lr=alpha)
+    statistics = sampler.sample()
+    sgd = calculate_steepest_gradient_direction(
+        policy,
+        statistics,
+        optimizer,
+    )
+    assert isinstance(sgd, Tensor)
+    assert sgd.shape == expected_shape
+
+    eps = 0.001
+    policy_dict = copy.deepcopy(policy.get_state())
+
+    fd_grad_parts = []
+    for param_name, param_value in policy_dict.items():
+        param_shape = tuple(param_value.shape)
+        fd_grad = np.zeros(param_shape)
+
+        for idx in np.ndindex(param_shape):
+            dict_plus = copy.deepcopy(policy_dict)
+            dict_minus = copy.deepcopy(policy_dict)
+
+            dict_plus[param_name][idx] += eps
+            policy.set_state(dict_plus)
+            log_probs_plus = policy.calculate_log_probabilities(statistics.observations, statistics.actions)
+            loss_plus = (-torch_sum(log_probs_plus * tensor(statistics.total_expected_rewards))).item()
+
+            dict_minus[param_name][idx] -= eps
+            policy.set_state(dict_minus)
+            log_probs_minus = policy.calculate_log_probabilities(statistics.observations, statistics.actions)
+            loss_minus = (-torch_sum(log_probs_minus * tensor(statistics.total_expected_rewards))).item()
+
+            fd_grad[idx] = (loss_plus - loss_minus) / (2 * eps)
+
+        fd_grad_parts.append(fd_grad.flatten())
+
+    policy.set_state(policy_dict)
+
+    fd_gradient = np.concatenate(fd_grad_parts)[np.newaxis, :]
+
+    np.testing.assert_almost_equal(
+        sgd.detach().numpy(),
+        fd_gradient,
+        decimal=2,
+    )

From 5e4e74001953be970412f0bfea2623627a5127a8 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 12 Jul 2026 09:11:43 +0100
Subject: [PATCH 09/19] Add tests for FIM-vector product

---
 src/tfrlrl/policies/utils.py                  |   8 +-
 src/tfrlrl/training_algorithms/tnpg.py        |  29 +++--
 tests/tfrlrl/policies/test_utils.py           |   4 +-
 tests/tfrlrl/training_algorithms/test_tnpg.py | 103 ++++++++++++++++--
 4 files changed, 121 insertions(+), 23 deletions(-)

diff --git a/src/tfrlrl/policies/utils.py b/src/tfrlrl/policies/utils.py
index 4c0d7ce..59d8934 100644
--- a/src/tfrlrl/policies/utils.py
+++ b/src/tfrlrl/policies/utils.py
@@ -19,15 +19,13 @@ def flatten_tensor_dict(x: Dict[str, Tensor], dim: int = 0) -> Tensor:
     Returns:
         The flattened tensor. The flattened tensor will have the following shape:
                     (start_dims, n_elems / n_start_dim_elems).
-        start_dims is dimensions of the dimensions precceding, n_start_dim_elems is the number of
-        dimensions in the starting dimensions (or 1 when dim is zero) and  n_elems is the number of
-        elements across the different tensors in the input dictionary.
+        start_dims is the dimensions preceeding dim, which is empty when dim is zero, n_start_dim_elems
+        is the number of dimensions in the starting dimensions (or 1 when dim is zero) and  n_elems is
+        the number of elements across the different tensors in the input dictionary.
 
     Raises:
         RuntimeError: When the starting dimensions of the tensors in x, i.e. the dimensions less than
         dim, then a RuntimeError will be thrown by the call to cat.
 
     """
-    if dim == 0:
-        return cat([j.flatten(start_dim=dim) for j in x.values()], dim=dim)[None, :]
     return cat([j.flatten(start_dim=dim) for j in x.values()], dim=dim)
diff --git a/src/tfrlrl/training_algorithms/tnpg.py b/src/tfrlrl/training_algorithms/tnpg.py
index 03ee208..962672a 100644
--- a/src/tfrlrl/training_algorithms/tnpg.py
+++ b/src/tfrlrl/training_algorithms/tnpg.py
@@ -5,7 +5,7 @@
 
 import numpy as np
 import ray
-from torch import Tensor, no_grad, sum, tensor
+from torch import no_grad, sum, tensor
 from torch.optim import (
     Optimizer,
 )
@@ -30,7 +30,7 @@
 
 def calculate_steepest_gradient_direction(
     policy: BasePyTorchPolicy, statistics: EpisodePolicyGradientStatistics, optimizer: Optimizer
-) -> Tensor:
+) -> np.ndarray:
     """
     Calculate the direction of steepest gradient ascent of the policy.
 
@@ -44,7 +44,7 @@ def calculate_steepest_gradient_direction(
         optimizer: A PyTorch optimizer class that will be used to calculate the policy gradient.
 
     Returns:
-        A PyTorch tensor containing the policy gradient.
+        A NumPy array containing the policy gradient.
 
     """
     optimizer.zero_grad()
@@ -55,8 +55,12 @@ def calculate_steepest_gradient_direction(
     loss = -sum(log_probabilities * tensor(statistics.total_expected_rewards))
     loss.backward()
 
-    return flatten_tensor_dict(
-        {name: param.grad for name, param in policy.network.named_parameters()},
+    return (
+        flatten_tensor_dict(
+            {name: param.grad for name, param in policy.network.named_parameters() if param.grad is not None},
+        )
+        .detach()
+        .numpy()
     )
 
 
@@ -85,8 +89,19 @@ def construct_fim_vector_product_fn(
         observations=statistics.observations,
         actions=statistics.actions,
     )
-    # TODO: Convert Jacobian to a NumPy array
-    jacobian_matrix = flatten_tensor_dict(jacobian)
+    jacobian_matrix = (
+        flatten_tensor_dict(
+            jacobian,
+            dim=statistics.actions.ndim,
+        )
+        .detach()
+        .numpy()
+    )
+
+    if jacobian_matrix.ndim > 3 or (jacobian_matrix.ndim > 2 and jacobian_matrix.shape[0] > 1):
+        raise RuntimeError('Unsupported shape of Jacobian matrix: %s', jacobian_matrix.shape)
+    elif jacobian_matrix.ndim > 2:
+        jacobian_matrix = jacobian_matrix.squeeze()
 
     def calculate_fim_vector_product(v: np.ndarray):
         return np.matmul(jacobian_matrix.T, np.matmul(jacobian_matrix, v))
diff --git a/tests/tfrlrl/policies/test_utils.py b/tests/tfrlrl/policies/test_utils.py
index 8c5862f..9b455af 100644
--- a/tests/tfrlrl/policies/test_utils.py
+++ b/tests/tfrlrl/policies/test_utils.py
@@ -42,10 +42,10 @@ def test_flatten_tensor_dict(env_id: str):
     indx = 0
     for _, param in policy.network.named_parameters():
         for p_indx in np.ndindex(param.shape):
-            assert param[p_indx] == flattend_tensor[0, indx]
+            assert param[p_indx] == flattend_tensor[indx]
             indx += 1
 
-    assert indx == flattend_tensor.shape[1]
+    assert indx == flattend_tensor.shape[0]
 
 
 @pytest.mark.parametrize(
diff --git a/tests/tfrlrl/training_algorithms/test_tnpg.py b/tests/tfrlrl/training_algorithms/test_tnpg.py
index e66c1c7..1b23846 100644
--- a/tests/tfrlrl/training_algorithms/test_tnpg.py
+++ b/tests/tfrlrl/training_algorithms/test_tnpg.py
@@ -7,11 +7,10 @@
 from hypothesis import given, settings
 from hypothesis import strategies as st
 from torch import (
-    Tensor,
-    tensor,
+    sum as torch_sum,
 )
 from torch import (
-    sum as torch_sum,
+    tensor,
 )
 from torch.optim import (
     SGD,
@@ -22,14 +21,14 @@
 from tfrlrl.policies.linear_soft_max import LinearSoftMax
 from tfrlrl.sampling.episodic_sampler import EpisodicSampler
 from tfrlrl.sampling.statistics_collection import EpisocidPolicyGradientStatisticsCollector
-from tfrlrl.training_algorithms.tnpg import calculate_steepest_gradient_direction
+from tfrlrl.training_algorithms.tnpg import calculate_steepest_gradient_direction, construct_fim_vector_product_fn
 
 
 @pytest.mark.parametrize(
     'env_id, expected_shape',
     [
-        ('FrozenLake-v1', (1, 49)),
-        ('InvertedPendulum-v5', (1, 690)),
+        ('FrozenLake-v1', (49,)),
+        ('InvertedPendulum-v5', (690,)),
     ],
 )
 @given(
@@ -84,7 +83,7 @@ def test_calculate_steepest_gradient_direction(env_id: str, expected_shape: Tupl
         statistics,
         optimizer,
     )
-    assert isinstance(sgd, Tensor)
+    assert isinstance(sgd, np.ndarray)
     assert sgd.shape == expected_shape
 
     eps = 0.001
@@ -115,10 +114,96 @@ def test_calculate_steepest_gradient_direction(env_id: str, expected_shape: Tupl
 
     policy.set_state(policy_dict)
 
-    fd_gradient = np.concatenate(fd_grad_parts)[np.newaxis, :]
+    fd_gradient = np.concatenate(fd_grad_parts)
 
     np.testing.assert_almost_equal(
-        sgd.detach().numpy(),
+        sgd,
         fd_gradient,
         decimal=2,
     )
+
+
+@pytest.mark.parametrize(
+    'env_id, expected_shape',
+    [
+        ('FrozenLake-v1', (49,)),
+        ('InvertedPendulum-v5', (50,)),
+    ],
+)
+@given(
+    n_episodes=st.integers(min_value=10, max_value=100),
+)
+@settings(deadline=None)
+def test_construct_fim_vector_product_fn(env_id: str, expected_shape: Tuple[int], n_episodes: int):
+    """
+    Test that construct_fim_vector_product_fn returns a function for calculate the FIM-vector product.
+
+    The production of the Fisher Information matrix (FIM) against a vector will be verified by
+    performing the multiplication directly.
+
+    Args:
+        env_id: The Gymnasium environment ID with a discrete action space.
+        expected_shape: The expected shape of the gradient.
+        n_episodes: The number of episodes to sample when estimating the gradient.
+
+    """
+    env = gym.make(env_id)
+    statistics_collector = EpisocidPolicyGradientStatisticsCollector(env_id)
+
+    if env_id == 'FrozenLake-v1':
+        feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+        policy = LinearSoftMax(env_id, feature_fn)
+        sampler = EpisodicSampler(
+            env_id=env_id,
+            n_episodes=n_episodes,
+            policy=policy,
+            statistics_collector=statistics_collector,
+            is_slippery=False,
+        )
+    elif env_id == 'InvertedPendulum-v5':
+        policy = DenseNetworkPolicy(
+            env_id=env_id,
+            hidden_space_dims=[4, 4],
+        )
+        sampler = EpisodicSampler(
+            env_id=env_id,
+            n_episodes=n_episodes,
+            policy=policy,
+            statistics_collector=statistics_collector,
+        )
+    else:
+        raise ValueError('Unexpected environment: %s', env_id)
+
+    alpha = 0.1
+    optimizer = SGD(policy.get_parameters(), lr=alpha)
+    statistics = sampler.sample()
+    sgd = calculate_steepest_gradient_direction(
+        policy,
+        statistics,
+        optimizer,
+    )
+    assert isinstance(sgd, np.ndarray)
+    assert sgd.shape == expected_shape
+
+    calculate_fim_vector_product = construct_fim_vector_product_fn(
+        policy=policy,
+        statistics=statistics,
+    )
+
+    v = calculate_fim_vector_product(sgd)
+
+    jacobian = policy.calculate_jacobian(
+        statistics.observations,
+        statistics.actions,
+    )
+
+    jac_blocks = []
+    for _, jac_param in jacobian.items():
+        n_steps = jac_param.shape[1]
+        jac_block = jac_param.detach().numpy()[0].reshape(n_steps, -1)
+        jac_blocks.append(jac_block)
+    J = np.concatenate(jac_blocks, axis=1)
+
+    v_expected = np.matmul(J.T, np.matmul(J, sgd))
+
+    np.testing.assert_allclose(v, v_expected, rtol=1e-1)

From 79f6af7f4d94e708ce180330c4e2aa5cbd669d36 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 12 Jul 2026 22:09:27 +0100
Subject: [PATCH 10/19] loosen test settings

---
 tests/tfrlrl/sampling/test_episodic_sampler.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/tfrlrl/sampling/test_episodic_sampler.py b/tests/tfrlrl/sampling/test_episodic_sampler.py
index b34ab59..1e00043 100644
--- a/tests/tfrlrl/sampling/test_episodic_sampler.py
+++ b/tests/tfrlrl/sampling/test_episodic_sampler.py
@@ -307,7 +307,7 @@ def test_ray_sample_n_episodes_without_limit(self, env_id: str, n_episodes: int,
 
     @pytest.mark.slow
     @given(n_episodes=st.integers(min_value=2, max_value=10))
-    @settings(deadline=5000)
+    @settings(deadline=7500)
     def test_ray_sample_episode_with_env_kwargs(self, n_episodes: int, test_ray_cluster):
         """
         Test that environment kwargs are correctly passed through to the environment construction.
@@ -351,7 +351,7 @@ def test_ray_sample_episode_with_env_kwargs(self, n_episodes: int, test_ray_clus
 
     @pytest.mark.slow
     @given(n_episodes=st.integers(min_value=2, max_value=10))
-    @settings(deadline=5000)
+    @settings(deadline=7500)
     def test_ray_reset_allows_reuse_as_iterator(self, n_episodes: int, test_ray_cluster):
         """
         Test that the reset method allows the EpisodicSampler to be used as an iterator multiple times.
@@ -425,7 +425,7 @@ def test_ray_reset_allows_reuse_as_iterator(self, n_episodes: int, test_ray_clus
     @pytest.mark.slow
     @pytest.mark.parametrize('env_id', ['FrozenLake-v1'])
     @given(n_episodes=st.integers(min_value=2, max_value=10))
-    @settings(deadline=5000)
+    @settings(deadline=7500)
     def test_ray_sample_without_limits(self, env_id: str, n_episodes: int, test_ray_cluster):
         """
         Test the sample function of the EpisodicSampler class and that the outputs follow the expected format.

From 181859e21d98d5fc27cc260a266cc352f6991b76 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 26 Jul 2026 12:30:45 +0100
Subject: [PATCH 11/19] Add update of policy parameters

---
 src/tfrlrl/policies/utils.py           |  36 ++++++++
 src/tfrlrl/training_algorithms/tnpg.py |  32 +++++--
 tests/tfrlrl/policies/test_utils.py    | 121 ++++++++++++++++++++++++-
 3 files changed, 181 insertions(+), 8 deletions(-)

diff --git a/src/tfrlrl/policies/utils.py b/src/tfrlrl/policies/utils.py
index 59d8934..33713cd 100644
--- a/src/tfrlrl/policies/utils.py
+++ b/src/tfrlrl/policies/utils.py
@@ -29,3 +29,39 @@ def flatten_tensor_dict(x: Dict[str, Tensor], dim: int = 0) -> Tensor:
 
     """
     return cat([j.flatten(start_dim=dim) for j in x.values()], dim=dim)
+
+
+def unflatten_tensor_dict(x: Tensor, reference: Dict[str, Tensor], dim: int = 0) -> Dict[str, Tensor]:
+    """
+    Unflatten the given PyTorch tensor into a dictionary of PyTorch tensors.
+
+    This function is the inverse of flatten_tensor_dict. It takes a tensor produced by flattening
+    the tensors in reference along dim, e.g. via flatten_tensor_dict(reference, dim=dim), and splits
+    it back into a dictionary of tensors with the same keys, order and shapes as reference.
+
+    Args:
+        x: The tensor to be unflattened.
+        reference: A dictionary of tensors whose keys, order and shapes (from dim onwards) are used
+        to split and reshape x back into a dictionary of tensors.
+        dim: The starting dimension on which x was flattened and concatenated.
+
+    Returns:
+        A dictionary of tensors with the same keys, order and shapes as reference.
+
+    Raises:
+        RuntimeError: When the total size of x along dim does not match the sum of the flattened
+        sizes (from dim onwards) of the tensors in reference, a RuntimeError will be thrown by the
+        call to narrow.
+
+    """
+    leading_shape = x.shape[:dim]
+
+    unflattened = {}
+    start = 0
+    for name, ref_tensor in reference.items():
+        tail_shape = ref_tensor.shape[dim:]
+        n = tail_shape.numel()
+        unflattened[name] = x.narrow(dim, start, n).reshape(*leading_shape, *tail_shape)
+        start += n
+
+    return unflattened
diff --git a/src/tfrlrl/training_algorithms/tnpg.py b/src/tfrlrl/training_algorithms/tnpg.py
index 962672a..c54474a 100644
--- a/src/tfrlrl/training_algorithms/tnpg.py
+++ b/src/tfrlrl/training_algorithms/tnpg.py
@@ -7,15 +7,17 @@
 import ray
 from torch import no_grad, sum, tensor
 from torch.optim import (
+    SGD,
     Optimizer,
 )
+from torch.optim.lr_scheduler import LRScheduler
 
 from tfrlrl import settings
 from tfrlrl.baselines.linear import Baseline
 from tfrlrl.data_models.reward_models import AverageEpisodicReward, DiscountedReward
 from tfrlrl.optimisation.conjugate_gradients import calculate_conjugate_gradient
 from tfrlrl.policies.base import BasePyTorchPolicy
-from tfrlrl.policies.utils import flatten_tensor_dict
+from tfrlrl.policies.utils import flatten_tensor_dict, unflatten_tensor_dict
 from tfrlrl.sampling.episodic_sampler import (
     EpisodicSampler,
     RayEpisodicSampler,
@@ -114,7 +116,8 @@ def train_policy_gradient(
     policy: BasePyTorchPolicy,
     n_iterations: int,
     n_episodes: int,
-    optimizer: Optimizer,
+    lr: float,
+    lr_scheduler_fn: Optional[Callable[[Optimizer], LRScheduler]] = None,
     n_samplers: int = 1,
     baseline: Optional[Baseline] = None,
     reward_model: Optional[Union[AverageEpisodicReward, DiscountedReward]] = None,
@@ -122,14 +125,16 @@ def train_policy_gradient(
     **kwargs,
 ) -> BasePyTorchPolicy:
     """
-    Train a policy using stochastic gradient ascent on the policy gradient.
+    Train a policy using truncated natural policy gradient ascent.
 
     Args:
         env_id: Gymnasium environment ID (e.g., CartPole-v1, MountainCar-v0).
         policy: The policy to train. Must have get_parameters() and set_parameters() methods.
         n_iterations: The number of policy updates to perform.
         n_episodes: The number of episodes to sample during each policy update.
-        optimizer: An instance of a PyTorch optimizer class that will be used to optimise the policy.
+        lr: The base learning rate for the SGD optimizer used to apply the natural policy gradient.
+        lr_scheduler_fn: An optional factory that, given the SGD optimizer instantiated internally,
+        returns an LRScheduler wrapping it. When not given, the learning rate stays constant at lr.
         n_samplers: The number of samplers to used to sample from the environment.
         baseline: An instance of a baseline class, if one is given.
         reward_model: The reward model to use when computing total expected rewards. Defaults to
@@ -141,6 +146,9 @@ def train_policy_gradient(
         The trained policy.
 
     """
+    optimizer = SGD(policy.network.parameters(), lr=lr, maximize=True)
+    lr_scheduler = lr_scheduler_fn(optimizer) if lr_scheduler_fn is not None else None
+
     statistics_collector = EpisocidPolicyGradientStatisticsCollector(
         env_id,
         baseline=baseline,
@@ -175,13 +183,13 @@ def train_policy_gradient(
         with no_grad():
             statistics = sampler.sample()
 
+        logger.debug('Calculate truncated-natural policy gradient.')
         sgd = calculate_steepest_gradient_direction(
             policy=policy,
             statistics=statistics,
             optimizer=optimizer,
         )
-
-        calculate_conjugate_gradient(
+        tngd = calculate_conjugate_gradient(
             mat_v_mult_fn=construct_fim_vector_product_fn(
                 policy=policy,
                 statistics=statistics,
@@ -190,7 +198,17 @@ def train_policy_gradient(
             n_iters=10,
         )
 
-        # optimizer.step()
+        logger.debug('Update policy parameters.')
+        tngd_dict = unflatten_tensor_dict(
+            tensor(tngd),
+            reference={name: param for name, param in policy.network.named_parameters()},
+            dim=0,
+        )
+        for name, param in policy.network.named_parameters():
+            param.grad = tngd_dict[name]
+        optimizer.step()
+        if lr_scheduler is not None:
+            lr_scheduler.step()
 
         if n % n_iteration_logging == 0:
             logger.info('Policy update: %s', n)
diff --git a/tests/tfrlrl/policies/test_utils.py b/tests/tfrlrl/policies/test_utils.py
index 9b455af..bab3ff0 100644
--- a/tests/tfrlrl/policies/test_utils.py
+++ b/tests/tfrlrl/policies/test_utils.py
@@ -7,7 +7,7 @@
 from tfrlrl.features.onehot import OneHotFeatureFunction
 from tfrlrl.policies.dense_neural_network import DenseNetworkPolicy
 from tfrlrl.policies.linear_soft_max import LinearSoftMax
-from tfrlrl.policies.utils import flatten_tensor_dict
+from tfrlrl.policies.utils import flatten_tensor_dict, unflatten_tensor_dict
 
 
 @pytest.mark.parametrize('env_id', ['FrozenLake-v1', 'InvertedPendulum-v5'])
@@ -48,6 +48,42 @@ def test_flatten_tensor_dict(env_id: str):
     assert indx == flattend_tensor.shape[0]
 
 
+@pytest.mark.parametrize('env_id', ['FrozenLake-v1', 'InvertedPendulum-v5'])
+def test_unflatten_tensor_dict(env_id: str):
+    """
+    Test unflatten_tensor_dict function on policy parameters.
+
+    This function tests that unflatten_tensor_dict is the inverse of flatten_tensor_dict, i.e. that
+    flattening a dictionary of policy parameters and then unflattening the result recovers the
+    original dictionary of tensors.
+
+    Args:
+        env_id: The environment I.D. on which to run the test. This is used to determine the policy on which
+        to perform the test.
+
+    """
+    if env_id == 'FrozenLake-v1':
+        env = gym.make(env_id)
+        feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+        policy = LinearSoftMax(env_id, feature_fn)
+    elif env_id == 'InvertedPendulum-v5':
+        policy = DenseNetworkPolicy(
+            env_id=env_id,
+            hidden_space_dims=[16, 32],
+        )
+    else:
+        raise ValueError('Unexpected environment: %s', env_id)
+
+    reference = {name: param for name, param in policy.network.named_parameters()}
+    flattend_tensor = flatten_tensor_dict(reference)
+    unflattened = unflatten_tensor_dict(flattend_tensor, reference)
+
+    assert list(unflattened.keys()) == list(reference.keys())
+    for name, param in reference.items():
+        assert unflattened[name].shape == param.shape
+        assert (unflattened[name] == param).all()
+
+
 @pytest.mark.parametrize(
     'env_id, extend_actions',
     [
@@ -130,3 +166,86 @@ def test_flatten_tensor_dict_jacobian(env_id: str, extend_actions, n_observation
                     assert jac_param[(0, i) + p_indx] == flattend_jacobian[0, i, indx]
             indx += 1
     assert indx == flattend_jacobian.shape[actions.ndim]
+
+
+@pytest.mark.parametrize(
+    'env_id, extend_actions',
+    [
+        (
+            'CliffWalking-v1',
+            False,
+        ),
+        (
+            'InvertedPendulum-v5',
+            True,
+        ),
+        (
+            'InvertedPendulum-v5',
+            False,
+        ),
+    ],
+)
+@given(
+    n_observations=st.integers(min_value=5, max_value=10),
+)
+@settings(deadline=None)
+def test_unflatten_tensor_dict_jacobian(env_id: str, extend_actions, n_observations: int):
+    """
+    Test unflatten_tensor_dict function on the policy Jacobian.
+
+    This function tests that unflatten_tensor_dict is the inverse of flatten_tensor_dict, i.e. that
+    flattening the policy Jacobian and then unflattening the result recovers the original dictionary
+    of Jacobian tensors.
+
+    Args:
+        env_id: The environment I.D. on which to run the test. This is used to determine the policy on which
+        to perform the test.
+        extend_actions: A Boolean indicating whether to extend the action dimensions. This is used to capture
+        the case of a 1-dimensional conintuous action space, in which the actions could be stored as either
+        a (n_observations) or an (1, n_observations) array.
+        n_observations: The number of observations.
+
+    """
+    env = gym.make(env_id)
+
+    if env_id == 'CliffWalking-v1':
+        feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+        policy = LinearSoftMax(env_id, feature_fn)
+        observations = np.random.randint(low=0, high=47, size=(1, n_observations))
+        actions = np.random.randint(low=0, high=3, size=(1, n_observations))
+    elif env_id == 'InvertedPendulum-v5':
+        policy = DenseNetworkPolicy(
+            env_id=env_id,
+            hidden_space_dims=[16, 32],
+        )
+        observations = np.concatenate(
+            [env.observation_space.sample()[..., np.newaxis] for _ in range(n_observations)],
+            axis=1,
+        )
+        if extend_actions:
+            actions = np.concatenate(
+                [env.action_space.sample()[..., np.newaxis] for _ in range(n_observations)],
+                axis=1,
+            )
+        else:
+            actions = np.concatenate(
+                [env.action_space.sample() for _ in range(n_observations)],
+            )
+    else:
+        raise ValueError('Unexpected environment: %s', env_id)
+
+    jacobian = policy.calculate_jacobian(observations, actions)
+    flattend_jacobian = flatten_tensor_dict(
+        jacobian,
+        dim=actions.ndim,
+    )
+    unflattened_jacobian = unflatten_tensor_dict(
+        flattend_jacobian,
+        jacobian,
+        dim=actions.ndim,
+    )
+
+    assert list(unflattened_jacobian.keys()) == list(jacobian.keys())
+    for name, jac_param in jacobian.items():
+        assert unflattened_jacobian[name].shape == jac_param.shape
+        assert (unflattened_jacobian[name] == jac_param).all()

From 611b3f1e8e71f5e4f72b6f9685560397d853dddb Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 26 Jul 2026 13:30:03 +0100
Subject: [PATCH 12/19] Add unit testa

---
 src/tfrlrl/training_algorithms/tnpg.py        |   4 +-
 tests/tfrlrl/training_algorithms/test_tnpg.py | 130 +++++++++++++++++-
 2 files changed, 130 insertions(+), 4 deletions(-)

diff --git a/src/tfrlrl/training_algorithms/tnpg.py b/src/tfrlrl/training_algorithms/tnpg.py
index c54474a..b9ffed8 100644
--- a/src/tfrlrl/training_algorithms/tnpg.py
+++ b/src/tfrlrl/training_algorithms/tnpg.py
@@ -195,9 +195,9 @@ def train_policy_gradient(
                 statistics=statistics,
             ),
             b=sgd,
-            n_iters=10,
+            n_iters=1,
         )
-
+        print(tngd)
         logger.debug('Update policy parameters.')
         tngd_dict = unflatten_tensor_dict(
             tensor(tngd),
diff --git a/tests/tfrlrl/training_algorithms/test_tnpg.py b/tests/tfrlrl/training_algorithms/test_tnpg.py
index 1b23846..a93451f 100644
--- a/tests/tfrlrl/training_algorithms/test_tnpg.py
+++ b/tests/tfrlrl/training_algorithms/test_tnpg.py
@@ -16,12 +16,17 @@
     SGD,
 )
 
+from tfrlrl.baselines.linear import LinearBaseline
 from tfrlrl.features.onehot import OneHotFeatureFunction
 from tfrlrl.policies.dense_neural_network import DenseNetworkPolicy
 from tfrlrl.policies.linear_soft_max import LinearSoftMax
 from tfrlrl.sampling.episodic_sampler import EpisodicSampler
 from tfrlrl.sampling.statistics_collection import EpisocidPolicyGradientStatisticsCollector
-from tfrlrl.training_algorithms.tnpg import calculate_steepest_gradient_direction, construct_fim_vector_product_fn
+from tfrlrl.training_algorithms.tnpg import (
+    calculate_steepest_gradient_direction,
+    construct_fim_vector_product_fn,
+    train_policy_gradient,
+)
 
 
 @pytest.mark.parametrize(
@@ -131,7 +136,7 @@ def test_calculate_steepest_gradient_direction(env_id: str, expected_shape: Tupl
     ],
 )
 @given(
-    n_episodes=st.integers(min_value=10, max_value=100),
+    n_episodes=st.integers(min_value=2, max_value=10),
 )
 @settings(deadline=None)
 def test_construct_fim_vector_product_fn(env_id: str, expected_shape: Tuple[int], n_episodes: int):
@@ -207,3 +212,124 @@ def test_construct_fim_vector_product_fn(env_id: str, expected_shape: Tuple[int]
     v_expected = np.matmul(J.T, np.matmul(J, sgd))
 
     np.testing.assert_allclose(v, v_expected, rtol=1e-1)
+
+
+@pytest.mark.parametrize(
+    'env_id, use_baseline',
+    [
+        # TODO: Fix numerical errors in this test
+        # (
+        #     'FrozenLake-v1',
+        #     False,
+        # ),
+        (
+            'InvertedPendulum-v5',
+            False,
+        ),
+        (
+            'InvertedPendulum-v5',
+            True,
+        ),
+    ],
+)
+@given(
+    n_iterations=st.integers(min_value=2, max_value=5),
+    n_episodes=st.integers(min_value=2, max_value=5),
+    lr=st.floats(min_value=0.0000001, max_value=0.000001),
+)
+@settings(deadline=5000)
+def test_train_policy_gradient_returns_policy(
+    env_id: str,
+    use_baseline: bool,
+    n_iterations: int,
+    n_episodes: int,
+    lr: float,
+):
+    """
+    Test that train_policy_gradient executes successfully and returns a policy.
+
+    Args:
+        env_id: The Gym environment ID to be used in training.
+        use_baseline: A Boolean indicating whether to use a linear baseline.
+        n_iterations: The number of policy updates to perform.
+        n_episodes: The number of episodes to sample during each policy update.
+        lr: The base learning rate for the SGD optimizer used to apply the natural policy gradient.
+
+    """
+    env = gym.make(env_id)
+
+    if env_id == 'InvertedPendulum-v5':
+        policy = DenseNetworkPolicy(
+            env_id=env_id,
+            hidden_space_dims=[16, 32],
+        )
+    else:
+        feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+        policy = LinearSoftMax(env_id, feature_fn)
+
+    if use_baseline:
+        baseline = LinearBaseline(env_id=env_id)
+    else:
+        baseline = None
+
+    trained_policy = train_policy_gradient(
+        env_id=env_id,
+        policy=policy,
+        n_iterations=n_iterations,
+        n_episodes=n_episodes,
+        lr=lr,
+        baseline=baseline,
+    )
+
+    assert trained_policy is not None
+    if env_id == 'InvertedPendulum-v5':
+        assert isinstance(trained_policy, DenseNetworkPolicy)
+    else:
+        assert isinstance(trained_policy, LinearSoftMax)
+
+    assert trained_policy is policy
+
+
+@pytest.mark.parametrize('env_id', ['FrozenLake-v1'])
+@given(
+    n_iterations=st.integers(min_value=2, max_value=5),
+    n_episodes=st.integers(min_value=2, max_value=5),
+    lr=st.floats(min_value=0.00001, max_value=0.0001),
+)
+@settings(deadline=5000)
+def test_train_policy_gradient_updates_policy(env_id: str, n_iterations: int, n_episodes: int, lr: float):
+    """
+    Test that train_policy_gradient executes successfully and updates the policy.
+
+    Args:
+        env_id: The Gym environment ID to be used in training.
+        n_iterations: The number of policy updates to perform.
+        n_episodes: The number of episodes to sample during each policy update.
+        lr: The base learning rate for the SGD optimizer used to apply the natural policy gradient.
+
+    """
+    env = gym.make(env_id)
+
+    feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+    policy = LinearSoftMax(env_id, feature_fn)
+
+    original_parameters = copy.deepcopy(list(policy.get_parameters()))
+
+    # Train the policy - We set reward_schedule to ensure the policy is updated.
+    trained_policy = train_policy_gradient(
+        env_id=env_id,
+        policy=policy,
+        n_iterations=n_iterations,
+        n_episodes=n_episodes,
+        lr=lr,
+        is_slippery=False,
+        reward_schedule=(1, 1, 1),
+    )
+
+    assert trained_policy is not None
+    assert isinstance(trained_policy, LinearSoftMax)
+    assert trained_policy is policy
+
+    updated_parameters = list(policy.get_parameters())
+    parameter_diff = original_parameters[0].detach().numpy() - updated_parameters[0].detach().numpy()
+    assert np.sum(np.abs(parameter_diff)) > 0

From 8866e42398449e56ed95884bfb4b6677899ee9d5 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 26 Jul 2026 13:56:02 +0100
Subject: [PATCH 13/19] Add subsampling of samples used for FIM product

---
 src/tfrlrl/training_algorithms/tnpg.py        | 35 +++++++--
 tests/tfrlrl/training_algorithms/test_tnpg.py | 72 +++++++++++++++++++
 2 files changed, 100 insertions(+), 7 deletions(-)

diff --git a/src/tfrlrl/training_algorithms/tnpg.py b/src/tfrlrl/training_algorithms/tnpg.py
index b9ffed8..c29b380 100644
--- a/src/tfrlrl/training_algorithms/tnpg.py
+++ b/src/tfrlrl/training_algorithms/tnpg.py
@@ -67,7 +67,9 @@ def calculate_steepest_gradient_direction(
 
 
 def construct_fim_vector_product_fn(
-    policy: BasePyTorchPolicy, statistics: EpisodePolicyGradientStatistics
+    policy: BasePyTorchPolicy,
+    statistics: EpisodePolicyGradientStatistics,
+    n_samples_fim: Optional[int] = None,
 ) -> Callable[[np.ndarray], np.ndarray]:
     """
     Construct function for calculating the product of the Fisher Information matrix with a vector.
@@ -81,20 +83,33 @@ def construct_fim_vector_product_fn(
         policy: The policy for which the Fisher Infromation matrix is to be calculated.
         statistics: The statistics over which the expectation within the Fisher Information matrix calculation
         is to be performed.
+        n_samples_fim: The number of state-action pairs to randomly subsample (without replacement) from
+        statistics for use in the Fisher Information matrix calculation. When not given, or when it is at
+        least as large as the number of sampled state-action pairs, all of the state-action pairs in
+        statistics will be used.
 
     Returns:
         A callable that takes an input vector as an argument and returns the product of the Fisher Information
         matrix with that vector.
 
     """
+    observations = statistics.observations
+    actions = statistics.actions
+
+    n_steps = actions.shape[-1]
+    if n_samples_fim is not None and n_samples_fim < n_steps:
+        idx = np.random.default_rng().choice(n_steps, size=n_samples_fim, replace=False)
+        observations = observations[..., idx]
+        actions = actions[..., idx]
+
     jacobian = policy.calculate_jacobian(
-        observations=statistics.observations,
-        actions=statistics.actions,
+        observations=observations,
+        actions=actions,
     )
     jacobian_matrix = (
         flatten_tensor_dict(
             jacobian,
-            dim=statistics.actions.ndim,
+            dim=actions.ndim,
         )
         .detach()
         .numpy()
@@ -103,7 +118,7 @@ def construct_fim_vector_product_fn(
     if jacobian_matrix.ndim > 3 or (jacobian_matrix.ndim > 2 and jacobian_matrix.shape[0] > 1):
         raise RuntimeError('Unsupported shape of Jacobian matrix: %s', jacobian_matrix.shape)
     elif jacobian_matrix.ndim > 2:
-        jacobian_matrix = jacobian_matrix.squeeze()
+        jacobian_matrix = jacobian_matrix.squeeze(axis=0)
 
     def calculate_fim_vector_product(v: np.ndarray):
         return np.matmul(jacobian_matrix.T, np.matmul(jacobian_matrix, v))
@@ -122,6 +137,8 @@ def train_policy_gradient(
     baseline: Optional[Baseline] = None,
     reward_model: Optional[Union[AverageEpisodicReward, DiscountedReward]] = None,
     n_iteration_logging: int = 10,
+    n_iters_cg: Optional[int] = None,
+    n_samples_fim: Optional[int] = None,
     **kwargs,
 ) -> BasePyTorchPolicy:
     """
@@ -140,6 +157,10 @@ def train_policy_gradient(
         reward_model: The reward model to use when computing total expected rewards. Defaults to
         AverageEpisodicReward if not specified.
         n_iteration_logging: The number of algorithm iterations between logging algorithm performance.
+        n_iters_cg: The maximum number of conjugate-gradient iterations to perform when calculating the
+        truncated natural policy gradient direction. Defaults to the default of calculate_conjugate_gradient.
+        n_samples_fim: The number of state-action pairs to randomly subsample when calculating the Fisher
+        Information matrix-vector product. When not given, all sampled state-action pairs are used.
         kwargs: Additional keyword arguments to pass to the EpisodicSampler (e.g., is_slippery).
 
     Returns:
@@ -193,11 +214,11 @@ def train_policy_gradient(
             mat_v_mult_fn=construct_fim_vector_product_fn(
                 policy=policy,
                 statistics=statistics,
+                n_samples_fim=n_samples_fim,
             ),
             b=sgd,
-            n_iters=1,
+            n_iters=n_iters_cg,
         )
-        print(tngd)
         logger.debug('Update policy parameters.')
         tngd_dict = unflatten_tensor_dict(
             tensor(tngd),
diff --git a/tests/tfrlrl/training_algorithms/test_tnpg.py b/tests/tfrlrl/training_algorithms/test_tnpg.py
index a93451f..1837d68 100644
--- a/tests/tfrlrl/training_algorithms/test_tnpg.py
+++ b/tests/tfrlrl/training_algorithms/test_tnpg.py
@@ -214,6 +214,74 @@ def test_construct_fim_vector_product_fn(env_id: str, expected_shape: Tuple[int]
     np.testing.assert_allclose(v, v_expected, rtol=1e-1)
 
 
+def test_construct_fim_vector_product_fn_subsamples_state_action_pairs(monkeypatch):
+    """
+    Test that construct_fim_vector_product_fn only uses a subsample of state-action pairs when n_samples_fim is given.
+
+    The subsample of state-action pairs used within the calculation is verified by mocking the random number
+    generator used to select the subsample indices, and comparing the result directly against the Fisher
+    Information matrix-vector product calculated using only the observations and actions at those indices.
+
+    Args:
+        monkeypatch: The PyTest monkeypatch fixture, used to control the random subsample of indices selected.
+
+    """
+    env_id = 'FrozenLake-v1'
+    n_episodes = 10
+    env = gym.make(env_id)
+    feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+    policy = LinearSoftMax(env_id, feature_fn)
+    statistics_collector = EpisocidPolicyGradientStatisticsCollector(env_id)
+    sampler = EpisodicSampler(
+        env_id=env_id,
+        n_episodes=n_episodes,
+        policy=policy,
+        statistics_collector=statistics_collector,
+        is_slippery=False,
+    )
+    statistics = sampler.sample()
+
+    n_steps = statistics.actions.shape[-1]
+    n_samples_fim = max(1, n_steps // 2)
+    idx = np.arange(n_samples_fim)
+
+    class _FakeRNG:
+        def choice(self, n, size, replace):
+            return idx
+
+    monkeypatch.setattr(np.random, 'default_rng', lambda: _FakeRNG())
+
+    alpha = 0.1
+    optimizer = SGD(policy.get_parameters(), lr=alpha)
+    sgd = calculate_steepest_gradient_direction(
+        policy,
+        statistics,
+        optimizer,
+    )
+
+    calculate_fim_vector_product = construct_fim_vector_product_fn(
+        policy=policy,
+        statistics=statistics,
+        n_samples_fim=n_samples_fim,
+    )
+    v = calculate_fim_vector_product(sgd)
+
+    jacobian = policy.calculate_jacobian(
+        statistics.observations[..., idx],
+        statistics.actions[..., idx],
+    )
+    jac_blocks = []
+    for _, jac_param in jacobian.items():
+        n_steps_sub = jac_param.shape[1]
+        jac_block = jac_param.detach().numpy()[0].reshape(n_steps_sub, -1)
+        jac_blocks.append(jac_block)
+    J = np.concatenate(jac_blocks, axis=1)
+
+    v_expected = np.matmul(J.T, np.matmul(J, sgd))
+
+    np.testing.assert_allclose(v, v_expected, rtol=1e-6)
+
+
 @pytest.mark.parametrize(
     'env_id, use_baseline',
     [
@@ -279,6 +347,8 @@ def test_train_policy_gradient_returns_policy(
         n_episodes=n_episodes,
         lr=lr,
         baseline=baseline,
+        n_iters_cg=2,
+        n_samples_fim=20,
     )
 
     assert trained_policy is not None
@@ -324,6 +394,8 @@ def test_train_policy_gradient_updates_policy(env_id: str, n_iterations: int, n_
         lr=lr,
         is_slippery=False,
         reward_schedule=(1, 1, 1),
+        n_iters_cg=2,
+        n_samples_fim=20,
     )
 
     assert trained_policy is not None

From 3ad318ab0f086f2379ff82fa5359412f3d7ca2bf Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 26 Jul 2026 14:09:12 +0100
Subject: [PATCH 14/19] add CLI for truncated natural policy gradients

---
 README.md      | 34 ++++++++++++++++++++++++++++++++++
 pyproject.toml |  1 +
 2 files changed, 35 insertions(+)

diff --git a/README.md b/README.md
index 41f972e..0f9b679 100644
--- a/README.md
+++ b/README.md
@@ -89,6 +89,40 @@ poetry run tfrlrl-sgd --env-id FrozenLake-v1 --policy-class linear --n-iteration
 - `--policy-class`: The class of policy to use in the environment. Allowed values are `linear` and `dense`.
 - `--n-hidden`: The number of hidden dimensions to use in the case of a dense policy.
 
+### tfrlrl-tnpg
+
+Train a policy using truncated natural policy gradient ascent, i.e. the natural policy gradient direction is
+approximated via a truncated run of conjugate-gradients, rather than by inverting the Fisher Information matrix
+directly.
+
+**Basic Usage:**
+
+```bash
+# Perform truncated natural policy gradient ascent on the given environment
+poetry run tfrlrl-tnpg --env-id FrozenLake-v1 --policy-class linear --n-iterations 100
+
+# With environment-specific configuration
+poetry run tfrlrl-tnpg --env-id FrozenLake-v1 --policy-class linear --n-iterations 100 --env-kwargs '{"is_slippery": false}'
+
+# With custom hyperparameters
+poetry run tfrlrl-tnpg --env-id FrozenLake-v1 --policy-class linear --n-iterations 50 --n-episodes 200 --alpha 10.0 --n-iters-cg 10 --n-samples-fim 200
+```
+
+**Options:**
+
+- `--env-id`: Gymnasium environment ID (e.g., FrozenLake-v1)
+- `--n-iterations`: Total number of policy updates to perform (default: 100)
+- `--n-episodes`: Total number of episodes to sample during each policy update (default: 100)
+- `--alpha`: The base learning rate for the SGD optimizer used to apply the natural policy gradient (default: 100.0)
+- `--n-samplers`: The number of samplers to use during sampling (default: 1)
+- `--env-kwargs`: Environment-specific keyword arguments as a JSON string (default: `{}`). For example, `'{"is_slippery": false}'` for FrozenLake-v1
+- `--policy-class`: The class of policy to use in the environment. Allowed values are `linear` and `dense`.
+- `--n-hidden`: The number of hidden dimensions to use in the case of a dense policy.
+- `--reward-model`: The reward model to use when computing returns. Allowed values are `average-episodic` and `discounted` (default: `average-episodic`)
+- `--gamma`: Discount factor for the discounted reward model, must be in (0, 1). Required when `--reward-model=discounted`.
+- `--n-iters-cg`: The maximum number of conjugate-gradient iterations to perform when calculating the truncated natural policy gradient direction. Defaults to the default of `calculate_conjugate_gradient`.
+- `--n-samples-fim`: The number of state-action pairs to randomly subsample when calculating the Fisher Information matrix-vector product. When not given, all sampled state-action pairs are used.
+
 ## Configuration
 
 The library uses Dynaconf for configuration management. Settings can be controlled via:
diff --git a/pyproject.toml b/pyproject.toml
index becc1ed..2a15b50 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,6 +34,7 @@ mypy = "^1.20.0"
 [tool.poetry.scripts]
 tfrlrl-sample = "tfrlrl.cli.sample:main"
 tfrlrl-sgd = "tfrlrl.cli.sgd:main"
+tfrlrl-tnpg = "tfrlrl.cli.tnpg:main"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

From 176c4b2c5cdbfe4ef088c7475bd2379b9cfd15f9 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 26 Jul 2026 14:10:52 +0100
Subject: [PATCH 15/19] add CLI for truncated natural policy gradients

---
 src/tfrlrl/cli/tnpg.py        | 166 ++++++++++++++
 tests/tfrlrl/cli/test_tnpg.py | 403 ++++++++++++++++++++++++++++++++++
 2 files changed, 569 insertions(+)
 create mode 100644 src/tfrlrl/cli/tnpg.py
 create mode 100644 tests/tfrlrl/cli/test_tnpg.py

diff --git a/src/tfrlrl/cli/tnpg.py b/src/tfrlrl/cli/tnpg.py
new file mode 100644
index 0000000..862a06d
--- /dev/null
+++ b/src/tfrlrl/cli/tnpg.py
@@ -0,0 +1,166 @@
+import argparse
+import json
+import logging
+
+import gymnasium as gym
+
+from tfrlrl.data_models.reward_models import AverageEpisodicReward, DiscountedReward
+from tfrlrl.features.onehot import OneHotFeatureFunction
+from tfrlrl.policies.dense_neural_network import DenseNetworkPolicy
+from tfrlrl.policies.linear_soft_max import LinearSoftMax
+from tfrlrl.training_algorithms.tnpg import train_policy_gradient
+
+logging.basicConfig(format='%(asctime)s %(message)s', force=True)
+logger = logging.getLogger(__name__)
+
+
+def parse_args(args=None):
+    """
+    Parse command line arguments for the truncated natural policy gradient CLI.
+
+    :param args: Command line arguments to parse. If None, uses sys.argv.
+    :return: Parsed arguments.
+    """
+    parser = argparse.ArgumentParser(
+        description='Train a policy using truncated natural policy gradient ascent',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        '--env-id',
+        type=str,
+        required=True,
+        help='Gymnasium environment ID (e.g., CartPole-v1, MountainCar-v0)',
+    )
+    parser.add_argument(
+        '--n-iterations',
+        type=int,
+        default=100,
+        help='The number of policy updates to perform.',
+    )
+    parser.add_argument(
+        '--n-episodes',
+        type=int,
+        default=100,
+        help='The number of episodes to sample during each policy update.',
+    )
+    parser.add_argument(
+        '--alpha',
+        type=float,
+        default=100.0,
+        help='The base learning rate for the SGD optimizer used to apply the natural policy gradient.',
+    )
+    parser.add_argument(
+        '--n-samplers',
+        type=int,
+        default=1,
+        help='The number of samplers to use when sampling episodes.',
+    )
+    parser.add_argument(
+        '--env-kwargs',
+        type=str,
+        default='{}',
+        help='Environment-specific keyword arguments as a JSON string (e.g., \'{"is_slippery": false}\').',
+    )
+    parser.add_argument(
+        '--policy-class',
+        type=str,
+        required=True,
+        choices=['linear', 'dense'],
+        help='The type of policy class to use in the truncated natural policy gradient algorithm.',
+    )
+    parser.add_argument(
+        '--n-hidden',
+        type=int,
+        nargs='+',
+        default=[16, 32],
+        help='The number of hidden dimensions to use in a dense policy network.',
+    )
+    parser.add_argument(
+        '--reward-model',
+        type=str,
+        default='average-episodic',
+        choices=['average-episodic', 'discounted'],
+        help='Reward model to use when computing returns.',
+    )
+    parser.add_argument(
+        '--gamma',
+        type=float,
+        default=None,
+        help='Discount factor for the discounted reward model, must be in (0, 1). '
+        'Required when --reward-model=discounted.',
+    )
+    parser.add_argument(
+        '--n-iters-cg',
+        type=int,
+        default=None,
+        help='The maximum number of conjugate-gradient iterations to perform when calculating the truncated '
+        'natural policy gradient direction. Defaults to the default of calculate_conjugate_gradient.',
+    )
+    parser.add_argument(
+        '--n-samples-fim',
+        type=int,
+        default=None,
+        help='The number of state-action pairs to randomly subsample when calculating the Fisher Information '
+        'matrix-vector product. When not given, all sampled state-action pairs are used.',
+    )
+    parsed = parser.parse_args(args)
+    if parsed.reward_model == 'discounted' and parsed.gamma is None:
+        parser.error('--gamma is required when --reward-model=discounted')
+    return parsed
+
+
+def main(args=None):
+    """
+    Entry point for the truncated natural policy gradient CLI.
+
+    :param args: Command line arguments. If None, uses sys.argv.
+    :return: Exit code (0 for success, 1 for failure).
+    """
+    parsed_args = parse_args(args)
+
+    # Parse environment kwargs from JSON string
+    try:
+        env_kwargs = json.loads(parsed_args.env_kwargs)
+    except json.JSONDecodeError as e:
+        logger.error('Failed to parse --env-kwargs as JSON: %s', e)
+        return 1
+
+    if not isinstance(env_kwargs, dict):
+        logger.error('--env-kwargs must be a JSON object (dictionary), got: %s', type(env_kwargs).__name__)
+        return 1
+    if env_kwargs is not None:
+        logger.info('Environment Arguments: %s', env_kwargs)
+
+    if parsed_args.reward_model == 'discounted':
+        try:
+            reward_model = DiscountedReward(gamma=parsed_args.gamma)
+        except (TypeError, ValueError) as e:
+            logger.error('Invalid --gamma value: %s', e)
+            return 1
+    else:
+        reward_model = AverageEpisodicReward()
+
+    if parsed_args.policy_class == 'linear':
+        logger.info('Using a linear policy with a one-hot feature encoding.')
+        env = gym.make(parsed_args.env_id)
+        feature_fn = OneHotFeatureFunction(env.observation_space.n, env.action_space.n)
+        policy = LinearSoftMax(parsed_args.env_id, feature_fn)
+    else:
+        logger.info('Using a dense policy with hidden dimensions: %s', parsed_args.n_hidden)
+        policy = DenseNetworkPolicy(
+            env_id=parsed_args.env_id,
+            hidden_space_dims=parsed_args.n_hidden,
+        )
+
+    train_policy_gradient(
+        env_id=parsed_args.env_id,
+        policy=policy,
+        n_iterations=parsed_args.n_iterations,
+        n_episodes=parsed_args.n_episodes,
+        lr=parsed_args.alpha,
+        n_samplers=parsed_args.n_samplers,
+        reward_model=reward_model,
+        n_iters_cg=parsed_args.n_iters_cg,
+        n_samples_fim=parsed_args.n_samples_fim,
+        **env_kwargs,
+    )
diff --git a/tests/tfrlrl/cli/test_tnpg.py b/tests/tfrlrl/cli/test_tnpg.py
new file mode 100644
index 0000000..a2d362a
--- /dev/null
+++ b/tests/tfrlrl/cli/test_tnpg.py
@@ -0,0 +1,403 @@
+"""Tests for the truncated natural policy gradient CLI."""
+
+import pytest
+from hypothesis import given, settings
+from hypothesis import strategies as st
+
+from tfrlrl.cli.tnpg import main, parse_args
+
+
+class TestParseArgs:
+    """Tests for the parse_args function."""
+
+    @pytest.mark.parametrize(
+        'args,expected',
+        [
+            (
+                [
+                    '--env-id',
+                    'FrozenLake-v1',
+                    '--n-iterations',
+                    '10',
+                    '--n-episodes',
+                    '20',
+                    '--alpha',
+                    '1.0',
+                    '--policy-class',
+                    'dense',
+                ],
+                {
+                    'env_id': 'FrozenLake-v1',
+                    'n_iterations': 10,
+                    'n_episodes': 20,
+                    'alpha': 1.0,
+                    'env_kwargs': '{}',
+                    'policy_class': 'dense',
+                    'n_iters_cg': None,
+                    'n_samples_fim': None,
+                },
+            ),
+            (
+                ['--env-id', 'FrozenLake-v1', '--env-kwargs', '{"is_slippery": false}', '--policy-class', 'dense'],
+                {
+                    'env_id': 'FrozenLake-v1',
+                    'n_iterations': 100,
+                    'n_episodes': 100,
+                    'alpha': 100.0,
+                    'env_kwargs': '{"is_slippery": false}',
+                    'policy_class': 'dense',
+                    'n_iters_cg': None,
+                    'n_samples_fim': None,
+                },
+            ),
+            (
+                [
+                    '--env-id',
+                    'FrozenLake-v1',
+                    '--n-iterations',
+                    '10',
+                    '--n-episodes',
+                    '20',
+                    '--alpha',
+                    '1.0',
+                    '--policy-class',
+                    'linear',
+                    '--n-iters-cg',
+                    '5',
+                    '--n-samples-fim',
+                    '50',
+                ],
+                {
+                    'env_id': 'FrozenLake-v1',
+                    'n_iterations': 10,
+                    'n_episodes': 20,
+                    'alpha': 1.0,
+                    'env_kwargs': '{}',
+                    'policy_class': 'linear',
+                    'n_iters_cg': 5,
+                    'n_samples_fim': 50,
+                },
+            ),
+        ],
+    )
+    def test_parse_args(self, args, expected):
+        """Test parsing command line arguments with various combinations."""
+        parsed = parse_args(args)
+        assert parsed.env_id == expected['env_id']
+        assert parsed.n_iterations == expected['n_iterations']
+        assert parsed.n_episodes == expected['n_episodes']
+        assert parsed.alpha == expected['alpha']
+        assert parsed.env_kwargs == expected['env_kwargs']
+        assert parsed.n_iters_cg == expected['n_iters_cg']
+        assert parsed.n_samples_fim == expected['n_samples_fim']
+
+    def test_parse_args_missing_required(self):
+        """Test that missing required arguments raises SystemExit."""
+        with pytest.raises(SystemExit):
+            parse_args(['--n-iterations', '10'])
+
+    def test_parse_args_reward_model_has_default(self):
+        """Test that --reward-model has a default (i.e. can be omitted)."""
+        parsed = parse_args(['--env-id', 'FrozenLake-v1', '--policy-class', 'dense'])
+        assert parsed.reward_model is not None
+
+    @pytest.mark.parametrize(
+        'args,expected_reward_model',
+        [
+            (
+                ['--env-id', 'FrozenLake-v1', '--policy-class', 'dense', '--reward-model', 'average-episodic'],
+                'average-episodic',
+            ),
+            (
+                [
+                    '--env-id',
+                    'FrozenLake-v1',
+                    '--policy-class',
+                    'dense',
+                    '--reward-model',
+                    'discounted',
+                    '--gamma',
+                    '0.95',
+                ],
+                'discounted',
+            ),
+        ],
+    )
+    def test_parse_args_reward_model(self, args, expected_reward_model):
+        """Test that --reward-model argument is parsed correctly."""
+        parsed = parse_args(args)
+        assert parsed.reward_model == expected_reward_model
+
+    def test_parse_args_invalid_reward_model(self):
+        """Test that an invalid --reward-model value raises SystemExit."""
+        with pytest.raises(SystemExit):
+            parse_args(['--env-id', 'FrozenLake-v1', '--policy-class', 'dense', '--reward-model', 'invalid'])
+
+
+class TestMain:
+    """Tests for the main function."""
+
+    @pytest.mark.parametrize(
+        'env_id, policy_class',
+        [
+            (
+                'InvertedPendulum-v5',
+                'dense',
+            ),
+            # TODO: Fix numerical errors in this test
+            # (
+            #     'FrozenLake-v1',
+            #     'linear',
+            # ),
+        ],
+    )
+    @given(
+        n_iterations=st.integers(min_value=1, max_value=3),
+        n_episodes=st.integers(min_value=5, max_value=15),
+        alpha=st.floats(min_value=0.0000001, max_value=0.000001),
+    )
+    @settings(deadline=10000)
+    def test_main_with_different_parameters(
+        self, env_id: str, policy_class: str, n_iterations: int, n_episodes: int, alpha: float
+    ):
+        """
+        Test main function with different parameter values.
+
+        Args:
+            env_id: The Gym environment ID to be used in training.
+            policy_class: The policy class to use.
+            n_iterations: The number of policy updates to perform.
+            n_episodes: The number of episodes to sample during each policy update.
+            alpha: The base learning rate for the SGD optimizer used to apply the natural policy gradient.
+
+        """
+        args = [
+            '--env-id',
+            env_id,
+            '--policy-class',
+            policy_class,
+            '--n-iterations',
+            str(n_iterations),
+            '--n-episodes',
+            str(n_episodes),
+            '--alpha',
+            str(alpha),
+        ]
+
+        exit_code = main(args)
+
+        assert exit_code is None or exit_code == 0
+
+    @pytest.mark.parametrize(
+        'env_id, policy_class, env_kwargs',
+        [
+            ('InvertedPendulum-v5', 'dense', '{"reset_noise_scale": 0.001}'),
+            # TODO: Fix numerical errors in this test
+            # (
+            #     'FrozenLake-v1',
+            #     'linear',
+            #     '{"is_slippery": false}',
+            # ),
+        ],
+    )
+    def test_main_with_env_kwargs(self, env_id: str, policy_class: str, env_kwargs: str):
+        """
+        Test main function with env-kwargs set.
+
+        Args:
+            env_id: The Gym environment ID to be used in training.
+            policy_class: The policy class to use.
+            env_kwargs: Any key-words passed to the environment.
+
+        """
+        args = [
+            '--env-id',
+            env_id,
+            '--policy-class',
+            policy_class,
+            '--n-iterations',
+            '2',
+            '--n-episodes',
+            '5',
+            '--alpha',
+            '0.0000001',
+            '--env-kwargs',
+            env_kwargs,
+        ]
+
+        exit_code = main(args)
+
+        assert exit_code is None or exit_code == 0
+
+    @pytest.mark.parametrize(
+        'env_id, policy_class',
+        [
+            (
+                'InvertedPendulum-v5',
+                'dense',
+            ),
+            # TODO: Fix numerical errors in this test
+            # (
+            #     'FrozenLake-v1',
+            #     'linear',
+            # ),
+        ],
+    )
+    def test_main_with_default_env_kwargs(self, env_id: str, policy_class: str):
+        """
+        Test main function with default (empty) env-kwargs.
+
+        Args:
+            env_id: The Gym environment ID to be used in training.
+            policy_class: The policy class to use.
+
+        """
+        args = [
+            '--env-id',
+            env_id,
+            '--policy-class',
+            policy_class,
+            '--n-iterations',
+            '2',
+            '--n-episodes',
+            '5',
+            '--alpha',
+            '0.0000001',
+        ]
+
+        exit_code = main(args)
+
+        assert exit_code is None or exit_code == 0
+
+    @pytest.mark.parametrize(
+        'env_id, policy_class',
+        [
+            # TODO: Fix numerical errors in this test
+            # ('FrozenLake-v1', 'linear'),
+        ],
+    )
+    def test_main_discounted_reward_model(self, env_id: str, policy_class: str):
+        """Test main function with the discounted reward model."""
+        args = [
+            '--env-id',
+            env_id,
+            '--policy-class',
+            policy_class,
+            '--n-iterations',
+            '2',
+            '--n-episodes',
+            '5',
+            '--alpha',
+            '0.0000001',
+            '--reward-model',
+            'discounted',
+            '--gamma',
+            '0.95',
+        ]
+        exit_code = main(args)
+        assert exit_code is None or exit_code == 0
+
+    def test_main_discounted_reward_model_requires_gamma(self):
+        """Test that using --reward-model=discounted without --gamma exits with code 2."""
+        args = [
+            '--env-id',
+            'FrozenLake-v1',
+            '--policy-class',
+            'linear',
+            '--n-iterations',
+            '1',
+            '--n-episodes',
+            '2',
+            '--alpha',
+            '0.0000001',
+            '--reward-model',
+            'discounted',
+        ]
+        with pytest.raises(SystemExit) as exc_info:
+            main(args)
+        assert exc_info.value.code == 2
+
+    def test_main_discounted_reward_model_invalid_gamma(self):
+        """Test that an invalid --gamma value (e.g. >= 1.0) returns exit code 1."""
+        args = [
+            '--env-id',
+            'FrozenLake-v1',
+            '--policy-class',
+            'linear',
+            '--n-iterations',
+            '1',
+            '--n-episodes',
+            '2',
+            '--alpha',
+            '0.0000001',
+            '--reward-model',
+            'discounted',
+            '--gamma',
+            '1.0',
+        ]
+        exit_code = main(args)
+        assert exit_code == 1
+
+    @pytest.mark.parametrize(
+        'malformed_json',
+        [
+            '{"is_slippery": false',  # Missing closing brace
+            '{is_slippery: false}',  # Missing quotes around key
+            '{"is_slippery": False}',  # Python bool instead of JSON bool
+            'not valid json',  # Completely invalid
+            '["is_slippery", false]',  # JSON array instead of object
+            '123',  # JSON number instead of object
+            '"string"',  # JSON string instead of object
+        ],
+    )
+    def test_main_with_malformed_env_kwargs(self, malformed_json: str):
+        """
+        Test main function with malformed env-kwargs returns exit code 1.
+
+        Args:
+            malformed_json: Malformed JSON string to test.
+
+        """
+        args = [
+            '--env-id',
+            'FrozenLake-v1',
+            '--n-iterations',
+            '1',
+            '--n-episodes',
+            '2',
+            '--alpha',
+            '0.0000001',
+            '--env-kwargs',
+            malformed_json,
+            '--policy-class',
+            'dense',
+        ]
+
+        exit_code = main(args)
+
+        assert exit_code == 1
+
+    def test_main_with_n_iters_cg_and_n_samples_fim(self):
+        """Test main function with --n-iters-cg and --n-samples-fim set."""
+        args = [
+            '--env-id',
+            'FrozenLake-v1',
+            '--policy-class',
+            'linear',
+            '--n-iterations',
+            '2',
+            '--n-episodes',
+            '5',
+            '--alpha',
+            '0.0000001',
+            '--env-kwargs',
+            '{"is_slippery": false}',
+            '--n-iters-cg',
+            '2',
+            '--n-samples-fim',
+            '30',
+        ]
+
+        exit_code = main(args)
+
+        assert exit_code is None or exit_code == 0

From b857fe21a3f06c1e136e8e57f4004571f37fe623 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 26 Jul 2026 14:11:23 +0100
Subject: [PATCH 16/19] update gitignore

---
 .gitignore | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index b176143..2c940d2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -149,4 +149,6 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
\ No newline at end of file
+#.idea/
+
+.DS_Store
\ No newline at end of file

From db66382e99c658af3aee08197a4fe34c2e526ada Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 26 Jul 2026 16:49:54 +0100
Subject: [PATCH 17/19] fix numerical issue

---
 .../optimisation/conjugate_gradients.py       |  3 +++
 tests/tfrlrl/cli/test_tnpg.py                 | 23 ++++++++-----------
 .../optimisation/test_conjugate_gradients.py  | 20 ++++++++++++++++
 tests/tfrlrl/training_algorithms/test_tnpg.py |  9 ++++----
 4 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/src/tfrlrl/optimisation/conjugate_gradients.py b/src/tfrlrl/optimisation/conjugate_gradients.py
index 4176224..68b877c 100644
--- a/src/tfrlrl/optimisation/conjugate_gradients.py
+++ b/src/tfrlrl/optimisation/conjugate_gradients.py
@@ -36,6 +36,9 @@ def calculate_conjugate_gradient(
     x = np.zeros_like(b)
     rdotr = np.inner(r, r)
 
+    if rdotr < tol:
+        return x
+
     for _ in range(n_iters):
         z = mat_v_mult_fn(p)
 
diff --git a/tests/tfrlrl/cli/test_tnpg.py b/tests/tfrlrl/cli/test_tnpg.py
index a2d362a..e13d595 100644
--- a/tests/tfrlrl/cli/test_tnpg.py
+++ b/tests/tfrlrl/cli/test_tnpg.py
@@ -144,11 +144,10 @@ class TestMain:
                 'InvertedPendulum-v5',
                 'dense',
             ),
-            # TODO: Fix numerical errors in this test
-            # (
-            #     'FrozenLake-v1',
-            #     'linear',
-            # ),
+            (
+                'FrozenLake-v1',
+                'linear',
+            ),
         ],
     )
     @given(
@@ -192,12 +191,11 @@ def test_main_with_different_parameters(
         'env_id, policy_class, env_kwargs',
         [
             ('InvertedPendulum-v5', 'dense', '{"reset_noise_scale": 0.001}'),
-            # TODO: Fix numerical errors in this test
-            # (
-            #     'FrozenLake-v1',
-            #     'linear',
-            #     '{"is_slippery": false}',
-            # ),
+            (
+                'FrozenLake-v1',
+                'linear',
+                '{"is_slippery": false}',
+            ),
         ],
     )
     def test_main_with_env_kwargs(self, env_id: str, policy_class: str, env_kwargs: str):
@@ -272,8 +270,7 @@ def test_main_with_default_env_kwargs(self, env_id: str, policy_class: str):
     @pytest.mark.parametrize(
         'env_id, policy_class',
         [
-            # TODO: Fix numerical errors in this test
-            # ('FrozenLake-v1', 'linear'),
+            ('FrozenLake-v1', 'linear'),
         ],
     )
     def test_main_discounted_reward_model(self, env_id: str, policy_class: str):
diff --git a/tests/tfrlrl/optimisation/test_conjugate_gradients.py b/tests/tfrlrl/optimisation/test_conjugate_gradients.py
index b167354..654b74a 100644
--- a/tests/tfrlrl/optimisation/test_conjugate_gradients.py
+++ b/tests/tfrlrl/optimisation/test_conjugate_gradients.py
@@ -98,6 +98,26 @@ def mat_v_mult_fn(v):
     assert np.dot(e_approx, np.matmul(A, e_approx)) <= np.dot(e_init, np.matmul(A, e_init))
 
 
+def test_zero_b_returns_zero_vector_without_nans():
+    """
+    Test that calculate_conjugate_gradient returns a zero vector when b is exactly zero.
+
+    When b is the zero vector, x=0 already satisfies Ax=b, so the algorithm should return immediately
+    rather than dividing rdotr (which is 0) by p.T.dot(z) (also 0), which would produce NaNs.
+
+    """
+    n = 5
+    A = _make_spd_matrix(n)
+    b = np.zeros(n)
+
+    def mat_v_mult_fn(v):
+        return np.matmul(A, v)
+
+    x = calculate_conjugate_gradient(mat_v_mult_fn, b)
+
+    np.testing.assert_array_equal(x, np.zeros(n))
+
+
 @pytest.mark.parametrize(
     'b_shape',
     [
diff --git a/tests/tfrlrl/training_algorithms/test_tnpg.py b/tests/tfrlrl/training_algorithms/test_tnpg.py
index 1837d68..aa9f507 100644
--- a/tests/tfrlrl/training_algorithms/test_tnpg.py
+++ b/tests/tfrlrl/training_algorithms/test_tnpg.py
@@ -285,11 +285,10 @@ def choice(self, n, size, replace):
 @pytest.mark.parametrize(
     'env_id, use_baseline',
     [
-        # TODO: Fix numerical errors in this test
-        # (
-        #     'FrozenLake-v1',
-        #     False,
-        # ),
+        (
+            'FrozenLake-v1',
+            False,
+        ),
         (
             'InvertedPendulum-v5',
             False,

From ab3c3bbbbf30085d4dc3084a82df7348dad281c9 Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 26 Jul 2026 19:13:02 +0100
Subject: [PATCH 18/19] fix incorrect negation

---
 src/tfrlrl/training_algorithms/tnpg.py        |  2 +-
 tests/tfrlrl/training_algorithms/test_tnpg.py | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/tfrlrl/training_algorithms/tnpg.py b/src/tfrlrl/training_algorithms/tnpg.py
index c29b380..9510868 100644
--- a/src/tfrlrl/training_algorithms/tnpg.py
+++ b/src/tfrlrl/training_algorithms/tnpg.py
@@ -54,7 +54,7 @@ def calculate_steepest_gradient_direction(
         observations=statistics.observations,
         actions=statistics.actions,
     )
-    loss = -sum(log_probabilities * tensor(statistics.total_expected_rewards))
+    loss = sum(log_probabilities * tensor(statistics.total_expected_rewards))
     loss.backward()
 
     return (
diff --git a/tests/tfrlrl/training_algorithms/test_tnpg.py b/tests/tfrlrl/training_algorithms/test_tnpg.py
index aa9f507..909612e 100644
--- a/tests/tfrlrl/training_algorithms/test_tnpg.py
+++ b/tests/tfrlrl/training_algorithms/test_tnpg.py
@@ -44,8 +44,8 @@ def test_calculate_steepest_gradient_direction(env_id: str, expected_shape: Tupl
     """
     Test that calculate_steepest_gradient_direction returns the correct policy gradient.
 
-    The steepest gradient direction is verified against central finite differences of the loss function
-    L(θ) = -sum_t log π_θ(a_t | s_t) * R_t with respect to each policy parameter element.
+    The steepest gradient direction is verified against central finite differences of the objective
+    J(θ) = sum_t log π_θ(a_t | s_t) * R_t with respect to each policy parameter element.
 
     Args:
         env_id: The Gymnasium environment ID with a discrete action space.
@@ -106,14 +106,14 @@ def test_calculate_steepest_gradient_direction(env_id: str, expected_shape: Tupl
             dict_plus[param_name][idx] += eps
             policy.set_state(dict_plus)
             log_probs_plus = policy.calculate_log_probabilities(statistics.observations, statistics.actions)
-            loss_plus = (-torch_sum(log_probs_plus * tensor(statistics.total_expected_rewards))).item()
+            objective_plus = torch_sum(log_probs_plus * tensor(statistics.total_expected_rewards)).item()
 
             dict_minus[param_name][idx] -= eps
             policy.set_state(dict_minus)
             log_probs_minus = policy.calculate_log_probabilities(statistics.observations, statistics.actions)
-            loss_minus = (-torch_sum(log_probs_minus * tensor(statistics.total_expected_rewards))).item()
+            objective_minus = torch_sum(log_probs_minus * tensor(statistics.total_expected_rewards)).item()
 
-            fd_grad[idx] = (loss_plus - loss_minus) / (2 * eps)
+            fd_grad[idx] = (objective_plus - objective_minus) / (2 * eps)
 
         fd_grad_parts.append(fd_grad.flatten())
 

From 556ebbd0a476dc2f4de1debb837a24cd0b7c8dbf Mon Sep 17 00:00:00 2001
From: tfurmston <tfurmston@users.noreply.github.com>
Date: Sun, 26 Jul 2026 19:36:12 +0100
Subject: [PATCH 19/19] Grant pull-requests write permission to Claude Code
 Review workflow

The review job could read PR data but not post comments/reviews back,
so runs completed successfully with no visible output on the PR.
---
 .github/workflows/claude-code-review.yml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/claude-code-review.yml b/.github/workflows/claude-code-review.yml
index 45b5f24..51f514f 100644
--- a/.github/workflows/claude-code-review.yml
+++ b/.github/workflows/claude-code-review.yml
@@ -10,8 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     permissions:
       contents: read
-      pull-requests: read
-      issues: read
+      pull-requests: write
       id-token: write
 
     steps: