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: 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 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/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"] 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/src/tfrlrl/optimisation/conjugate_gradients.py b/src/tfrlrl/optimisation/conjugate_gradients.py new file mode 100644 index 0000000..68b877c --- /dev/null +++ b/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 diff --git a/src/tfrlrl/policies/base.py b/src/tfrlrl/policies/base.py index 97d1fc4..8b2d290 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 @@ -7,6 +7,7 @@ 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) diff --git a/src/tfrlrl/policies/dense_neural_network.py b/src/tfrlrl/policies/dense_neural_network.py index 5c4ed0e..e9f2250 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 @@ -126,15 +127,42 @@ 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( + 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]: + """ + 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 - ) + ).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 020924c..4d1f0a5 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 ( @@ -142,7 +143,35 @@ 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)) + + 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/src/tfrlrl/policies/utils.py b/src/tfrlrl/policies/utils.py new file mode 100644 index 0000000..33713cd --- /dev/null +++ b/src/tfrlrl/policies/utils.py @@ -0,0 +1,67 @@ +from typing import Dict + +from torch import ( + Tensor, + cat, +) + + +def flatten_tensor_dict(x: Dict[str, Tensor], dim: int = 0) -> 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. + dim: The starting dimension on which to flatten and concatenate the tensors. + + Returns: + The flattened tensor. The flattened tensor will have the following shape: + (start_dims, n_elems / n_start_dim_elems). + 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. + + """ + 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 new file mode 100644 index 0000000..9510868 --- /dev/null +++ b/src/tfrlrl/training_algorithms/tnpg.py @@ -0,0 +1,247 @@ +"""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 ( + 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, unflatten_tensor_dict +from tfrlrl.sampling.episodic_sampler import ( + EpisodicSampler, + RayEpisodicSampler, +) +from tfrlrl.sampling.statistics_collection import ( + EpisocidPolicyGradientStatisticsCollector, + EpisodePolicyGradientStatistics, +) + +logger = logging.getLogger(__name__) + + +def calculate_steepest_gradient_direction( + policy: BasePyTorchPolicy, statistics: EpisodePolicyGradientStatistics, 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() + + return ( + flatten_tensor_dict( + {name: param.grad for name, param in policy.network.named_parameters() if param.grad is not None}, + ) + .detach() + .numpy() + ) + + +def construct_fim_vector_product_fn( + 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. + + 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. + 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=observations, + actions=actions, + ) + jacobian_matrix = ( + flatten_tensor_dict( + jacobian, + dim=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(axis=0) + + def calculate_fim_vector_product(v: np.ndarray): + return np.matmul(jacobian_matrix.T, np.matmul(jacobian_matrix, v)) + + return calculate_fim_vector_product + + +def train_policy_gradient( + env_id: str, + policy: BasePyTorchPolicy, + n_iterations: int, + n_episodes: int, + 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, + n_iteration_logging: int = 10, + n_iters_cg: Optional[int] = None, + n_samples_fim: Optional[int] = None, + **kwargs, +) -> BasePyTorchPolicy: + """ + 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. + 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 + 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: + 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, + 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() + + logger.debug('Calculate truncated-natural policy gradient.') + sgd = calculate_steepest_gradient_direction( + policy=policy, + statistics=statistics, + optimizer=optimizer, + ) + tngd = calculate_conjugate_gradient( + mat_v_mult_fn=construct_fim_vector_product_fn( + policy=policy, + statistics=statistics, + n_samples_fim=n_samples_fim, + ), + b=sgd, + n_iters=n_iters_cg, + ) + 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) + 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 diff --git a/tests/tfrlrl/cli/test_tnpg.py b/tests/tfrlrl/cli/test_tnpg.py new file mode 100644 index 0000000..e13d595 --- /dev/null +++ b/tests/tfrlrl/cli/test_tnpg.py @@ -0,0 +1,400 @@ +"""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', + ), + ( + '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}'), + ( + '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', + [ + ('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 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..654b74a --- /dev/null +++ b/tests/tfrlrl/optimisation/test_conjugate_gradients.py @@ -0,0 +1,147 @@ +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)) + + +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', + [ + (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) diff --git a/tests/tfrlrl/policies/test_dense_neural_network.py b/tests/tfrlrl/policies/test_dense_neural_network.py index 1bb94f9..272b385 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 @@ -169,6 +171,233 @@ 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,) + + +@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, + ) + + +@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 daa6b2f..16c2dd0 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), @@ -387,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, + ) diff --git a/tests/tfrlrl/policies/test_utils.py b/tests/tfrlrl/policies/test_utils.py new file mode 100644 index 0000000..bab3ff0 --- /dev/null +++ b/tests/tfrlrl/policies/test_utils.py @@ -0,0 +1,251 @@ +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, unflatten_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[indx] + indx += 1 + + 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', + [ + ( + '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] + + +@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() 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. diff --git a/tests/tfrlrl/training_algorithms/test_tnpg.py b/tests/tfrlrl/training_algorithms/test_tnpg.py new file mode 100644 index 0000000..909612e --- /dev/null +++ b/tests/tfrlrl/training_algorithms/test_tnpg.py @@ -0,0 +1,406 @@ +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 ( + sum as torch_sum, +) +from torch import ( + tensor, +) +from torch.optim import ( + 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, + train_policy_gradient, +) + + +@pytest.mark.parametrize( + 'env_id, expected_shape', + [ + ('FrozenLake-v1', (49,)), + ('InvertedPendulum-v5', (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 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. + 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, np.ndarray) + 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) + 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) + objective_minus = torch_sum(log_probs_minus * tensor(statistics.total_expected_rewards)).item() + + fd_grad[idx] = (objective_plus - objective_minus) / (2 * eps) + + fd_grad_parts.append(fd_grad.flatten()) + + policy.set_state(policy_dict) + + fd_gradient = np.concatenate(fd_grad_parts) + + np.testing.assert_almost_equal( + 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=2, max_value=10), +) +@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) + + +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', + [ + ( + '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, + n_iters_cg=2, + n_samples_fim=20, + ) + + 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), + n_iters_cg=2, + n_samples_fim=20, + ) + + 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