Port soem fixes / additions learned from implementign itpal_jax

.gitignore

@@ -1,5 +1,6 @@
 __pycache__
 .venv
+.vscode
 wandb
 *.egg-info/
 test.py

README.md

@@ -8,6 +8,9 @@
 
 Fancy RL provides a minimalistic and efficient implementation of Proximal Policy Optimization (PPO) and Trust Region Policy Layers (TRPL) using primitives from [torchrl](https://pypi.org/project/torchrl/). This library focuses on providing clean, understandable code and reusable modules while leveraging the powerful functionalities of torchrl.
 
+| :exclamation: This project is still WIP and not ready to be used. (Problems with torchdict routing are a pain to debug...) |
+| -------------------------------------------------------------------------------------------------------------------------- |
+
 ## Installation
 
 Fancy RL requires Python 3.7-3.11. (TorchRL currently does not support Python 3.12)
@@ -33,8 +36,9 @@ Fancy RL provides two main components:
    For environments, you can pass any [gymnasium](https://gymnasium.farama.org/) or [Fancy Gym](https://alrhub.github.io/fancy_gym/) environment ID as a string, a function returning a gymnasium or torchrl environment, an already instantiated gymnasium or torchrl environment, or a dict that will be passed to gymnasium.make. Check 'example/example.py' for a more complete usage example.
 
 2. **Additional Modules for TRPL**: Designed to integrate with torchrl's primitives-first approach, these modules are ideal for building custom algorithms with precise trust region projections.
+   Oh, you want documentation for these? To bad... (TODO)
 
-### Background on Trust Region Policy Layers (TRPL)
+## Background on Trust Region Policy Layers (TRPL)
 
 Trust region methods are essential in reinforcement learning for ensuring robust policy updates. Traditional methods like TRPO and PPO use approximations, which can sometimes violate constraints or fail to find optimal solutions. To address these issues, TRPL provides differentiable neural network layers that enforce trust regions through closed-form projections for deep Gaussian policies. These layers formalize trust regions individually for each state and complement existing reinforcement learning algorithms.
 
@@ -45,9 +49,38 @@ The TRPL implementation in Fancy RL includes projections based on the Kullback-L
 To run the test suite:
 
 ```bash
-pytest test/test_ppo.py
+pytest test/
 ```
 
+## Status
+
+### Implemented Features
+- Proximal Policy Optimization (PPO) algorithm
+- Trust Region Policy Layers (TRPL) algorithm (WIP)
+- Support for continuous and discrete action spaces
+- Multiple projection methods (Rewritten for MIT License Compatability):
+  - KL Divergence projection
+  - Frobenius norm projection
+  - Wasserstein distance projection
+  - Identity projection (Eq to PPO)
+- Configurable neural network architectures for actor and critic
+- Logging support (Terminal and WandB, extendable)
+
+### TODO
+- [ ] All PPO Tests green
+- [ ] Better / more logging
+- [ ] Test / Benchmark PPO
+- [ ] Refactor Modules for TRPL
+- [ ] Get TRPL working
+- [ ] All TRPL Tests green
+- [ ] Make contextual covariance optional
+- [ ] Allow full-cov via chol
+- [ ] Test / Benchmark TRPL
+- [ ] Write docs / extend README
+- [ ] Test func of non-gym envs
+- [ ] Implement SAC
+- [ ] Implement VLEARN
+
 ## Contributing
 
 Contributions are welcome! Feel free to open issues or submit pull requests to enhance the library.

fancy_rl/__init__.py

@@ -1,9 +1,3 @@
-import gymnasium
+from fancy_rl.algos import PPO, TRPL #, VLEARN
-try:
-    import fancy_gym
-except ImportError:
-    pass
 
-from fancy_rl.algos import PPO
+__all__ = ['PPO', 'TRPL']
-
-__all__ = ["PPO"]

fancy_rl/algos/__init__.py

@@ -1 +1,5 @@
 from fancy_rl.algos.ppo import PPO
+from fancy_rl.algos.trpl import TRPL
+#from fancy_rl.algos.vlearn import VLEARN
+
+__all__ = ['PPO', 'TRPL']

fancy_rl/algos/algo.py

@@ -0,0 +1,115 @@
+import torch
+import gymnasium as gym
+from torchrl.envs import GymEnv, TransformedEnv, Compose, RewardSum, StepCounter, SerialEnv
+from torchrl.record import VideoRecorder
+from abc import ABC
+import pdb
+import numpy as np
+from tensordict import TensorDict
+from torchrl.envs import GymWrapper, TransformedEnv
+from torchrl.envs import BatchSizeTransform
+
+from fancy_rl.loggers import TerminalLogger
+
+class Algo(ABC):
+    def __init__(
+        self,
+        env_spec,
+        loggers,
+        optimizers,
+        learning_rate,
+        n_steps,
+        batch_size,
+        n_epochs,
+        gamma,
+        total_timesteps,
+        eval_interval,
+        eval_deterministic,
+        entropy_coef,
+        critic_coef,
+        normalize_advantage,
+        device=None,
+        eval_episodes=10,
+        env_spec_eval=None,
+    ):
+        self.env_spec = env_spec
+        self.env_spec_eval = env_spec_eval if env_spec_eval is not None else env_spec
+        self.loggers = loggers if loggers != None else [TerminalLogger(None, None)]
+        self.optimizers = optimizers
+        self.learning_rate = learning_rate
+        self.n_steps = n_steps
+        self.batch_size = batch_size
+        self.n_epochs = n_epochs
+        self.gamma = gamma
+        self.total_timesteps = total_timesteps
+        self.eval_interval = eval_interval
+        self.eval_deterministic = eval_deterministic
+        self.entropy_coef = entropy_coef
+        self.critic_coef = critic_coef
+        self.normalize_advantage = normalize_advantage
+        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
+        self.eval_episodes = eval_episodes
+
+    def make_env(self, eval=False):
+        env_spec = self.env_spec_eval if eval else self.env_spec
+        env = self._wrap_env(env_spec)
+        env.reset()
+        return env
+
+    def _wrap_env(self, env_spec):
+        # If given an existing env, ensure it's properly batched
+        if isinstance(env_spec, (GymEnv, GymWrapper)):
+            if not env_spec.batch_size:
+                raise ValueError("Environment must be batched")
+            return env_spec
+        
+        # Handle callable without wrapping the recursive call
+        if callable(env_spec):
+            return self._wrap_env(env_spec())
+        
+        # Create new batched environment using SerialEnv
+        if isinstance(env_spec, str):
+            env = SerialEnv(1, lambda: GymEnv(env_spec, device=self.device))
+        elif isinstance(env_spec, gym.Env):
+            wrapped_env = GymWrapper(env_spec, device=self.device)
+            if wrapped_env.batch_size:
+                env = wrapped_env
+            else:
+                env = SerialEnv(1, lambda: wrapped_env)
+        else:
+            raise ValueError(
+                f"env_spec must be a string, callable, Gymnasium environment, or GymEnv, "
+                f"got {type(env_spec)}"
+            )
+        
+        return env
+
+    def train_step(self, batch):
+        raise NotImplementedError("train_step method must be implemented in subclass.")
+
+    def train(self):
+        raise NotImplementedError("train method must be implemented in subclass.")
+
+    def evaluate(self, epoch):
+        raise NotImplementedError("evaluate method must be implemented in subclass.")
+
+    def predict(
+        self,
+        tensordict,
+        state=None,
+        deterministic=False
+    ):
+        with torch.no_grad():
+            # If numpy array, convert to TensorDict
+            if isinstance(tensordict, np.ndarray):
+                tensordict = TensorDict(
+                    {"observation": torch.from_numpy(tensordict).float()}, 
+                    batch_size=[]
+                )
+            
+            # Move to device
+            tensordict = tensordict.to(self.device)
+            
+            # Get action from policy
+            action_td = self.prob_actor(tensordict)
+            return action_td

fancy_rl/algos/on_policy.py

@@ -5,59 +5,57 @@ from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
 from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
 from torchrl.envs.libs.gym import GymWrapper
 from torchrl.envs import ExplorationType, set_exploration_type
-from torchrl.record import VideoRecorder
-from abc import ABC, abstractmethod
 
+from fancy_rl.loggers import TerminalLogger
+from fancy_rl.algos.algo import Algo
 
-class OnPolicy(ABC):
+class OnPolicy(Algo):
     def __init__(
         self,
-        policy,
         env_spec,
-        loggers,
         optimizers,
-        learning_rate,
+        loggers=None,
-        n_steps,
+        learning_rate=3e-4,
-        batch_size,
+        n_steps=2048,
-        n_epochs,
+        batch_size=64,
-        gamma,
+        n_epochs=10,
-        gae_lambda,
+        gamma=0.99,
-        total_timesteps,
+        total_timesteps=1e6,
-        eval_interval,
+        eval_interval=2048,
-        eval_deterministic,
+        eval_deterministic=True,
-        entropy_coef,
+        entropy_coef=0.01,
-        critic_coef,
+        critic_coef=0.5,
-        normalize_advantage,
+        normalize_advantage=True,
-        clip_range=0.2,
-        device=None,
-        eval_episodes=10,
         env_spec_eval=None,
+        eval_episodes=10,
+        device=None,
     ):
-        self.policy = policy
+        device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.env_spec = env_spec
+
-        self.env_spec_eval = env_spec_eval if env_spec_eval is not None else env_spec
+        super().__init__(
-        self.loggers = loggers
+            env_spec=env_spec,
-        self.optimizers = optimizers
+            loggers=loggers,
-        self.learning_rate = learning_rate
+            optimizers=optimizers,
-        self.n_steps = n_steps
+            learning_rate=learning_rate,
-        self.batch_size = batch_size
+            n_steps=n_steps,
-        self.n_epochs = n_epochs
+            batch_size=batch_size,
-        self.gamma = gamma
+            n_epochs=n_epochs,
-        self.gae_lambda = gae_lambda
+            gamma=gamma,
-        self.total_timesteps = total_timesteps
+            total_timesteps=total_timesteps,
-        self.eval_interval = eval_interval
+            eval_interval=eval_interval,
-        self.eval_deterministic = eval_deterministic
+            eval_deterministic=eval_deterministic,
-        self.entropy_coef = entropy_coef
+            entropy_coef=entropy_coef,
-        self.critic_coef = critic_coef
+            critic_coef=critic_coef,
-        self.normalize_advantage = normalize_advantage
+            normalize_advantage=normalize_advantage,
-        self.clip_range = clip_range
+            device=device,
-        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
+            env_spec_eval=env_spec_eval,
-        self.eval_episodes = eval_episodes
+            eval_episodes=eval_episodes,
+        )
 
         # Create collector
         self.collector = SyncDataCollector(
             create_env_fn=lambda: self.make_env(eval=False),
-            policy=self.policy,
+            policy=self.prob_actor,
             frames_per_batch=self.n_steps,
             total_frames=self.total_timesteps,
             device=self.device,
@@ -73,29 +71,25 @@ class OnPolicy(ABC):
             batch_size=self.batch_size,
         )
 
-    def make_env(self, eval=False):
+    def pre_process_batch(self, batch):
-        """Creates an environment and wraps it if necessary."""
+        return batch
-        env_spec = self.env_spec_eval if eval else self.env_spec
+
-        if isinstance(env_spec, str):
+    def post_process_batch(self, batch):
-            env = gym.make(env_spec)
+        pass
-            env = GymWrapper(env)
-        elif callable(env_spec):
-            env = env_spec()
-            if isinstance(env, gym.Env):
-                env = GymWrapper(env)
-        elif isinstance(env, gym.Env):
-            env = GymWrapper(env)
-        else:
-            raise ValueError("env_spec must be a string or a callable that returns an environment.")
-        return env
 
     def train_step(self, batch):
+        batch = self.pre_process_batch(batch)
+        
         for optimizer in self.optimizers.values():
             optimizer.zero_grad()
-        loss = self.loss_module(batch)
+        losses = self.loss_module(batch)
+        loss = sum(losses.values())  # Sum all losses
         loss.backward()
         for optimizer in self.optimizers.values():
             optimizer.step()
+        
+        self.post_process_batch(batch)
+        
         return loss
 
     def train(self):
@@ -115,7 +109,7 @@ class OnPolicy(ABC):
                     batch = batch.to(self.device)
                     loss = self.train_step(batch)
                     for logger in self.loggers:
-                        logger.log_scalar({"loss": loss.item()}, step=collected_frames)
+                        logger.log_scalar("loss", loss.item(), step=collected_frames)
 
             if (t + 1) % self.eval_interval == 0:
                 self.evaluate(t)
@@ -130,7 +124,7 @@ class OnPolicy(ABC):
         for _ in range(self.eval_episodes):
             with torch.no_grad(), set_exploration_type(ExplorationType.MODE):
                 td_test = eval_env.rollout(
-                    policy=self.policy,
+                    policy=self.actor,
                     auto_reset=True,
                     auto_cast_to_device=True,
                     break_when_any_done=True,
@@ -143,7 +137,3 @@ class OnPolicy(ABC):
         avg_return = torch.cat(test_rewards, 0).mean().item()
         for logger in self.loggers:
             logger.log_scalar({"eval_avg_return": avg_return}, step=epoch)
-
-def dump_video(module):
-    if isinstance(module, VideoRecorder):
-        module.dump()

fancy_rl/algos/ppo.py

@@ -1,9 +1,10 @@
 import torch
-from torchrl.modules import ActorValueOperator, ProbabilisticActor
+from torchrl.modules import ProbabilisticActor
 from torchrl.objectives import ClipPPOLoss
 from torchrl.objectives.value.advantages import GAE
+from torchrl.data.tensor_specs import DiscreteTensorSpec
 from fancy_rl.algos.on_policy import OnPolicy
-from fancy_rl.policy import Actor, Critic, SharedModule
+from fancy_rl.policy import Actor, Critic
 
 class PPO(OnPolicy):
     def __init__(
@@ -14,7 +15,6 @@ class PPO(OnPolicy):
         critic_hidden_sizes=[64, 64],
         actor_activation_fn="Tanh",
         critic_activation_fn="Tanh",
-        shared_stem_sizes=[64],
         learning_rate=3e-4,
         n_steps=2048,
         batch_size=64,
@@ -31,34 +31,49 @@ class PPO(OnPolicy):
         device=None,
         env_spec_eval=None,
         eval_episodes=10,
+        full_covariance=False,
     ):
+        self.clip_range = clip_range
+
         device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = device
 
         # Initialize environment to get observation and action space sizes
         self.env_spec = env_spec
         env = self.make_env()
-        obs_space = env.observation_space
-        act_space = env.action_space
         
-        # Define the shared, actor, and critic modules
+        # Get spaces from specs for parallel env
-        self.shared_module = SharedModule(obs_space, shared_stem_sizes, actor_activation_fn, device)
+        self.obs_space = env.observation_spec
-        self.actor = Actor(self.shared_module, act_space, actor_hidden_sizes, actor_activation_fn, device)
+        self.act_space = env.action_spec
-        self.critic = Critic(self.shared_module, critic_hidden_sizes, critic_activation_fn, device)
         
-        # Combine into an ActorValueOperator
+        self.discrete = isinstance(self.act_space, DiscreteTensorSpec)
-        self.ac_module = ActorValueOperator(
-            self.shared_module,
-            self.actor,
-            self.critic
-        )
 
-        # Define the policy as a ProbabilisticActor
+        self.critic = Critic(self.obs_space, critic_hidden_sizes, critic_activation_fn, device)
-        policy = ProbabilisticActor(
+        self.actor = Actor(self.obs_space, self.act_space, actor_hidden_sizes, actor_activation_fn, device, full_covariance=full_covariance)
-            module=self.ac_module.get_policy_operator(),
+
-            in_keys=["loc", "scale"],
+        if self.discrete:
+            distribution_class = torch.distributions.Categorical
+            self.prob_actor = ProbabilisticActor(
+                module=self.actor,
+                distribution_class=distribution_class,
+                return_log_prob=True,
+                in_keys=["logits"],
                 out_keys=["action"],
-            distribution_class=torch.distributions.Normal,
+            )
-            return_log_prob=True
+        else:
+            if full_covariance:
+                distribution_class = torch.distributions.MultivariateNormal
+                in_keys = ["loc", "scale_tril"]
+            else:
+                distribution_class = torch.distributions.Normal
+                in_keys = ["loc", "scale"]
+
+            self.prob_actor = ProbabilisticActor(
+                module=self.actor,
+                distribution_class=distribution_class,
+                return_log_prob=True,
+                in_keys=in_keys,
+                out_keys=["action"]
             )
 
         optimizers = {
@@ -66,25 +81,7 @@ class PPO(OnPolicy):
             "critic": torch.optim.Adam(self.critic.parameters(), lr=learning_rate)
         }
 
-        self.adv_module = GAE(
-            gamma=self.gamma,
-            lmbda=self.gae_lambda,
-            value_network=self.critic,
-            average_gae=False,
-        )
-
-        self.loss_module = ClipPPOLoss(
-            actor_network=self.actor,
-            critic_network=self.critic,
-            clip_epsilon=self.clip_range,
-            loss_critic_type='MSELoss',
-            entropy_coef=self.entropy_coef,
-            critic_coef=self.critic_coef,
-            normalize_advantage=self.normalize_advantage,
-        )
-
         super().__init__(
-            policy=policy,
             env_spec=env_spec,
             loggers=loggers,
             optimizers=optimizers,
@@ -93,15 +90,30 @@ class PPO(OnPolicy):
             batch_size=batch_size,
             n_epochs=n_epochs,
             gamma=gamma,
-            gae_lambda=gae_lambda,
             total_timesteps=total_timesteps,
             eval_interval=eval_interval,
             eval_deterministic=eval_deterministic,
             entropy_coef=entropy_coef,
             critic_coef=critic_coef,
             normalize_advantage=normalize_advantage,
-            clip_range=clip_range,
             device=device,
             env_spec_eval=env_spec_eval,
             eval_episodes=eval_episodes,
         )
+
+        self.adv_module = GAE(
+            gamma=self.gamma,
+            lmbda=gae_lambda,
+            value_network=self.critic,
+            average_gae=False,
+        )
+
+        self.loss_module = ClipPPOLoss(
+            actor_network=self.actor,
+            critic_network=self.critic,
+            clip_epsilon=self.clip_range,
+            loss_critic_type='l2',
+            entropy_coef=self.entropy_coef,
+            critic_coef=self.critic_coef,
+            normalize_advantage=self.normalize_advantage,
+        )

fancy_rl/algos/trpl.py

@@ -0,0 +1,200 @@
+import torch
+from torch import nn
+from typing import Dict, Any, Optional
+from torchrl.data.tensor_specs import DiscreteTensorSpec
+from torchrl.modules import ProbabilisticActor, ValueOperator
+from torchrl.objectives import ClipPPOLoss
+from torchrl.collectors import SyncDataCollector
+from torchrl.data import TensorDictReplayBuffer, LazyTensorStorage, SamplerWithoutReplacement
+from torchrl.objectives.value import GAE
+from fancy_rl.algos.on_policy import OnPolicy
+from fancy_rl.policy import Actor, Critic
+from fancy_rl.projections import get_projection, BaseProjection
+from fancy_rl.objectives import TRPLLoss
+from copy import deepcopy
+from tensordict.nn import TensorDictModule
+from tensordict import TensorDict
+from torch.distributions import Categorical, MultivariateNormal, Normal
+
+
+class ProjectedActor(TensorDictModule):
+    def __init__(self, raw_actor, old_actor, projection):
+        combined_module = self.CombinedModule(raw_actor, old_actor, projection)
+        super().__init__(
+            module=combined_module,
+            in_keys=raw_actor.in_keys,
+            out_keys=raw_actor.out_keys
+        )
+        self.raw_actor = raw_actor
+        self.old_actor = old_actor
+        self.projection = projection
+        self.discrete = raw_actor.discrete
+        self.full_covariance = raw_actor.full_covariance
+
+    class CombinedModule(nn.Module):
+        def __init__(self, raw_actor, old_actor, projection):
+            super().__init__()
+            self.raw_actor = raw_actor
+            self.old_actor = old_actor
+            self.projection = projection
+
+        def forward(self, tensordict):
+            # Convert the tuple outputs to TensorDict
+            raw_params = self.raw_actor(tensordict)
+            if isinstance(raw_params, tuple):
+                raw_params = TensorDict({
+                    "loc": raw_params[0], 
+                    "scale": raw_params[1]
+                }, batch_size=[raw_params[0].shape[0]])  # Use the first dimension of the tensor as batch size
+                    
+            old_params = self.old_actor(tensordict)
+            if isinstance(old_params, tuple):
+                old_params = TensorDict({
+                    "loc": old_params[0], 
+                    "scale": old_params[1]
+                }, batch_size=[old_params[0].shape[0]])  # Use the first dimension of the tensor as batch size
+
+            # Now combine them
+            combined_params = TensorDict({
+                **raw_params,
+                **{f"old_{key}": value for key, value in old_params.items()}
+            }, batch_size=[raw_params["loc"].shape[0]])  # Use the first dimension of loc tensor as batch size
+            
+            projected_params = self.projection(combined_params)
+            return projected_params
+
+    def get_dist(self, tensordict):
+        # Forward the observation through the network
+        out = self.forward(tensordict)
+        if self.discrete:
+            return Categorical(logits=out["logits"])
+        else:
+            if self.full_covariance:
+                return MultivariateNormal(loc=out["loc"], scale_tril=out["scale_tril"])
+            else:
+                return Normal(loc=out["loc"], scale=out["scale"])
+
+class TRPL(OnPolicy):
+    def __init__(
+        self,
+        env_spec,
+        loggers=None,
+        actor_hidden_sizes=[64, 64],
+        critic_hidden_sizes=[64, 64],
+        actor_activation_fn="Tanh",
+        critic_activation_fn="Tanh",
+        learning_rate=3e-4,
+        n_steps=2048,
+        batch_size=64,
+        n_epochs=10,
+        gamma=0.99,
+        gae_lambda=0.95,
+        projection_class="identity_projection",
+        trust_region_coef=10.0,
+        trust_region_bound_mean=0.1,
+        trust_region_bound_cov=0.001,
+        total_timesteps=1e6,
+        eval_interval=2048,
+        eval_deterministic=True,
+        entropy_coef=0.01,
+        critic_coef=0.5,
+        normalize_advantage=False,
+        device=None,
+        env_spec_eval=None,
+        eval_episodes=10,
+        full_covariance=False,
+    ):
+        device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = device
+
+        # Initialize environment to get observation and action space sizes
+        self.env_spec = env_spec
+        env = self.make_env()
+
+        # Get spaces from specs for parallel env
+        self.obs_space = env.observation_spec
+        self.act_space = env.action_spec
+
+        assert not isinstance(self.act_space, DiscreteTensorSpec), "TRPL does not support discrete action spaces"
+
+        self.critic = Critic(self.obs_space, critic_hidden_sizes, critic_activation_fn, device)
+        self.raw_actor = Actor(self.obs_space, self.act_space, actor_hidden_sizes, actor_activation_fn, device, full_covariance=full_covariance)
+        self.old_actor = Actor(self.obs_space, self.act_space, actor_hidden_sizes, actor_activation_fn, device, full_covariance=full_covariance)
+
+        # Handle projection_class
+        if isinstance(projection_class, str):
+            projection_class = get_projection(projection_class)
+        elif not issubclass(projection_class, BaseProjection):
+            raise ValueError("projection_class must be a string or a subclass of BaseProjection")
+
+        self.projection = projection_class(
+            in_keys=["loc", "scale_tril", "old_loc", "old_scale_tril"] if full_covariance else ["loc", "scale", "old_loc", "old_scale"],
+            out_keys=["loc", "scale_tril"] if full_covariance else ["loc", "scale"],
+            mean_bound=trust_region_bound_mean,
+            cov_bound=trust_region_bound_cov
+        )
+
+        self.actor = ProjectedActor(self.raw_actor, self.old_actor, self.projection)
+
+        if full_covariance:
+            distribution_class = torch.distributions.MultivariateNormal
+            distribution_kwargs = {"loc": "loc", "scale_tril": "scale_tril"}
+        else:
+            distribution_class = torch.distributions.Normal
+            distribution_kwargs = {"loc": "loc", "scale": "scale"}
+
+        self.prob_actor = ProbabilisticActor(
+            module=self.actor,
+            distribution_class=distribution_class,
+            return_log_prob=True,
+            in_keys=distribution_kwargs,
+        )
+
+        self.trust_region_coef = trust_region_coef
+        self.loss_module = TRPLLoss(
+            actor_network=self.actor,
+            old_actor_network=self.old_actor,
+            critic_network=self.critic,
+            projection=self.projection,
+            entropy_coef=entropy_coef,
+            critic_coef=critic_coef,
+            trust_region_coef=trust_region_coef,
+            normalize_advantage=normalize_advantage,
+        )
+
+        optimizers = {
+            "actor": torch.optim.Adam(self.raw_actor.parameters(), lr=learning_rate),
+            "critic": torch.optim.Adam(self.critic.parameters(), lr=learning_rate)
+        }
+
+        super().__init__(
+            env_spec=env_spec,
+            loggers=loggers,
+            optimizers=optimizers,
+            learning_rate=learning_rate,
+            n_steps=n_steps,
+            batch_size=batch_size,
+            n_epochs=n_epochs,
+            gamma=gamma,
+            total_timesteps=total_timesteps,
+            eval_interval=eval_interval,
+            eval_deterministic=eval_deterministic,
+            entropy_coef=entropy_coef,
+            critic_coef=critic_coef,
+            normalize_advantage=normalize_advantage,
+            device=device,
+            env_spec_eval=env_spec_eval,
+            eval_episodes=eval_episodes,
+        )
+        self.adv_module = GAE(
+            gamma=self.gamma,
+            lmbda=gae_lambda,
+            value_network=self.critic,
+            average_gae=False,
+        )
+
+    def update_old_policy(self):
+        self.old_actor.load_state_dict(self.raw_actor.state_dict())
+
+    def post_update(self):
+        self.update_old_policy()

fancy_rl/algos/vlearn.py

@@ -0,0 +1,113 @@
+import torch
+from torch import nn
+from typing import Dict, Any, Optional
+from torchrl.modules import ProbabilisticActor, ValueOperator
+from torchrl.collectors import SyncDataCollector
+from torchrl.data import TensorDictReplayBuffer, LazyMemmapStorage
+
+from fancy_rl.objectives.vlearn import VLEARNLoss
+from fancy_rl.projections import get_vlearn_projection
+from fancy_rl.utils import get_squashed_normal
+
+class VLEARN:
+    def __init__(self, env_id: str, device: str = "cpu", **kwargs: Any):
+        self.device = torch.device(device)
+        self.env = get_env(env_id)
+        
+        self.projection = get_vlearn_projection(**kwargs.get("projection", {}))
+        
+        actor = get_actor(self.env, **kwargs.get("actor", {}))
+        self.actor = ProbabilisticActor(
+            actor,
+            in_keys=["observation"],
+            out_keys=["loc", "scale"],
+            distribution_class=get_squashed_normal(),
+            return_log_prob=True
+        ).to(self.device)
+        self.old_actor = self.actor.clone()
+        
+        self.critic = ValueOperator(
+            module=get_critic(self.env, **kwargs.get("critic", {})),
+            in_keys=["observation"]
+        ).to(self.device)
+        
+        self.collector = SyncDataCollector(
+            self.env,
+            self.actor,
+            frames_per_batch=kwargs.get("frames_per_batch", 1000),
+            total_frames=kwargs.get("total_frames", -1),
+            device=self.device,
+        )
+        
+        self.replay_buffer = TensorDictReplayBuffer(
+            storage=LazyMemmapStorage(kwargs.get("buffer_size", 100000)),
+            batch_size=kwargs.get("batch_size", 256),
+        )
+        
+        self.loss_module = VLEARNLoss(
+            actor_network=self.actor,
+            critic_network=self.critic,
+            old_actor_network=self.old_actor,
+            projection=self.projection,
+            **kwargs.get("loss", {})
+        )
+        
+        self.optimizers = nn.ModuleDict({
+            "policy": torch.optim.Adam(self.actor.parameters(), lr=kwargs.get("lr_policy", 3e-4)),
+            "critic": torch.optim.Adam(self.critic.parameters(), lr=kwargs.get("lr_critic", 3e-4))
+        })
+
+        self.update_policy_interval = kwargs.get("update_policy_interval", 1)
+        self.update_critic_interval = kwargs.get("update_critic_interval", 1)
+        self.target_update_interval = kwargs.get("target_update_interval", 1)
+        self.polyak_weight_critic = kwargs.get("polyak_weight_critic", 0.995)
+
+    def train(self, num_iterations: int = 1000) -> None:
+        for i in range(num_iterations):
+            data = next(self.collector)
+            self.replay_buffer.extend(data)
+            
+            batch = self.replay_buffer.sample().to(self.device)
+            loss_dict = self.loss_module(batch)
+            
+            if i % self.update_policy_interval == 0:
+                self.optimizers["policy"].zero_grad()
+                loss_dict["policy_loss"].backward()
+                self.optimizers["policy"].step()
+            
+            if i % self.update_critic_interval == 0:
+                self.optimizers["critic"].zero_grad()
+                loss_dict["critic_loss"].backward()
+                self.optimizers["critic"].step()
+            
+            if i % self.target_update_interval == 0:
+                self.critic.update_target_params(self.polyak_weight_critic)
+            
+            self.old_actor.load_state_dict(self.actor.state_dict())
+            self.collector.update_policy_weights_()
+
+            if i % 100 == 0:
+                eval_reward = self.eval()
+                print(f"Iteration {i}, Eval reward: {eval_reward}")
+
+    def eval(self, num_episodes: int = 10) -> float:
+        total_reward = 0
+        for _ in range(num_episodes):
+            td = self.env.reset()
+            done = False
+            while not done:
+                with torch.no_grad():
+                    action = self.actor(td.to(self.device))["action"]
+                td = self.env.step(action)
+                total_reward += td["reward"].item()
+                done = td["done"].item()
+        return total_reward / num_episodes
+
+    def save_policy(self, path: str) -> None:
+        torch.save(self.actor.state_dict(), f"{path}/actor.pth")
+        torch.save(self.critic.state_dict(), f"{path}/critic.pth")
+
+    def load_policy(self, path: str) -> None:
+        self.actor.load_state_dict(torch.load(f"{path}/actor.pth"))
+        self.critic.load_state_dict(torch.load(f"{path}/critic.pth"))
+        self.old_actor.load_state_dict(self.actor.state_dict())

fancy_rl/loggers/__init__.py

@@ -0,0 +1 @@
+from fancy_rl.loggers.terminal import TerminalLogger

fancy_rl/loggers/terminal.py

@@ -0,0 +1,41 @@
+from typing import Dict, Sequence, Union, Optional
+from torch import Tensor
+
+from torchrl.record.loggers.common import Logger
+
+class TerminalLogger(Logger):
+    """Logger that prints to the terminal."""
+
+    def __init__(self, exp_name: str, log_dir: str) -> None:
+        super().__init__(exp_name, log_dir)
+
+    def _create_experiment(self):
+        # No need to create any experiment object for terminal logging
+        pass
+
+    def log_scalar(self, name: str, value: float, step: Optional[int] = None) -> None:
+        """Logs a scalar value to the terminal."""
+        if step is not None:
+            print(f"Step {step}: {name} - {value}")
+        else:
+            print(f"{name}: {value}")
+
+    def log_video(self, name: str, video: Tensor, step: Optional[int] = None, **kwargs) -> None:
+        """Logs video information to the terminal."""
+        if step is not None:
+            print(f"Step {step}: Logging video {name}")
+        else:
+            print(f"Logging video {name}")
+
+    def log_hparams(self, cfg: Union[Dict, Sequence]) -> None:
+        """Logs hyperparameters to the terminal."""
+        print("Hyperparameters:")
+        for key, value in cfg.items():
+            print(f"{key}: {value}")
+
+    def __repr__(self) -> str:
+        return "TerminalLogger"
+
+    def log_histogram(self, name: str, data: Sequence, **kwargs) -> None:
+        """Logs histogram data to the terminal."""
+        print(f"Logging histogram {name}")

fancy_rl/norm.py

@@ -0,0 +1,27 @@
+import torch as th
+from torch.distributions.multivariate_normal import _batch_mahalanobis
+
+
+def mahalanobis_alt(u, v, std):
+    delta = u - v
+    return th.triangular_solve(delta, std, upper=False)[0].pow(2).sum([-2, -1])
+
+
+def mahalanobis(u, v, chol):
+    delta = u - v
+    return _batch_mahalanobis(chol, delta)
+
+
+def frob_sq(diff, is_spd=False):
+    # If diff is spd, we can use a (probably) more performant algorithm
+    if is_spd:
+        return _frob_sq_spd(diff)
+    return th.norm(diff, p='fro', dim=tuple(range(1, diff.dim()))).pow(2)
+
+
+def _frob_sq_spd(diff):
+    return _batch_trace(diff @ diff)
+
+
+def _batch_trace(x):
+    return th.diagonal(x, dim1=-2, dim2=-1).sum(-1)

fancy_rl/objectives/__init__.py

@@ -0,0 +1 @@
+from fancy_rl.objectives.trpl import TRPLLoss

fancy_rl/objectives/trpl.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import contextlib
+
+import math
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Tuple
+
+import torch
+from tensordict import TensorDict, TensorDictBase, TensorDictParams
+from tensordict.nn import (
+    dispatch,
+    ProbabilisticTensorDictModule,
+    ProbabilisticTensorDictSequential,
+    TensorDictModule,
+)
+from tensordict.utils import NestedKey
+from torch import distributions as d
+
+from torchrl.objectives.common import LossModule
+
+from torchrl.objectives.utils import (
+    _cache_values,
+    _clip_value_loss,
+    _GAMMA_LMBDA_DEPREC_ERROR,
+    _reduce,
+    default_value_kwargs,
+    distance_loss,
+    ValueEstimators,
+)
+from torchrl.objectives.value import (
+    GAE,
+    TD0Estimator,
+    TD1Estimator,
+    TDLambdaEstimator,
+    VTrace,
+)
+
+from torchrl.objectives.ppo import PPOLoss
+from fancy_rl.projections import get_projection
+
+class TRPLLoss(PPOLoss):
+    def __init__(
+        self,
+        actor_network: ProbabilisticTensorDictSequential,
+        old_actor_network: ProbabilisticTensorDictSequential,
+        critic_network: TensorDictModule,
+        projection: any,
+        entropy_coef: float = 0.01,
+        critic_coef: float = 1.0,
+        trust_region_coef: float = 10.0,
+        normalize_advantage: bool = False,
+        **kwargs,
+    ):
+        super().__init__(
+            actor_network=actor_network,
+            critic_network=critic_network,
+            entropy_coef=entropy_coef,
+            critic_coef=critic_coef,
+            normalize_advantage=normalize_advantage,
+            **kwargs,
+        )
+        self.old_actor_network = old_actor_network
+        self.projection = projection
+        self.trust_region_coef = trust_region_coef
+
+    @property
+    def out_keys(self):
+        if self._out_keys is None:
+            keys = ["loss_objective", "tr_loss"]
+            if self.entropy_bonus:
+                keys.extend(["entropy", "loss_entropy"])
+            if self.critic_coef:
+                keys.append("loss_critic")
+            keys.append("ESS")
+            self._out_keys = keys
+        return self._out_keys
+
+    @out_keys.setter
+    def out_keys(self, values):
+        self._out_keys = values
+
+    def _trust_region_loss(self, tensordict):
+        old_distribution = self.old_actor_network(tensordict)
+        new_distribution = self.actor_network(tensordict)
+        return self.projection.get_trust_region_loss(new_distribution, old_distribution)
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDict:
+        tensordict = tensordict.clone(False)
+        advantage = tensordict.get(self.tensor_keys.advantage, None)
+        if advantage is None:
+            self.value_estimator(
+                tensordict,
+                params=self._cached_critic_network_params_detached,
+                target_params=self.target_critic_network_params,
+            )
+            advantage = tensordict.get(self.tensor_keys.advantage)
+        if self.normalize_advantage and advantage.numel() > 1:
+            loc = advantage.mean()
+            scale = advantage.std().clamp_min(1e-6)
+            advantage = (advantage - loc) / scale
+
+        log_weight, dist, kl_approx = self._log_weight(tensordict)
+        trust_region_loss_unscaled = self._trust_region_loss(tensordict)
+        
+        with torch.no_grad():
+            lw = log_weight.squeeze()
+            ess = (2 * lw.logsumexp(0) - (2 * lw).logsumexp(0)).exp()
+            batch = log_weight.shape[0]
+
+        surrogate_gain = log_weight.exp() * advantage
+        trust_region_loss = trust_region_loss_unscaled * self.trust_region_coef
+
+        loss = -surrogate_gain + trust_region_loss
+        td_out = TensorDict({"loss_objective": loss}, batch_size=[])
+        td_out.set("tr_loss", trust_region_loss)
+
+        if self.entropy_bonus:
+            entropy = self.get_entropy_bonus(dist)
+            td_out.set("entropy", entropy.detach().mean())  # for logging
+            td_out.set("kl_approx", kl_approx.detach().mean())  # for logging
+            td_out.set("loss_entropy", -self.entropy_coef * entropy)
+        if self.critic_coef:
+            loss_critic, value_clip_fraction = self.loss_critic(tensordict)
+            td_out.set("loss_critic", loss_critic)
+            if value_clip_fraction is not None:
+                td_out.set("value_clip_fraction", value_clip_fraction)
+
+        td_out.set("ESS", _reduce(ess, self.reduction) / batch)
+        td_out = td_out.named_apply(
+            lambda name, value: _reduce(value, reduction=self.reduction).squeeze(-1)
+            if name.startswith("loss_")
+            else value,
+            batch_size=[],
+        )
+        return td_out

fancy_rl/objectives/vlearn.py

@@ -0,0 +1,106 @@
+import torch
+from torchrl.objectives import LossModule
+from torch.distributions import Normal
+
+class VLEARNLoss(LossModule):
+    def __init__(
+        self,
+        actor_network,
+        critic_network,
+        old_actor_network,
+        gamma=0.99,
+        lmbda=0.95,
+        entropy_coef=0.01,
+        critic_coef=0.5,
+        normalize_advantage=True,
+        eps=1e-8,
+        delta=0.1
+    ):
+        super().__init__()
+        self.actor_network = actor_network
+        self.critic_network = critic_network
+        self.old_actor_network = old_actor_network
+        self.gamma = gamma
+        self.lmbda = lmbda
+        self.entropy_coef = entropy_coef
+        self.critic_coef = critic_coef
+        self.normalize_advantage = normalize_advantage
+        self.eps = eps
+        self.delta = delta
+
+    def forward(self, tensordict):
+        # Compute returns and advantages
+        with torch.no_grad():
+            returns = self.compute_returns(tensordict)
+            values = self.critic_network(tensordict)["state_value"]
+            advantages = returns - values
+            if self.normalize_advantage:
+                advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+
+        # Compute actor loss
+        new_td = self.actor_network(tensordict)
+        old_td = self.old_actor_network(tensordict)
+
+        new_dist = Normal(new_td["loc"], new_td["scale"])
+        old_dist = Normal(old_td["loc"], old_td["scale"])
+
+        new_log_prob = new_dist.log_prob(tensordict["action"]).sum(-1)
+        old_log_prob = old_dist.log_prob(tensordict["action"]).sum(-1)
+
+        ratio = torch.exp(new_log_prob - old_log_prob)
+        
+        # Compute projection
+        kl = torch.distributions.kl.kl_divergence(new_dist, old_dist).sum(-1)
+        alpha = torch.where(kl > self.delta, 
+                            torch.sqrt(self.delta / (kl + self.eps)), 
+                            torch.ones_like(kl))
+        proj_loc = alpha.unsqueeze(-1) * new_td["loc"] + (1 - alpha.unsqueeze(-1)) * old_td["loc"]
+        proj_scale = torch.sqrt(alpha.unsqueeze(-1)**2 * new_td["scale"]**2 + (1 - alpha.unsqueeze(-1))**2 * old_td["scale"]**2)
+        proj_dist = Normal(proj_loc, proj_scale)
+
+        proj_log_prob = proj_dist.log_prob(tensordict["action"]).sum(-1)
+        proj_ratio = torch.exp(proj_log_prob - old_log_prob)
+        
+        policy_loss = -torch.min(
+            ratio * advantages,
+            proj_ratio * advantages
+        ).mean()
+
+        # Compute critic loss
+        value_pred = self.critic_network(tensordict)["state_value"]
+        critic_loss = 0.5 * (returns - value_pred).pow(2).mean()
+
+        # Compute entropy loss
+        entropy_loss = -self.entropy_coef * new_dist.entropy().mean()
+
+        # Combine losses
+        loss = policy_loss + self.critic_coef * critic_loss + entropy_loss
+
+        return {
+            "loss": loss,
+            "policy_loss": policy_loss,
+            "critic_loss": critic_loss,
+            "entropy_loss": entropy_loss,
+        }
+
+    def compute_returns(self, tensordict):
+        rewards = tensordict["reward"]
+        dones = tensordict["done"]
+        values = self.critic_network(tensordict)["state_value"]
+        
+        returns = torch.zeros_like(rewards)
+        advantages = torch.zeros_like(rewards)
+        last_gae_lam = 0
+        
+        for t in reversed(range(len(rewards))):
+            if t == len(rewards) - 1:
+                next_value = 0
+            else:
+                next_value = values[t + 1]
+            
+            delta = rewards[t] + self.gamma * next_value * (1 - dones[t]) - values[t]
+            advantages[t] = last_gae_lam = delta + self.gamma * self.lmbda * (1 - dones[t]) * last_gae_lam
+        
+        returns = advantages + values
+        
+        return returns

fancy_rl/policy.py

@@ -1,59 +1,87 @@
 import torch.nn as nn
 from tensordict.nn import TensorDictModule
 from torchrl.modules import MLP
+from torchrl.data.tensor_specs import DiscreteTensorSpec
 from tensordict.nn.distributions import NormalParamExtractor
-from fancy_rl.utils import is_discrete_space, get_space_shape
+from tensordict import TensorDict
-
+from torch.distributions import Categorical, MultivariateNormal, Normal
-class SharedModule(TensorDictModule):
-    def __init__(self, obs_space, hidden_sizes, activation_fn, device):
-        if hidden_sizes:
-            shared_module = MLP(
-                in_features=get_space_shape(obs_space)[-1],
-                out_features=hidden_sizes[-1],
-                num_cells=hidden_sizes[:-1],
-                activation_class=getattr(nn, activation_fn),
-                device=device
-            )
-            out_features = hidden_sizes[-1]
-        else:
-            shared_module = nn.Identity()
-            out_features = get_space_shape(obs_space)[-1]
-        
-        super().__init__(
-            module=shared_module,
-            in_keys=["observation"],
-            out_keys=["shared"],
-        )
-        self.out_features = out_features
 
 class Actor(TensorDictModule):
-    def __init__(self, shared_module, act_space, hidden_sizes, activation_fn, device):
+    def __init__(self, obs_space, act_space, hidden_sizes, activation_fn, device, full_covariance=False):
-        act_space_shape = get_space_shape(act_space)
+        self.discrete = isinstance(act_space, DiscreteTensorSpec)
-        if is_discrete_space(act_space):
-            out_features = act_space_shape[-1]
-        else:
-            out_features = act_space_shape[-1] * 2
         
-        actor_module = nn.Sequential(
+        obs_space = obs_space["observation"]
-            MLP(
+        act_space_shape = act_space.shape[1:]
-                in_features=shared_module.out_features,
+        obs_space_shape = obs_space.shape[1:]
+        
+        if self.discrete and full_covariance:
+            raise ValueError("Full covariance is not applicable for discrete action spaces.")
+        
+        self.full_covariance = full_covariance
+
+        if self.discrete:
+            out_features = act_space_shape[0]
+            out_keys = ["logits"]
+        else:
+            if full_covariance:
+                out_features = act_space_shape[0] + (act_space_shape[0] * (act_space_shape[0] + 1)) // 2
+                out_keys = ["loc", "scale_tril"]
+            else:
+                out_features = act_space_shape[0] * 2
+                out_keys = ["loc", "scale"]
+
+        actor_module = MLP(
+            in_features=obs_space_shape[0],
             out_features=out_features,
             num_cells=hidden_sizes,
             activation_class=getattr(nn, activation_fn),
             device=device
-            ),
-            NormalParamExtractor() if not is_discrete_space(act_space) else nn.Identity(),
         ).to(device)
+
+        if not self.discrete:
+            if full_covariance:
+                param_extractor = FullCovarianceNormalParamExtractor(act_space_shape[0])
+            else:
+                param_extractor = NormalParamExtractor()
+            actor_module = nn.Sequential(actor_module, param_extractor)
+
         super().__init__(
             module=actor_module,
-            in_keys=["shared"],
+            in_keys=["observation"],
-            out_keys=["loc", "scale"] if not is_discrete_space(act_space) else ["action_logits"],
+            out_keys=out_keys
         )
 
+    def get_dist(self, tensordict):
+        # Forward the observation through the network
+        out = self.forward(tensordict)
+        if self.discrete:
+            return Categorical(logits=out["logits"])
+        else:
+            if self.full_covariance:
+                return MultivariateNormal(loc=out["loc"], scale_tril=out["scale_tril"])
+            else:
+                return Normal(loc=out["loc"], scale=out["scale"])
+
+class FullCovarianceNormalParamExtractor(nn.Module):
+    def __init__(self, action_dim):
+        super().__init__()
+        self.action_dim = action_dim
+
+    def forward(self, x):
+        loc = x[:, :self.action_dim]
+        scale_tril = torch.zeros(x.shape[0], self.action_dim, self.action_dim, device=x.device)
+        tril_indices = torch.tril_indices(row=self.action_dim, col=self.action_dim, offset=0)
+        scale_tril[:, tril_indices[0], tril_indices[1]] = x[:, self.action_dim:]
+        scale_tril.diagonal(dim1=-2, dim2=-1).exp_()
+        return TensorDict({"loc": loc, "scale_tril": scale_tril}, batch_size=x.shape[0])
+
 class Critic(TensorDictModule):
-    def __init__(self, shared_module, hidden_sizes, activation_fn, device):
+    def __init__(self, obs_space, hidden_sizes, activation_fn, device):
+        obs_space = obs_space["observation"]
+        obs_space_shape = obs_space.shape[1:]
+        
         critic_module = MLP(
-            in_features=shared_module.out_features,
+            in_features=obs_space_shape[0],
             out_features=1,
             num_cells=hidden_sizes,
             activation_class=getattr(nn, activation_fn),
@@ -61,6 +89,6 @@ class Critic(TensorDictModule):
         ).to(device)
         super().__init__(
             module=critic_module,
-            in_keys=["shared"],
+            in_keys=["observation"],
             out_keys=["state_value"],
         )

fancy_rl/projections/__init__.py

@@ -0,0 +1,20 @@
+from .base_projection import BaseProjection
+from .identity_projection import IdentityProjection
+from .kl_projection import KLProjection
+from .wasserstein_projection import WassersteinProjection
+from .frobenius_projection import FrobeniusProjection
+
+def get_projection(projection_name: str):
+    projections = {
+        "identity_projection": IdentityProjection,
+        "kl_projection": KLProjection,
+        "wasserstein_projection": WassersteinProjection,
+        "frobenius_projection": FrobeniusProjection,
+    }
+    
+    projection = projections.get(projection_name.lower())
+    if projection is None:
+        raise ValueError(f"Unknown projection: {projection_name}")
+    return projection
+
+__all__ = ["BaseProjection", "IdentityProjection", "KLProjection", "WassersteinProjection", "FrobeniusProjection", "get_projection"]

fancy_rl/projections/base_projection.py

@@ -0,0 +1,92 @@
+from abc import ABC, abstractmethod
+import torch
+from torch import nn
+from typing import Dict, List, Tuple
+
+class BaseProjection(nn.Module, ABC):
+    def __init__(self, in_keys: List[str], out_keys: List[str], trust_region_coeff: float = 1.0, mean_bound: float = 0.01, cov_bound: float = 0.01, contextual_std: bool = True):
+        super().__init__()
+        self._validate_in_keys(in_keys)
+        self._validate_out_keys(out_keys)
+        self.in_keys = in_keys
+        self.out_keys = out_keys
+        self.trust_region_coeff = trust_region_coeff
+        self.mean_bound = mean_bound
+        self.cov_bound = cov_bound
+        self.full_cov = "scale_tril" in in_keys
+        self.contextual_std = contextual_std
+
+    def _validate_in_keys(self, keys: List[str]):
+        valid_keys = {"loc", "scale", "scale_tril", "old_loc", "old_scale", "old_scale_tril"}
+        if not set(keys).issubset(valid_keys):
+            raise ValueError(f"Invalid in_keys: {keys}. Must be a subset of {valid_keys}")
+        if "loc" not in keys or "old_loc" not in keys:
+            raise ValueError("Both 'loc' and 'old_loc' must be included in in_keys")
+        if ("scale" in keys) != ("old_scale" in keys) or ("scale_tril" in keys) != ("old_scale_tril" in keys):
+            raise ValueError("in_keys must have matching 'scale'/'old_scale' or 'scale_tril'/'old_scale_tril'")
+
+    def _validate_out_keys(self, keys: List[str]):
+        valid_keys = {"loc", "scale", "scale_tril"}
+        if not set(keys).issubset(valid_keys):
+            raise ValueError(f"Invalid out_keys: {keys}. Must be a subset of {valid_keys}")
+        if "loc" not in keys:
+            raise ValueError("'loc' must be included in out_keys")
+        if "scale" not in keys and "scale_tril" not in keys:
+            raise ValueError("Either 'scale' or 'scale_tril' must be included in out_keys")
+
+    @abstractmethod
+    def project(self, policy_params: Dict[str, torch.Tensor], old_policy_params: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        pass
+
+    @abstractmethod
+    def get_trust_region_loss(self, policy_params: Dict[str, torch.Tensor], proj_policy_params: Dict[str, torch.Tensor]) -> torch.Tensor:
+        pass
+
+    def forward(self, tensordict):
+        policy_params = {}
+        old_policy_params = {}
+        
+        for key in self.in_keys:
+            if key not in tensordict:
+                raise KeyError(f"Key '{key}' not found in tensordict. Available keys: {tensordict.keys()}")
+            
+            if key.startswith("old_"):
+                old_policy_params[key[4:]] = tensordict[key]
+            else:
+                policy_params[key] = tensordict[key]
+
+        projected_params = self.project(policy_params, old_policy_params)
+        return projected_params
+
+    def _calc_covariance(self, params: Dict[str, torch.Tensor]) -> torch.Tensor:
+        if not self.full_cov:
+            return torch.diag_embed(params["scale"].pow(2))
+        else:
+            return torch.matmul(params["scale_tril"], params["scale_tril"].transpose(-1, -2))
+
+    def _calc_scale_or_scale_tril(self, cov: torch.Tensor) -> torch.Tensor:
+        if not self.full_cov:
+            return torch.sqrt(cov.diagonal(dim1=-2, dim2=-1))
+        else:
+            return torch.linalg.cholesky(cov)
+
+    def _mean_projection(self, mean: torch.Tensor, old_mean: torch.Tensor, mean_part: torch.Tensor) -> torch.Tensor:
+        """Project mean based on the Mahalanobis objective and trust region.
+        
+        Args:
+            mean: Current mean vectors
+            old_mean: Old mean vectors
+            mean_part: Mahalanobis/Euclidean distance between the two mean vectors
+            
+        Returns:
+            Projected mean that satisfies the trust region
+        """
+        mask = mean_part > self.mean_bound
+        omega = torch.ones_like(mean_part, device=mean.device)
+        omega = torch.where(mask, torch.sqrt(mean_part / self.mean_bound) - 1., omega)
+        omega = torch.maximum(-omega, omega)[..., None]
+        
+        # Use matrix operations instead of boolean indexing
+        m = (mean + omega * old_mean) / (1. + omega + 1e-16)
+        mask_matrix = mask[..., None].to(mean.dtype)
+        return mask_matrix * m + (1 - mask_matrix) * mean

fancy_rl/projections/frobenius_projection.py

@@ -0,0 +1,89 @@
+import torch
+from .base_projection import BaseProjection
+from tensordict.nn import TensorDictModule
+from typing import Dict, Tuple
+
+class FrobeniusProjection(BaseProjection):
+    def __init__(self, in_keys: list[str], out_keys: list[str], trust_region_coeff: float = 1.0, mean_bound: float = 0.01, cov_bound: float = 0.01, scale_prec: bool = False, contextual_std: bool = True):
+        super().__init__(in_keys=in_keys, out_keys=out_keys, trust_region_coeff=trust_region_coeff, mean_bound=mean_bound, cov_bound=cov_bound, contextual_std=contextual_std)
+        self.scale_prec = scale_prec
+
+    def project(self, policy_params: Dict[str, torch.Tensor], old_policy_params: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        mean = policy_params["loc"]
+        old_mean = old_policy_params["loc"]
+
+        # Convert to covariance representation
+        cov = self._calc_covariance(policy_params)
+        old_cov = self._calc_covariance(old_policy_params)
+
+        if not self.contextual_std:
+            cov = cov[:1]
+            old_cov = old_cov[:1]
+
+        mean_part, cov_part = self._gaussian_frobenius((mean, cov), (old_mean, old_cov))
+        proj_mean = self._mean_projection(mean, old_mean, mean_part)
+        proj_cov = self._cov_projection(cov, old_cov, cov_part)
+
+        scale_or_tril = self._calc_scale_or_scale_tril(proj_cov)
+        if not self.contextual_std:
+            scale_or_tril = scale_or_tril.expand(mean.shape[0], *scale_or_tril.shape[1:])
+
+        return {"loc": proj_mean, self.out_keys[1]: scale_or_tril}
+
+    def get_trust_region_loss(self, policy_params: Dict[str, torch.Tensor], proj_policy_params: Dict[str, torch.Tensor]) -> torch.Tensor:
+        mean = policy_params["loc"]
+        proj_mean = proj_policy_params["loc"]
+
+        cov = self._calc_covariance(policy_params)
+        proj_cov = self._calc_covariance(proj_policy_params)
+
+        mean_diff = torch.sum(torch.square(mean - proj_mean), dim=-1)
+        cov_diff = torch.sum(torch.square(cov - proj_cov), dim=(-2, -1))
+
+        return (mean_diff + cov_diff).mean() * self.trust_region_coeff
+
+    def _gaussian_frobenius(self, p: Tuple[torch.Tensor, torch.Tensor], q: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+        mean, cov = p
+        old_mean, old_cov = q
+
+        if self.scale_prec:
+            if self.full_cov:
+                # Use triangular solve instead of inverse for stability
+                diff = mean - old_mean
+                solved = torch.triangular_solve(diff.unsqueeze(-1), old_cov, upper=False)[0].squeeze(-1)
+                mean_part = torch.sum(torch.square(solved), dim=-1)
+            else:
+                # Diagonal case - direct division is stable
+                mean_part = torch.sum(torch.square((mean - old_mean) / old_cov), dim=-1)
+        else:
+            mean_part = torch.sum(torch.square(mean - old_mean), dim=-1)
+
+        # Covariance part
+        if self.full_cov:
+            diff = old_cov - cov
+            cov_part = torch.sum(torch.square(diff), dim=(-2, -1))
+        else:
+            diff = old_cov - cov
+            cov_part = torch.sum(torch.square(diff), dim=-1)
+
+        return mean_part, cov_part
+
+    def _cov_projection(self, cov: torch.Tensor, old_cov: torch.Tensor, cov_part: torch.Tensor) -> torch.Tensor:
+        batch_shape = cov.shape[:-2] if cov.ndim > 2 else cov.shape[:-1]
+        cov_mask = cov_part > self.cov_bound
+
+        eta = torch.ones(batch_shape, dtype=cov.dtype, device=cov.device)
+        eta = torch.where(cov_mask, torch.sqrt(cov_part / self.cov_bound) - 1., eta)
+        eta = torch.maximum(-eta, eta)
+
+        if self.full_cov:
+            new_cov = (cov + torch.einsum('...,...ij->...ij', eta, old_cov)) / \
+                     (1. + eta + 1e-16)[..., None, None]
+            mask_matrix = cov_mask[..., None, None].to(cov.dtype)
+            proj_cov = torch.where(mask_matrix, new_cov, cov)
+        else:
+            new_cov = (cov + eta[..., None] * old_cov) / (1. + eta + 1e-16)[..., None]
+            mask_matrix = cov_mask[..., None].to(cov.dtype)
+            proj_cov = torch.where(mask_matrix, new_cov, cov)
+
+        return proj_cov

fancy_rl/projections/identity_projection.py

@@ -0,0 +1,13 @@
+import torch
+from .base_projection import BaseProjection
+from typing import Dict
+
+class IdentityProjection(BaseProjection):
+    def __init__(self, in_keys: list[str], out_keys: list[str], trust_region_coeff: float = 1.0, mean_bound: float = 0.01, cov_bound: float = 0.01, contextual_std: bool = True):
+        super().__init__(in_keys=in_keys, out_keys=out_keys, trust_region_coeff=trust_region_coeff, mean_bound=mean_bound, cov_bound=cov_bound, contextual_std=contextual_std)
+
+    def project(self, policy_params: Dict[str, torch.Tensor], old_policy_params: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        return policy_params
+
+    def get_trust_region_loss(self, policy_params: Dict[str, torch.Tensor], proj_policy_params: Dict[str, torch.Tensor]) -> torch.Tensor:
+        return torch.tensor(0.0, device=next(iter(policy_params.values())).device)

fancy_rl/projections/kl_projection.py

@@ -0,0 +1,228 @@
+import torch
+try:
+    import cpp_projection
+    cpp_projection_available = True
+except ImportError:
+    cpp_projection_available = False
+import numpy as np
+from .base_projection import BaseProjection
+from tensordict.nn import TensorDictModule
+from typing import Dict, Tuple, Any
+
+MAX_EVAL = 1000
+
+def get_numpy(tensor):
+    return tensor.detach().cpu().numpy()
+
+class KLProjection(BaseProjection):
+    def __init__(self, in_keys: list[str], out_keys: list[str], trust_region_coeff: float = 1.0, mean_bound: float = 0.01, cov_bound: float = 0.01, contextual_std: bool = True):
+        super().__init__(in_keys=in_keys, out_keys=out_keys, trust_region_coeff=trust_region_coeff, mean_bound=mean_bound, cov_bound=cov_bound, contextual_std=contextual_std)
+
+    def project(self, policy_params: Dict[str, torch.Tensor], old_policy_params: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        self._validate_inputs(policy_params, old_policy_params)
+        
+        mean = policy_params["loc"]
+        old_mean = old_policy_params["loc"]
+        
+        if self.full_cov:
+            scale_or_tril = policy_params["scale_tril"]
+            old_scale_or_tril = old_policy_params["scale_tril"]
+        else:
+            scale_or_tril = policy_params["scale"]
+            old_scale_or_tril = old_policy_params["scale"]
+
+        mean_part, cov_part = self._gaussian_kl((mean, scale_or_tril), (old_mean, old_scale_or_tril))
+
+        if not self.contextual_std:
+            scale_or_tril = scale_or_tril[:1]
+            old_scale_or_tril = old_scale_or_tril[:1]
+            cov_part = cov_part[:1]
+
+        proj_mean = self._mean_projection(mean, old_mean, mean_part)
+        proj_scale_or_tril = self._cov_projection(scale_or_tril, old_scale_or_tril, cov_part)
+
+        if not self.contextual_std:
+            proj_scale_or_tril = proj_scale_or_tril.expand(mean.shape[0], *proj_scale_or_tril.shape[1:])
+
+        if self.full_cov:
+            return {"loc": proj_mean, "scale_tril": proj_scale_or_tril}
+        else:
+            return {"loc": proj_mean, "scale": proj_scale_or_tril}
+
+    def get_trust_region_loss(self, policy_params: Dict[str, torch.Tensor], proj_policy_params: Dict[str, torch.Tensor]) -> torch.Tensor:
+        mean = policy_params["loc"]
+        proj_mean = proj_policy_params["loc"]
+        
+        if self.full_cov:
+            scale_or_tril = policy_params["scale_tril"]
+            proj_scale_or_tril = proj_policy_params["scale_tril"]
+        else:
+            scale_or_tril = policy_params["scale"]
+            proj_scale_or_tril = proj_policy_params["scale"]
+
+        kl = sum(self._gaussian_kl((mean, scale_or_tril), (proj_mean, proj_scale_or_tril)))
+        return kl.mean() * self.trust_region_coeff
+
+    def _gaussian_kl(self, p: Tuple[torch.Tensor, torch.Tensor], q: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+        mean, scale_or_tril = p
+        mean_other, scale_or_tril_other = q
+        k = mean.shape[-1]
+
+        maha_part = 0.5 * self._maha(mean, mean_other, scale_or_tril_other)
+
+        det_term = self._log_determinant(scale_or_tril)
+        det_term_other = self._log_determinant(scale_or_tril_other)
+
+        if self.full_cov:
+            trace_part = self._batched_trace_square(
+                torch.triangular_solve(scale_or_tril, scale_or_tril_other, upper=False)[0]
+            )
+        else:
+            trace_part = torch.sum((scale_or_tril / scale_or_tril_other) ** 2, dim=-1)
+
+        cov_part = 0.5 * (trace_part - k + det_term_other - det_term)
+
+        return maha_part, cov_part
+
+    def _maha(self, x: torch.Tensor, y: torch.Tensor, scale_or_tril: torch.Tensor) -> torch.Tensor:
+        diff = x - y
+        if self.full_cov:
+            solved = torch.triangular_solve(diff.unsqueeze(-1), scale_or_tril, upper=False)[0]
+            return torch.sum(torch.square(solved.squeeze(-1)), dim=-1)
+        else:
+            return torch.sum(torch.square(diff / scale_or_tril), dim=-1)
+
+    def _log_determinant(self, scale_or_tril: torch.Tensor) -> torch.Tensor:
+        if self.full_cov:
+            return 2 * torch.sum(torch.log(torch.diagonal(scale_or_tril, dim1=-2, dim2=-1)), dim=-1)
+        else:
+            return 2 * torch.sum(torch.log(scale_or_tril), dim=-1)
+
+    def _batched_trace_square(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.sum(x ** 2, dim=(-2, -1))
+
+    def _cov_projection(self, scale_or_tril: torch.Tensor, old_scale_or_tril: torch.Tensor, cov_part: torch.Tensor) -> torch.Tensor:
+        if self.full_cov:
+            cov = torch.matmul(scale_or_tril, scale_or_tril.transpose(-1, -2))
+            old_cov = torch.matmul(old_scale_or_tril, old_scale_or_tril.transpose(-1, -2))
+        else:
+            cov = scale_or_tril ** 2
+            old_cov = old_scale_or_tril ** 2
+
+        mask = cov_part > self.cov_bound
+        proj_scale_or_tril = scale_or_tril  # Start with original scale
+        
+        if mask.any():
+            if self.full_cov:
+                proj_cov = project_full_covariance(cov, scale_or_tril, old_scale_or_tril, self.cov_bound)
+                is_invalid = torch.isnan(proj_cov.mean(dim=(-2, -1)))
+                proj_scale_or_tril = torch.where(is_invalid[..., None, None], old_scale_or_tril, scale_or_tril)
+                mask = mask & ~is_invalid
+                chol = torch.linalg.cholesky(proj_cov)
+                proj_scale_or_tril = torch.where(mask[..., None, None], chol, proj_scale_or_tril)
+            else:
+                proj_cov = project_diag_covariance(cov, old_cov, self.cov_bound)
+                is_invalid = (torch.isnan(proj_cov.mean(dim=-1)) | 
+                            torch.isinf(proj_cov.mean(dim=-1)) | 
+                            (proj_cov.min(dim=-1).values < 0))
+                proj_scale_or_tril = torch.where(is_invalid[..., None], old_scale_or_tril, scale_or_tril)
+                mask = mask & ~is_invalid
+                proj_scale_or_tril = torch.where(mask[..., None], torch.sqrt(proj_cov), scale_or_tril)
+
+        return proj_scale_or_tril
+
+    def _validate_inputs(self, policy_params, old_policy_params):
+        if self.full_cov:
+            required_keys = ["loc", "scale_tril"]
+        else:
+            required_keys = ["loc", "scale"]
+        
+        for key in required_keys:
+            if key not in policy_params or key not in old_policy_params:
+                raise KeyError(f"Missing required key '{key}' in policy parameters")
+
+
+class KLProjectionGradFunctionCovOnly(torch.autograd.Function):
+    projection_op = None
+
+    @staticmethod
+    def get_projection_op(batch_shape, dim, max_eval=MAX_EVAL):
+        if not KLProjectionGradFunctionCovOnly.projection_op:
+            KLProjectionGradFunctionCovOnly.projection_op = \
+                cpp_projection.BatchedCovOnlyProjection(batch_shape, dim, max_eval=max_eval)
+        return KLProjectionGradFunctionCovOnly.projection_op
+
+    @staticmethod
+    def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
+        cov, chol, old_chol, eps_cov = args
+
+        batch_shape = cov.shape[0]
+        dim = cov.shape[-1]
+
+        cov_np = get_numpy(cov)
+        chol_np = get_numpy(chol)
+        old_chol_np = get_numpy(old_chol)
+        eps = get_numpy(eps_cov) * np.ones(batch_shape)
+
+        p_op = KLProjectionGradFunctionCovOnly.get_projection_op(batch_shape, dim)
+        ctx.proj = p_op
+
+        proj_std = p_op.forward(eps, old_chol_np, chol_np, cov_np)
+
+        return cov.new(proj_std)
+
+    @staticmethod
+    def backward(ctx: Any, *grad_outputs: Any) -> Any:
+        projection_op = ctx.proj
+        d_cov, = grad_outputs
+
+        d_cov_np = get_numpy(d_cov)
+        d_cov_np = np.atleast_2d(d_cov_np)
+
+        df_stds = projection_op.backward(d_cov_np)
+        df_stds = np.atleast_2d(df_stds)
+
+        df_stds = d_cov.new(df_stds)
+
+        return df_stds, None, None, None
+
+
+class KLProjectionGradFunctionDiagCovOnly(torch.autograd.Function):
+    projection_op = None
+
+    @staticmethod
+    def get_projection_op(batch_shape, dim, max_eval=MAX_EVAL):
+        if not KLProjectionGradFunctionDiagCovOnly.projection_op:
+            KLProjectionGradFunctionDiagCovOnly.projection_op = \
+                cpp_projection.BatchedDiagCovOnlyProjection(batch_shape, dim, max_eval=max_eval)
+        return KLProjectionGradFunctionDiagCovOnly.projection_op
+
+    @staticmethod
+    def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
+        cov, old_cov, eps_cov = args
+
+        batch_shape = cov.shape[0]
+        dim = cov.shape[-1]
+
+        cov_np = get_numpy(cov)
+        old_cov_np = get_numpy(old_cov)
+        eps = get_numpy(eps_cov) * np.ones(batch_shape)
+
+        p_op = KLProjectionGradFunctionDiagCovOnly.get_projection_op(batch_shape, dim)
+        ctx.proj = p_op
+
+        proj_std = p_op.forward(eps, old_cov_np, cov_np)
+
+        return cov.new(proj_std)
+
+    @staticmethod
+    def backward(ctx: Any, *grad_outputs: Any) -> Any:
+        projection_op = ctx.proj
+        d_std, = grad_outputs
+
+        d_cov_np = get_numpy(d_std)
+        d_cov_np = np.atleast_2d(d_cov_np)
+        df_stds = projection_op.backward(d_cov_np)
+        df_stds = np.atleast_2d(df_stds)
+
+        return d_std.new(df_stds), None, None

fancy_rl/projections/wasserstein_projection.py

@@ -0,0 +1,161 @@
+import torch
+from .base_projection import BaseProjection
+from typing import Dict, Tuple
+
+def scale_tril_to_sqrt(scale_tril: torch.Tensor) -> torch.Tensor:
+    """
+    'Converts' scale_tril to scale_sqrt.
+    
+    For Wasserstein distance, we need the matrix square root, not the Cholesky decomposition.
+    But since both are lower triangular, we can treat the Cholesky decomposition as if it were the matrix square root.
+    """
+    return scale_tril
+
+def gaussian_wasserstein_commutative(p: Tuple[torch.Tensor, torch.Tensor],
+                                   q: Tuple[torch.Tensor, torch.Tensor], 
+                                   scale_prec: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
+    mean, scale_or_sqrt = p
+    mean_other, scale_or_sqrt_other = q
+
+    mean_part = torch.sum(torch.square(mean - mean_other), dim=-1)
+
+    if scale_or_sqrt.dim() == mean.dim():  # Diagonal case
+        if scale_prec:
+            # More stable implementation for precision scaling
+            scale_part = torch.sum(
+                scale_or_sqrt_other**2 + scale_or_sqrt**2 - 
+                2 * scale_or_sqrt_other * scale_or_sqrt, 
+                dim=-1
+            )
+        else:
+            # Standard W2 for diagonal case
+            scale_part = torch.sum(
+                scale_or_sqrt_other**2 + scale_or_sqrt**2 - 
+                2 * scale_or_sqrt_other * scale_or_sqrt, 
+                dim=-1
+            )
+    else:  # Full covariance case
+        # Note: scale_or_sqrt is treated as the matrix square root, not Cholesky decomposition
+        if scale_prec:
+            # More stable implementation using triangular solve
+            identity = torch.eye(mean.shape[-1], dtype=scale_or_sqrt.dtype, device=scale_or_sqrt.device)
+            sqrt_inv_other = torch.triangular_solve(identity, scale_or_sqrt_other, upper=False)[0]
+            c = torch.matmul(sqrt_inv_other, scale_or_sqrt)
+            scale_part = torch.sum(identity**2 + c**2 - 2 * c, dim=(-2, -1))
+        else:
+            # Standard W2 for full covariance
+            scale_part = torch.sum(
+                scale_or_sqrt_other**2 + scale_or_sqrt**2 - 
+                2 * torch.matmul(scale_or_sqrt_other, scale_or_sqrt.transpose(-1, -2)), 
+                dim=(-2, -1)
+            )
+
+    return mean_part, scale_part
+
+class WassersteinProjection(BaseProjection):
+    def __init__(self, in_keys: list[str], out_keys: list[str], trust_region_coeff: float = 1.0, 
+                 mean_bound: float = 0.01, cov_bound: float = 0.01, scale_prec: bool = False, 
+                 contextual_std: bool = True):
+        super().__init__(in_keys=in_keys, out_keys=out_keys, trust_region_coeff=trust_region_coeff, 
+                        mean_bound=mean_bound, cov_bound=cov_bound, contextual_std=contextual_std)
+        self.scale_prec = scale_prec
+
+    def project(self, policy_params: Dict[str, torch.Tensor], old_policy_params: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        mean = policy_params["loc"]
+        old_mean = old_policy_params["loc"]
+        
+        # scale_tril is already Cholesky of matrix sqrt
+        scale_sqrt = policy_params[self.in_keys[1]]
+        old_scale_sqrt = old_policy_params[self.in_keys[1]]
+
+        if not self.contextual_std:
+            scale_sqrt = scale_sqrt[:1]
+            old_scale_sqrt = old_scale_sqrt[:1]
+
+        mean_part, scale_part = self._gaussian_wasserstein(
+            (mean, scale_sqrt), 
+            (old_mean, old_scale_sqrt)
+        )
+
+        proj_mean = self._mean_projection(mean, old_mean, mean_part)
+        proj_scale_sqrt = self._scale_projection(scale_sqrt, old_scale_sqrt, scale_part)
+
+        if not self.contextual_std:
+            proj_scale_sqrt = proj_scale_sqrt.expand(mean.shape[0], *proj_scale_sqrt.shape[1:])
+
+        return {"loc": proj_mean, self.out_keys[1]: proj_scale_sqrt}
+
+    def _gaussian_wasserstein(self, p: Tuple[torch.Tensor, torch.Tensor], 
+                            q: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+        mean, scale_sqrt = p
+        mean_other, scale_sqrt_other = q
+
+        mean_part = torch.sum(torch.square(mean - mean_other), dim=-1)
+
+        if not self.full_cov:
+            # Diagonal case is simpler
+            scale_part = torch.sum(
+                scale_sqrt_other**2 + scale_sqrt**2 - 
+                2 * scale_sqrt_other * scale_sqrt, 
+                dim=-1
+            )
+        else:
+            # Full covariance case uses matrix operations
+            scale_part = torch.sum(
+                scale_sqrt_other**2 + scale_sqrt**2 - 
+                2 * torch.matmul(scale_sqrt_other, scale_sqrt.transpose(-1, -2)), 
+                dim=(-2, -1)
+            )
+
+        return mean_part, scale_part
+
+    def get_trust_region_loss(self, policy_params: Dict[str, torch.Tensor], proj_policy_params: Dict[str, torch.Tensor]) -> torch.Tensor:
+        mean = policy_params["loc"]
+        proj_mean = proj_policy_params["loc"]
+        scale_or_sqrt = scale_tril_to_sqrt(policy_params[self.in_keys[1]])
+        proj_scale_or_sqrt = scale_tril_to_sqrt(proj_policy_params[self.out_keys[1]])
+        
+        mean_part, scale_part = gaussian_wasserstein_commutative(
+            (mean, scale_or_sqrt),
+            (proj_mean, proj_scale_or_sqrt),
+            self.scale_prec
+        )
+        w2 = mean_part + scale_part
+        return w2.mean() * self.trust_region_coeff
+
+    def _scale_projection(self, scale_or_sqrt: torch.Tensor, old_scale_or_sqrt: torch.Tensor, scale_part: torch.Tensor) -> torch.Tensor:
+        """Project scale parameters using multiplicative update."""
+        if scale_or_sqrt.dim() == old_scale_or_sqrt.dim() == 2:  # Diagonal case
+            return self._diagonal_scale_projection(scale_or_sqrt, old_scale_or_sqrt, scale_part)
+        else:  # Full covariance case
+            return self._full_cov_scale_projection(scale_or_sqrt, old_scale_or_sqrt, scale_part)
+
+    def _diagonal_scale_projection(self, scale: torch.Tensor, old_scale: torch.Tensor, scale_part: torch.Tensor) -> torch.Tensor:
+        cov_mask = scale_part > self.cov_bound
+        
+        batch_shape = scale.shape[:-1]
+        eta = torch.ones(batch_shape, dtype=scale.dtype, device=scale.device)
+        eta = torch.where(cov_mask,
+                        torch.sqrt(scale_part / self.cov_bound) - 1.,
+                        eta)
+        eta = torch.maximum(-eta, eta)
+        
+        new_scale = (scale + eta[..., None] * old_scale) / \
+                   (1. + eta + 1e-16)[..., None]
+        mask_matrix = cov_mask[..., None].to(scale.dtype)
+        return torch.where(mask_matrix, new_scale, scale)
+
+    def _full_cov_scale_projection(self, scale_sqrt: torch.Tensor, old_scale_sqrt: torch.Tensor, scale_part: torch.Tensor) -> torch.Tensor:
+        cov_mask = scale_part > self.cov_bound
+        
+        batch_shape = scale_sqrt.shape[:-2]
+        eta = torch.ones(batch_shape, dtype=scale_sqrt.dtype, device=scale_sqrt.device)
+        eta = torch.where(cov_mask,
+                        torch.sqrt(scale_part / self.cov_bound) - 1.,
+                        eta)
+        eta = torch.maximum(-eta, eta)
+        
+        new_scale = (scale_sqrt + torch.einsum('...,...ij->...ij', eta, old_scale_sqrt)) / \
+                    (1. + eta + 1e-16)[..., None, None]
+        mask_matrix = cov_mask[..., None, None].to(scale_sqrt.dtype)
+        return torch.where(mask_matrix, new_scale, scale_sqrt)

fancy_rl/utils.py

@@ -1,53 +0,0 @@
-try:
-    import gym
-    from gym.spaces import Discrete as GymDiscrete, MultiDiscrete as GymMultiDiscrete, MultiBinary as GymMultiBinary, Box as GymBox
-except ImportError:
-    gym = None
-
-import gymnasium
-from gymnasium.spaces import Discrete as GymnasiumDiscrete, MultiDiscrete as GymnasiumMultiDiscrete, MultiBinary as GymnasiumMultiBinary, Box as GymnasiumBox
-from torchrl.data.tensor_specs import (
-    DiscreteTensorSpec, OneHotDiscreteTensorSpec, MultiDiscreteTensorSpec,
-    BinaryDiscreteTensorSpec, BoundedTensorSpec, UnboundedContinuousTensorSpec
-)
-
-def is_discrete_space(action_space):
-    discrete_types = (
-        GymDiscrete, GymMultiDiscrete, GymMultiBinary,
-        GymnasiumDiscrete, GymnasiumMultiDiscrete, GymnasiumMultiBinary,
-        DiscreteTensorSpec, OneHotDiscreteTensorSpec, MultiDiscreteTensorSpec, BinaryDiscreteTensorSpec
-    )
-    continuous_types = (
-        GymBox, GymnasiumBox, BoundedTensorSpec, UnboundedContinuousTensorSpec
-    )
-    
-    if isinstance(action_space, discrete_types):
-        return True
-    elif isinstance(action_space, continuous_types):
-        return False
-    else:
-        raise ValueError(f"Unsupported action space type: {type(action_space)}")
-
-def get_space_shape(action_space):
-    if gym is not None:
-        discrete_types = (GymDiscrete, GymMultiDiscrete, GymMultiBinary)
-        continuous_types = (GymBox,)
-    else:
-        discrete_types = ()
-        continuous_types = ()
-
-    discrete_types += (GymnasiumDiscrete, GymnasiumMultiDiscrete, GymnasiumMultiBinary,
-                       DiscreteTensorSpec, OneHotDiscreteTensorSpec, MultiDiscreteTensorSpec, BinaryDiscreteTensorSpec)
-    continuous_types += (GymnasiumBox, BoundedTensorSpec, UnboundedContinuousTensorSpec)
-
-    if isinstance(action_space, discrete_types):
-        if isinstance(action_space, (GymDiscrete, GymnasiumDiscrete, DiscreteTensorSpec, OneHotDiscreteTensorSpec)):
-            return (action_space.n,)
-        elif isinstance(action_space, (GymMultiDiscrete, GymnasiumMultiDiscrete, MultiDiscreteTensorSpec)):
-            return (sum(action_space.nvec),)
-        elif isinstance(action_space, (GymMultiBinary, GymnasiumMultiBinary, BinaryDiscreteTensorSpec)):
-            return (action_space.n,)
-    elif isinstance(action_space, continuous_types):
-        return action_space.shape
-    
-    raise ValueError(f"Unsupported action space type: {type(action_space)}")

pyproject.toml

@@ -1,13 +1,37 @@
-[build-system]
+   [build-system]
-requires = ["setuptools", "wheel"]
+   requires = ["setuptools>=61.0", "wheel"]
-build-backend = "setuptools.build_meta"
+   build-backend = "setuptools.build_meta"
 
-[project]
+   [project]
-name = "fancy_rl"
+   name = "fancy_rl"
-version = "0.1.0"
+   version = "0.1.0"
-dependencies = [
+   description = "Minimalistic and efficient implementations of PPO and TRPL for torchrl"
-    "gymnasium",
+   authors = [{name = "Dominik Roth", email = "mail@dominik-roth.eu"}]
-    "pyyaml",
+   readme = "README.md"
+   requires-python = ">=3.7"
+   classifiers = [
+       "Development Status :: 3 - Alpha",
+       "Intended Audience :: Developers",
+       "License :: OSI Approved :: MIT License",
+       "Programming Language :: Python :: 3",
+       "Programming Language :: Python :: 3.7",
+       "Programming Language :: Python :: 3.8",
+       "Programming Language :: Python :: 3.9",
+       "Programming Language :: Python :: 3.10",
+       "Programming Language :: Python :: 3.11",
+   ]
+   dependencies = [
+       "numpy",
        "torch",
-    "torchrl"
+       "gymnasium<1.0",
-]
+       "tensordict",
+       "torchrl",
+       "pytest",
+   ]
+
+   [project.urls]
+   Homepage = "https://git.dominik-roth.eu/dodox/fancy_rl"
+
+   [project.optional-dependencies]
+   dev = ["pytest"]
+   box2d = ["swig", "gymnasium[box2d]"]

test/test_ppo.py

@@ -1 +1,54 @@
-# TODO
+import pytest
+import numpy as np
+from fancy_rl import PPO
+import gymnasium as gym
+from torchrl.envs import GymEnv
+import torch as th
+from tensordict import TensorDict
+
+def simple_env():
+    return gym.make('LunarLander-v2')
+
+def test_ppo_instantiation():
+    ppo = PPO(simple_env)
+    assert isinstance(ppo, PPO)
+
+def test_ppo_instantiation_from_str():
+    ppo = PPO('CartPole-v1')
+    assert isinstance(ppo, PPO)
+
+def test_ppo_predict():
+    ppo = PPO(simple_env)
+    env = ppo.make_env()
+    obs = env.reset()
+    action = ppo.predict(obs)
+    assert isinstance(action, TensorDict)
+    
+    # Handle both single and composite action spaces
+    if isinstance(env.action_space, list):
+        expected_shape = (len(env.action_space),) + env.action_space[0].shape
+    else:
+        expected_shape = env.action_space.shape
+        
+    assert action["action"].shape == expected_shape
+
+def test_ppo_training():
+    ppo = PPO(simple_env, total_timesteps=100)
+    env = ppo.make_env()
+    
+    initial_performance = evaluate_policy(ppo, env)
+    ppo.train()
+    final_performance = evaluate_policy(ppo, env)
+
+def evaluate_policy(policy, env, n_eval_episodes=3):
+    total_reward = 0
+    for _ in range(n_eval_episodes):
+        tensordict = env.reset()
+        done = False
+        while not done:
+            action = policy.predict(tensordict)
+            next_tensordict = env.step(action).get("next")
+            total_reward += next_tensordict["reward"]
+            done = next_tensordict["done"]
+            tensordict = next_tensordict
+    return total_reward / n_eval_episodes

test/test_trpl.py

@@ -0,0 +1,75 @@
+import pytest
+import numpy as np
+from fancy_rl import TRPL
+import gymnasium as gym
+from tensordict import TensorDict
+
+def simple_env():
+    return gym.make('Pendulum-v1')
+
+def test_trpl_instantiation():
+    trpl = TRPL(simple_env)
+    assert isinstance(trpl, TRPL)
+
+def test_trpl_instantiation_from_str():
+    trpl = TRPL('MountainCarContinuous-v0')
+    assert isinstance(trpl, TRPL)
+
+@pytest.mark.parametrize("learning_rate", [1e-4, 3e-4, 1e-3])
+@pytest.mark.parametrize("n_steps", [1024, 2048])
+@pytest.mark.parametrize("batch_size", [32, 64, 128])
+@pytest.mark.parametrize("gamma", [0.95, 0.99])
+@pytest.mark.parametrize("trust_region_bound_mean", [0.05, 0.1])
+@pytest.mark.parametrize("trust_region_bound_cov", [0.0005, 0.001])
+def test_trpl_initialization_with_different_hps(learning_rate, n_steps, batch_size, gamma, trust_region_bound_mean, trust_region_bound_cov):
+    trpl = TRPL(
+        simple_env,
+        learning_rate=learning_rate,
+        n_steps=n_steps,
+        batch_size=batch_size,
+        gamma=gamma,
+        trust_region_bound_mean=trust_region_bound_mean,
+        trust_region_bound_cov=trust_region_bound_cov
+    )
+    assert trpl.learning_rate == learning_rate
+    assert trpl.n_steps == n_steps
+    assert trpl.batch_size == batch_size
+    assert trpl.gamma == gamma
+    assert trpl.projection.mean_bound == trust_region_bound_mean
+    assert trpl.projection.cov_bound == trust_region_bound_cov
+
+def test_trpl_predict():
+    trpl = TRPL(simple_env)
+    env = trpl.make_env()
+    obs = env.reset()
+    action = trpl.predict(obs)
+    assert isinstance(action, TensorDict)
+    
+    # Handle both single and composite action spaces
+    if isinstance(env.action_space, list):
+        expected_shape = (len(env.action_space),) + env.action_space[0].shape
+    else:
+        expected_shape = env.action_space.shape
+        
+    assert action["action"].shape == expected_shape
+
+def test_trpl_training():
+    trpl = TRPL(simple_env, total_timesteps=100)
+    env = trpl.make_env()
+    
+    initial_performance = evaluate_policy(trpl, env)
+    trpl.train()
+    final_performance = evaluate_policy(trpl, env)
+
+def evaluate_policy(policy, env, n_eval_episodes=3):
+    total_reward = 0
+    for _ in range(n_eval_episodes):
+        tensordict = env.reset()
+        done = False
+        while not done:
+            action = policy.predict(tensordict)
+            next_tensordict = env.step(action).get("next")
+            total_reward += next_tensordict["reward"]
+            done = next_tensordict["done"]
+            tensordict = next_tensordict
+    return total_reward / n_eval_episodes

test/test_vlearn.py.disabled

@@ -0,0 +1,81 @@
+import pytest
+import torch
+import numpy as np
+from fancy_rl import VLEARN
+import gymnasium as gym
+
+@pytest.fixture
+def simple_env():
+    return gym.make('CartPole-v1')
+
+def test_vlearn_instantiation():
+    vlearn = VLEARN("CartPole-v1")
+    assert isinstance(vlearn, VLEARN)
+
+@pytest.mark.parametrize("learning_rate", [1e-4, 3e-4, 1e-3])
+@pytest.mark.parametrize("n_steps", [1024, 2048])
+@pytest.mark.parametrize("batch_size", [32, 64, 128])
+@pytest.mark.parametrize("gamma", [0.95, 0.99])
+@pytest.mark.parametrize("mean_bound", [0.05, 0.1])
+@pytest.mark.parametrize("cov_bound", [0.0005, 0.001])
+def test_vlearn_initialization_with_different_hps(learning_rate, n_steps, batch_size, gamma, mean_bound, cov_bound):
+    vlearn = VLEARN(
+        "CartPole-v1",
+        learning_rate=learning_rate,
+        n_steps=n_steps,
+        batch_size=batch_size,
+        gamma=gamma,
+        mean_bound=mean_bound,
+        cov_bound=cov_bound
+    )
+    assert vlearn.learning_rate == learning_rate
+    assert vlearn.n_steps == n_steps
+    assert vlearn.batch_size == batch_size
+    assert vlearn.gamma == gamma
+    assert vlearn.mean_bound == mean_bound
+    assert vlearn.cov_bound == cov_bound
+
+def test_vlearn_predict(simple_env):
+    vlearn = VLEARN("CartPole-v1")
+    obs, _ = simple_env.reset()
+    action, _ = vlearn.predict(obs)
+    assert isinstance(action, np.ndarray)
+    assert action.shape == simple_env.action_space.shape
+
+def test_vlearn_learn():
+    vlearn = VLEARN("CartPole-v1", n_steps=64, batch_size=32)
+    env = gym.make("CartPole-v1")
+    obs, _ = env.reset()
+    for _ in range(64):
+        action, _ = vlearn.predict(obs)
+        next_obs, reward, done, truncated, _ = env.step(action)
+        vlearn.store_transition(obs, action, reward, done, next_obs)
+        obs = next_obs
+        if done or truncated:
+            obs, _ = env.reset()
+    
+    loss = vlearn.learn()
+    assert isinstance(loss, dict)
+    assert "policy_loss" in loss
+    assert "value_loss" in loss
+
+def test_vlearn_training(simple_env):
+    vlearn = VLEARN("CartPole-v1", total_timesteps=10000)
+    
+    initial_performance = evaluate_policy(vlearn, simple_env)
+    vlearn.train()
+    final_performance = evaluate_policy(vlearn, simple_env)
+    
+    assert final_performance > initial_performance, "VLearn should improve performance after training"
+
+def evaluate_policy(policy, env, n_eval_episodes=10):
+    total_reward = 0
+    for _ in range(n_eval_episodes):
+        obs, _ = env.reset()
+        done = False
+        while not done:
+            action, _ = policy.predict(obs)
+            obs, reward, terminated, truncated, _ = env.step(action)
+            total_reward += reward
+            done = terminated or truncated
+    return total_reward / n_eval_episodes