tfrlrl-38: Add truncated natural-policy gradients.

.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:

.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/
+#.idea/
+
+.DS_Store

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
 

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:

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"]

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,
+    )

src/tfrlrl/optimisation/conjugate_gradients.py

@@ -0,0 +1,57 @@
+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)
+
+    if rdotr < tol:
+        return x
+
+    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

src/tfrlrl/policies/base.py

@@ -1,12 +1,13 @@
 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
 from torch import (
     Tensor,
     nn,
 )
+from torch.func import jacfwd
 
 
 class PolicyException(Exception):
@@ -183,7 +184,49 @@ 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).
 
         """
         ...
+
+    @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.
+
+        """
+        ...
+
+    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)