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,8 +8,8 @@
 
 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. |
+| :exclamation: This project is still WIP and not ready to be used. (Problems with torchdict routing are a pain to debug...) |
-| ------------------------------------------------------------ |
+| -------------------------------------------------------------------------------------------------------------------------- |
 
 ## Installation
 
@@ -49,19 +49,37 @@ 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/
 ```
 
-## TODO
+## 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
-- [ ] Test / Benchmark TRPL
+- [ ] All TRPL Tests green
 - [ ] Make contextual covariance optional
 - [ ] Allow full-cov via chol
+- [ ] Test / Benchmark TRPL
 - [ ] Write docs / extend README
-- [ ] (Implement SAC?)
+- [ ] Test func of non-gym envs
+- [ ] Implement SAC
+- [ ] Implement VLEARN
 
 ## Contributing
 

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,55 +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 tensordict import LazyStackedTensorDict, TensorDict
-from abc import ABC
 
 from fancy_rl.loggers import TerminalLogger
+from fancy_rl.algos.algo import Algo
 
-class OnPolicy(ABC):
+class OnPolicy(Algo):
     def __init__(
         self,
         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,
-        total_timesteps,
+        gamma=0.99,
-        eval_interval,
+        total_timesteps=1e6,
-        eval_deterministic,
+        eval_interval=2048,
-        entropy_coef,
+        eval_deterministic=True,
-        critic_coef,
+        entropy_coef=0.01,
-        normalize_advantage,
+        critic_coef=0.5,
-        device=None,
+        normalize_advantage=True,
-        eval_episodes=10,
         env_spec_eval=None,
+        eval_episodes=10,
+        device=None,
     ):
-        self.env_spec = env_spec
+        device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        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)]
+        super().__init__(
-        self.optimizers = optimizers
+            env_spec=env_spec,
-        self.learning_rate = learning_rate
+            loggers=loggers,
-        self.n_steps = n_steps
+            optimizers=optimizers,
-        self.batch_size = batch_size
+            learning_rate=learning_rate,
-        self.n_epochs = n_epochs
+            n_steps=n_steps,
-        self.gamma = gamma
+            batch_size=batch_size,
-        self.total_timesteps = total_timesteps
+            n_epochs=n_epochs,
-        self.eval_interval = eval_interval
+            gamma=gamma,
-        self.eval_deterministic = eval_deterministic
+            total_timesteps=total_timesteps,
-        self.entropy_coef = entropy_coef
+            eval_interval=eval_interval,
-        self.critic_coef = critic_coef
+            eval_deterministic=eval_deterministic,
-        self.normalize_advantage = normalize_advantage
+            entropy_coef=entropy_coef,
-        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
+            critic_coef=critic_coef,
-        self.eval_episodes = eval_episodes
+            normalize_advantage=normalize_advantage,
+            device=device,
+            env_spec_eval=env_spec_eval,
+            eval_episodes=eval_episodes,
+        )
 
         # Create collector
         self.collector = SyncDataCollector(
             create_env_fn=lambda: self.make_env(eval=False),
-            policy=self.actor,
+            policy=self.prob_actor,
             frames_per_batch=self.n_steps,
             total_frames=self.total_timesteps,
             device=self.device,
@@ -69,30 +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).to(self.device)
-        elif callable(env_spec):
-            env = env_spec()
-            if isinstance(env, gym.Env):
-                env = GymWrapper(env).to(self.device)
-        elif isinstance(env, gym.Env):
-            env = GymWrapper(env).to(self.device)
-        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()
         losses = self.loss_module(batch)
-        loss = losses['loss_objective'] + losses["loss_entropy"] + losses["loss_critic"]
+        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):
@@ -140,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

@@ -2,6 +2,7 @@ import torch
 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
 
@@ -30,24 +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
         
-        self.critic = Critic(obs_space, critic_hidden_sizes, critic_activation_fn, device)
+        # Get spaces from specs for parallel env
-        actor_net = Actor(obs_space, act_space, actor_hidden_sizes, actor_activation_fn, device)
+        self.obs_space = env.observation_spec
-        self.actor = ProbabilisticActor(
+        self.act_space = env.action_spec
-            module=actor_net,
+        
-            in_keys=["loc", "scale"],
+        self.discrete = isinstance(self.act_space, DiscreteTensorSpec)
+
+        self.critic = Critic(self.obs_space, critic_hidden_sizes, critic_activation_fn, device)
+        self.actor = Actor(self.obs_space, self.act_space, actor_hidden_sizes, actor_activation_fn, device, full_covariance=full_covariance)
+
+        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 = {
@@ -85,7 +111,7 @@ class PPO(OnPolicy):
         self.loss_module = ClipPPOLoss(
             actor_network=self.actor,
             critic_network=self.critic,
-            clip_epsilon=clip_range,
+            clip_epsilon=self.clip_range,
             loss_critic_type='l2',
             entropy_coef=self.entropy_coef,
             critic_coef=self.critic_coef,

fancy_rl/algos/trpl.py

@@ -1,9 +1,78 @@
 import torch
-from torchrl.modules import ProbabilisticActor
+from torch import nn
-from torchrl.objectives.value.advantages import GAE
+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__(
@@ -14,49 +83,87 @@ class TRPL(OnPolicy):
         critic_hidden_sizes=[64, 64],
         actor_activation_fn="Tanh",
         critic_activation_fn="Tanh",
-        proj_layer_type=None,
         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,
-        trust_region_coef=10.0,
         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
 
-        self.trust_region_layer = None # TODO: from proj_layer_type
-        self.trust_region_coef = trust_region_coef
-
         # 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
 
-        self.critic = Critic(obs_space, critic_hidden_sizes, critic_activation_fn, device)
+        # Get spaces from specs for parallel env
-        actor_net = Actor(obs_space, act_space, actor_hidden_sizes, actor_activation_fn, device)
+        self.obs_space = env.observation_spec
-        raw_actor = ProbabilisticActor(
+        self.act_space = env.action_spec
-            module=actor_net,
+
-            in_keys=["loc", "scale"],
+        assert not isinstance(self.act_space, DiscreteTensorSpec), "TRPL does not support discrete action spaces"
-            out_keys=["action"],
+
-            distribution_class=torch.distributions.Normal,
+        self.critic = Critic(self.obs_space, critic_hidden_sizes, critic_activation_fn, device)
-            return_log_prob=True
+        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,
         )
-        self.actor = raw_actor # TODO: Proj here
 
         optimizers = {
-            "actor": torch.optim.Adam(self.actor.parameters(), lr=learning_rate),
+            "actor": torch.optim.Adam(self.raw_actor.parameters(), lr=learning_rate),
             "critic": torch.optim.Adam(self.critic.parameters(), lr=learning_rate)
         }
 
@@ -79,7 +186,6 @@ class TRPL(OnPolicy):
             env_spec_eval=env_spec_eval,
             eval_episodes=eval_episodes,
         )
-
         self.adv_module = GAE(
             gamma=self.gamma,
             lmbda=gae_lambda,
@@ -87,13 +193,8 @@ class TRPL(OnPolicy):
             average_gae=False,
         )
 
-        self.loss_module = TRPLLoss(
+    def update_old_policy(self):
-            actor_network=self.actor,
+        self.old_actor.load_state_dict(self.raw_actor.state_dict())
-            critic_network=self.critic,
+
-            trust_region_layer=self.trust_region_layer,
+    def post_update(self):
-            loss_critic_type='l2',
+        self.update_old_policy()
-            entropy_coef=self.entropy_coef,
-            critic_coef=self.critic_coef,
-            trust_region_coef=self.trust_region_coef,
-            normalize_advantage=self.normalize_advantage,
-        )

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/objectives/trpl.py

@@ -38,100 +38,40 @@ from torchrl.objectives.value import (
 )
 
 from torchrl.objectives.ppo import PPOLoss
+from fancy_rl.projections import get_projection
 
 class TRPLLoss(PPOLoss):
-    @dataclass
-    class _AcceptedKeys:
-        """Maintains default values for all configurable tensordict keys.
-
-        This class defines which tensordict keys can be set using '.set_keys(key_name=key_value)' and their
-        default values
-
-        Attributes:
-            advantage (NestedKey): The input tensordict key where the advantage is expected.
-                Will be used for the underlying value estimator. Defaults to ``"advantage"``.
-            value_target (NestedKey): The input tensordict key where the target state value is expected.
-                Will be used for the underlying value estimator Defaults to ``"value_target"``.
-            value (NestedKey): The input tensordict key where the state value is expected.
-                Will be used for the underlying value estimator. Defaults to ``"state_value"``.
-            sample_log_prob (NestedKey): The input tensordict key where the
-               sample log probability is expected.  Defaults to ``"sample_log_prob"``.
-            action (NestedKey): The input tensordict key where the action is expected.
-                Defaults to ``"action"``.
-            reward (NestedKey): The input tensordict key where the reward is expected.
-                Will be used for the underlying value estimator. Defaults to ``"reward"``.
-            done (NestedKey): The key in the input TensorDict that indicates
-                whether a trajectory is done. Will be used for the underlying value estimator.
-                Defaults to ``"done"``.
-            terminated (NestedKey): The key in the input TensorDict that indicates
-                whether a trajectory is terminated. Will be used for the underlying value estimator.
-                Defaults to ``"terminated"``.
-        """
-
-        advantage: NestedKey = "advantage"
-        value_target: NestedKey = "value_target"
-        value: NestedKey = "state_value"
-        sample_log_prob: NestedKey = "sample_log_prob"
-        action: NestedKey = "action"
-        reward: NestedKey = "reward"
-        done: NestedKey = "done"
-        terminated: NestedKey = "terminated"
-
-    default_keys = _AcceptedKeys()
-    default_value_estimator = ValueEstimators.GAE
-
-
-    actor_network: TensorDictModule
-    critic_network: TensorDictModule
-    actor_network_params: TensorDictParams
-    critic_network_params: TensorDictParams
-    target_actor_network_params: TensorDictParams
-    target_critic_network_params: TensorDictParams
-
     def __init__(
         self,
-        actor_network: ProbabilisticTensorDictSequential | None = None,
+        actor_network: ProbabilisticTensorDictSequential,
-        critic_network: TensorDictModule | None = None,
+        old_actor_network: ProbabilisticTensorDictSequential,
-        trust_region_layer: any | None = None,
+        critic_network: TensorDictModule,
-        entropy_bonus: bool = True,
+        projection: any,
-        samples_mc_entropy: int = 1,
         entropy_coef: float = 0.01,
         critic_coef: float = 1.0,
         trust_region_coef: float = 10.0,
-        loss_critic_type: str = "smooth_l1",
         normalize_advantage: bool = False,
-        gamma: float = None,
-        separate_losses: bool = False,
-        reduction: str = None,
-        clip_value: bool | float | None = None,
         **kwargs,
     ):
-        self.trust_region_layer = trust_region_layer
+        super().__init__(
-        self.trust_region_coef = trust_region_coef
+            actor_network=actor_network,
-
+            critic_network=critic_network,
-        super(TRPLLoss, self).__init__(
-            actor_network,
-            critic_network,
-            entropy_bonus=entropy_bonus,
-            samples_mc_entropy=samples_mc_entropy,
             entropy_coef=entropy_coef,
             critic_coef=critic_coef,
-            loss_critic_type=loss_critic_type,
             normalize_advantage=normalize_advantage,
-            gamma=gamma,
-            separate_losses=separate_losses,
-            reduction=reduction,
-            clip_value=clip_value,
             **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"]
+            keys = ["loss_objective", "tr_loss"]
             if self.entropy_bonus:
                 keys.extend(["entropy", "loss_entropy"])
-            if self.loss_critic:
+            if self.critic_coef:
                 keys.append("loss_critic")
             keys.append("ESS")
             self._out_keys = keys
@@ -141,8 +81,12 @@ class TRPLLoss(PPOLoss):
     def out_keys(self, values):
         self._out_keys = values
 
-    @dispatch
+    def _trust_region_loss(self, tensordict):
-    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        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:
@@ -159,11 +103,8 @@ class TRPLLoss(PPOLoss):
 
         log_weight, dist, kl_approx = self._log_weight(tensordict)
         trust_region_loss_unscaled = self._trust_region_loss(tensordict)
-        # ESS for logging
+        
         with torch.no_grad():
-            # In theory, ESS should be computed on particles sampled from the same source. Here we sample according
-            # to different, unrelated trajectories, which is not standard. Still it can give a idea of the dispersion
-            # of the weights.
             lw = log_weight.squeeze()
             ess = (2 * lw.logsumexp(0) - (2 * lw).logsumexp(0)).exp()
             batch = log_weight.shape[0]
@@ -194,8 +135,3 @@ class TRPLLoss(PPOLoss):
             batch_size=[],
         )
         return td_out
-
-    def _trust_region_loss(self, tensordict):
-        old_distribution = 
-        raw_distribution = 
-        return self.policy_projection.get_trust_region_loss(raw_distribution, old_distribution)

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,37 +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 Actor(TensorDictModule):
-    def __init__(self, obs_space, 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=get_space_shape(obs_space)[-1],
+        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=["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, obs_space, hidden_sizes, activation_fn, device):
+        obs_space = obs_space["observation"]
+        obs_space_shape = obs_space.shape[1:]
+        
         critic_module = MLP(
-            in_features=get_space_shape(obs_space)[-1],
+            in_features=obs_space_shape[0],
             out_features=1,
             num_cells=hidden_sizes,
             activation_class=getattr(nn, activation_fn),

fancy_rl/projections/__init__.py

@@ -1,6 +1,20 @@
-try:
+from .base_projection import BaseProjection
-    import cpp_projection
+from .identity_projection import IdentityProjection
-except ModuleNotFoundError:
+from .kl_projection import KLProjection
-    from .base_projection_layer import ITPALExceptionLayer as KLProjectionLayer
+from .wasserstein_projection import WassersteinProjection
-else:
+from .frobenius_projection import FrobeniusProjection
-    from .kl_projection_layer import KLProjectionLayer
+
+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/base_projection_layer.py

@@ -1,191 +0,0 @@
-from typing import Any, Dict, Optional, Type, Union, Tuple, final
-
-import torch as th
-
-from fancy_rl.norm import *
-
-class BaseProjectionLayer(object):
-    def __init__(self,
-                 mean_bound: float = 0.03,
-                 cov_bound: float = 1e-3,
-                 trust_region_coeff: float = 1.0,
-                 scale_prec: bool = False,
-                 ):
-        self.mean_bound = mean_bound
-        self.cov_bound = cov_bound
-        self.trust_region_coeff = trust_region_coeff
-        self.scale_prec = scale_prec
-        self.mean_eq = False
-
-    def __call__(self, p, q, **kwargs):
-        return self._projection(p, q, eps=self.mean_bound, eps_cov=self.cov_bound, beta=None, **kwargs)
-
-    @final
-    def _projection(self, p, q, eps: th.Tensor, eps_cov: th.Tensor, beta: th.Tensor, **kwargs):
-        return self._trust_region_projection(
-            p, q, eps, eps_cov, **kwargs)
-
-    def _trust_region_projection(self, p, q, eps: th.Tensor, eps_cov: th.Tensor, **kwargs):
-        """
-        Hook for implementing the specific trust region projection
-        Args:
-            p: current distribution
-            q: old distribution
-            eps: mean trust region bound
-            eps_cov: covariance trust region bound
-            **kwargs:
-
-        Returns:
-            projected
-        """
-        return p
-
-    def get_trust_region_loss(self, p, proj_p):
-        # p:
-        #   predicted distribution from network output
-        # proj_p:
-        #   projected distribution
-
-        proj_mean, proj_chol = get_mean_and_chol(proj_p)
-        p_target = new_dist_like(p, proj_mean, proj_chol)
-        kl_diff = self.trust_region_value(p, p_target)
-
-        kl_loss = kl_diff.mean()
-
-        return kl_loss * self.trust_region_coeff
-
-    def trust_region_value(self, p, q):
-        """
-        Computes the KL divergence between two Gaussian distributions p and q_values.
-        Returns:
-            full kl divergence
-        """
-        return kl_divergence(p, q)
-
-    def new_dist_like(self, orig_p, mean, cov_cholesky):
-        assert isinstance(orig_p, Distribution)
-        p = orig_p.distribution
-        if isinstance(p, th.distributions.Normal):
-            p_out = orig_p.__class__(orig_p.action_dim)
-            p_out.distribution = th.distributions.Normal(mean, cov_cholesky)
-        elif isinstance(p, th.distributions.Independent):
-            p_out = orig_p.__class__(orig_p.action_dim)
-            p_out.distribution = th.distributions.Independent(
-                th.distributions.Normal(mean, cov_cholesky), 1)
-        elif isinstance(p, th.distributions.MultivariateNormal):
-            p_out = orig_p.__class__(orig_p.action_dim)
-            p_out.distribution = th.distributions.MultivariateNormal(
-                mean, scale_tril=cov_cholesky)
-        else:
-            raise Exception('Dist-Type not implemented (of sb3 dist)')
-        return p_out
-
-def entropy_inequality_projection(p: th.distributions.Normal,
-                                  beta: Union[float, th.Tensor]):
-    """
-    Projects std to satisfy an entropy INEQUALITY constraint.
-    Args:
-        p: current distribution
-        beta: target entropy for EACH std or general bound for all stds
-
-    Returns:
-        projected std that satisfies the entropy bound
-    """
-    mean, std = p.mean, p.stddev
-    k = std.shape[-1]
-    batch_shape = std.shape[:-2]
-
-    ent = p.entropy()
-    mask = ent < beta
-
-    # if nothing has to be projected skip computation
-    if (~mask).all():
-        return p
-
-    alpha = th.ones(batch_shape, dtype=std.dtype, device=std.device)
-    alpha[mask] = th.exp((beta[mask] - ent[mask]) / k)
-
-    proj_std = th.einsum('ijk,i->ijk', std, alpha)
-    new_mean, new_std = mean, th.where(mask[..., None, None], proj_std, std)
-    return th.distributions.Normal(new_mean, new_std)
-
-
-def entropy_equality_projection(p: th.distributions.Normal,
-                                beta: Union[float, th.Tensor]):
-    """
-    Projects std to satisfy an entropy EQUALITY constraint.
-    Args:
-        p: current distribution
-        beta: target entropy for EACH std or general bound for all stds
-
-    Returns:
-        projected std that satisfies the entropy bound
-    """
-    mean, std = p.mean, p.stddev
-    k = std.shape[-1]
-
-    ent = p.entropy()
-    alpha = th.exp((beta - ent) / k)
-    proj_std = th.einsum('ijk,i->ijk', std, alpha)
-    new_mean, new_std = mean, proj_std
-    return th.distributions.Normal(new_mean, new_std)
-
-
-def mean_projection(mean: th.Tensor, old_mean: th.Tensor, maha: th.Tensor, eps: th.Tensor):
-    """
-    Projects the mean based on the Mahalanobis objective and trust region.
-    Args:
-        mean: current mean vectors
-        old_mean: old mean vectors
-        maha: Mahalanobis distance between the two mean vectors
-        eps: trust region bound
-
-    Returns:
-        projected mean that satisfies the trust region
-    """
-    batch_shape = mean.shape[:-1]
-    mask = maha > eps
-
-    ################################################################################################################
-    # mean projection maha
-
-    # if nothing has to be projected skip computation
-    if mask.any():
-        omega = th.ones(batch_shape, dtype=mean.dtype, device=mean.device)
-        omega[mask] = th.sqrt(maha[mask] / eps) - 1.
-        omega = th.max(-omega, omega)[..., None]
-
-        m = (mean + omega * old_mean) / (1 + omega + 1e-16)
-        proj_mean = th.where(mask[..., None], m, mean)
-    else:
-        proj_mean = mean
-
-    return proj_mean
-
-
-def mean_equality_projection(mean: th.Tensor, old_mean: th.Tensor, maha: th.Tensor, eps: th.Tensor):
-    """
-    Projections the mean based on the Mahalanobis objective and trust region for an EQUALITY constraint.
-    Args:
-        mean: current mean vectors
-        old_mean: old mean vectors
-        maha: Mahalanobis distance between the two mean vectors
-        eps: trust region bound
-    Returns:
-        projected mean that satisfies the trust region
-    """
-
-    maha[maha == 0] += 1e-16
-    omega = th.sqrt(maha / eps) - 1.
-    omega = omega[..., None]
-
-    proj_mean = (mean + omega * old_mean) / (1 + omega + 1e-16)
-
-    return proj_mean
-
-
-class ITPALExceptionLayer(BaseProjectionLayer):
-    def __init__(self,
-                 *args, **kwargs
-                 ):
-        raise Exception('To be able to use KL projections, ITPAL must be installed: https://github.com/ALRhub/ITPAL.')

fancy_rl/projections/frob_projection_layer.py

@@ -1,133 +0,0 @@
-import torch as th
-from typing import Tuple
-
-from .base_projection_layer import BaseProjectionLayer, mean_projection
-
-from ..misc.norm import mahalanobis, frob_sq
-from ..misc.distTools import get_mean_and_chol, get_cov, new_dist_like, has_diag_cov
-
-
-class FrobeniusProjectionLayer(BaseProjectionLayer):
-
-    def _trust_region_projection(self, p, q, eps: th.Tensor, eps_cov: th.Tensor, **kwargs):
-        """
-        Stolen from Fabian's Code (Public Version)
-
-        Runs Frobenius projection layer and constructs cholesky of covariance
-
-        Args:
-            policy: policy instance
-            p: current distribution
-            q: old distribution
-            eps: (modified) kl bound/ kl bound for mean part
-            eps_cov: (modified) kl bound for cov part
-            beta: (modified) entropy bound
-            **kwargs:
-        Returns: mean, cov cholesky
-        """
-
-        mean, chol = get_mean_and_chol(p, expand=True)
-        old_mean, old_chol = get_mean_and_chol(q, expand=True)
-        batch_shape = mean.shape[:-1]
-
-        ####################################################################################################################
-        # precompute mean and cov part of frob projection, which are used for the projection.
-        mean_part, cov_part, cov, cov_old = gaussian_frobenius(
-            p, q, self.scale_prec, True)
-
-        ################################################################################################################
-        # mean projection maha/euclidean
-
-        proj_mean = mean_projection(mean, old_mean, mean_part, eps)
-
-        ################################################################################################################
-        # cov projection frobenius
-
-        cov_mask = cov_part > eps_cov
-
-        if cov_mask.any():
-            eta = th.ones(batch_shape, dtype=chol.dtype, device=chol.device)
-            eta[cov_mask] = th.sqrt(cov_part[cov_mask] / eps_cov) - 1.
-            eta = th.max(-eta, eta)
-
-            new_cov = (cov + th.einsum('i,ijk->ijk', eta, cov_old)
-                       ) / (1. + eta + 1e-16)[..., None, None]
-            proj_chol = th.where(
-                cov_mask[..., None, None], th.linalg.cholesky(new_cov), chol)
-        else:
-            proj_chol = chol
-
-        if has_diag_cov(p):
-            proj_chol = th.diagonal(proj_chol, dim1=-2, dim2=-1)
-
-        proj_p = new_dist_like(p, proj_mean, proj_chol)
-        return proj_p
-
-    def trust_region_value(self, p, q):
-        """
-        Stolen from Fabian's Code (Public Version)
-
-        Computes the Frobenius metric between two Gaussian distributions p and q.
-        Args:
-            policy: policy instance
-            p: current distribution
-            q: old distribution
-        Returns:
-            mean and covariance part of Frobenius metric
-        """
-        return gaussian_frobenius(p, q, self.scale_prec)
-
-    def get_trust_region_loss(self, p, proj_p):
-        """
-        Stolen from Fabian's Code (Public Version)
-        """
-
-        mean_diff, _ = self.trust_region_value(p, proj_p)
-        if False and policy.contextual_std:
-            # Compute MSE here, because we found the Frobenius norm tends to generate values that explode for the cov
-            p_mean, proj_p_mean = p.mean, proj_p.mean
-            cov_diff = (p_mean - proj_p_mean).pow(2).sum([-1, -2])
-            delta_loss = (mean_diff + cov_diff).mean()
-        else:
-            delta_loss = mean_diff.mean()
-
-        return delta_loss * self.trust_region_coeff
-
-
-def gaussian_frobenius(p, q, scale_prec: bool = False, return_cov: bool = False):
-    """
-    Stolen from Fabian' Code (Public Version)
-
-    Compute (p - q_values) (L_oL_o^T)^-1 (p - 1)^T + |LL^T - L_oL_o^T|_F^2 with p,q_values ~ N(y, LL^T)
-    Args:
-        policy: current policy
-        p: mean and chol of gaussian p
-        q: mean and chol of gaussian q_values
-        return_cov: return cov matrices for further computations
-        scale_prec: scale objective with precision matrix
-    Returns: mahalanobis distance, squared frobenius norm
-    """
-
-    mean, chol = get_mean_and_chol(p)
-    mean_other, chol_other = get_mean_and_chol(q)
-
-    if scale_prec:
-        # maha objective for mean
-        mean_part = mahalanobis(mean, mean_other, chol_other)
-    else:
-        # euclidean distance for mean
-        # mean_part = ch.norm(mean_other - mean, ord=2, axis=1) ** 2
-        mean_part = ((mean_other - mean) ** 2).sum(1)
-
-    # frob objective for cov
-    cov = get_cov(p)
-    cov_other = get_cov(q)
-    diff = cov_other - cov
-    # Matrix is real symmetric PSD, therefore |A @ A^H|^2_F = tr{A @ A^H} = tr{A @ A}
-    #cov_part = torch_batched_trace(diff @ diff)
-    cov_part = frob_sq(diff, is_spd=True)
-
-    if return_cov:
-        return mean_part, cov_part, cov, cov_other
-
-    return mean_part, cov_part

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/identity_projection_layer.py

@@ -1,5 +0,0 @@
-from .base_projection_layer import BaseProjectionLayer
-
-class IdentityProjectionLayer(BaseProjectionLayer):
-    def project_from_rollouts(self, dist, rollout_data, **kwargs):
-        return dist, dist

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/kl_projection_layer.py

@@ -1,256 +0,0 @@
-from ..misc.distTools import get_diag_cov_vec, get_mean_and_chol, get_cov, is_contextual, new_dist_like, has_diag_cov
-from .base_projection_layer import BaseProjectionLayer, mean_projection, mean_equality_projection
-
-import cpp_projection
-import numpy as np
-import torch as th
-from typing import Tuple, Any
-
-from ..misc.norm import mahalanobis
-
-MAX_EVAL = 1000
-
-
-class KLProjectionLayer(BaseProjectionLayer):
-    """
-    Stolen from Fabian's Code (Private Version)
-    """
-
-    def _trust_region_projection(self, p, q, eps: th.Tensor, eps_cov: th.Tensor, **kwargs):
-        """
-        Stolen from Fabian's Code (Private Version)
-
-        runs kl projection layer and constructs sqrt of covariance
-        Args:
-            **kwargs:
-            policy: policy instance
-            p: current distribution
-            q: old distribution
-            eps: (modified) kl bound/ kl bound for mean part
-            eps_cov: (modified) kl bound for cov part
-
-        Returns:
-            mean, cov sqrt
-        """
-        mean, chol = get_mean_and_chol(p, expand=True)
-        old_mean, old_chol = get_mean_and_chol(q, expand=True)
-
-        ################################################################################################################
-        # project mean with closed form
-        # orig code: mean_part, _ = gaussian_kl(policy, p, q)
-        # But the mean_part is just the mahalanobis dist:
-        mean_part = mahalanobis(mean, old_mean, old_chol)
-        if self.mean_eq:
-            proj_mean = mean_equality_projection(
-                mean, old_mean, mean_part, eps)
-        else:
-            proj_mean = mean_projection(mean, old_mean, mean_part, eps)
-
-        if has_diag_cov(p):
-            cov_diag = get_diag_cov_vec(p)
-            old_cov_diag = get_diag_cov_vec(q)
-            proj_cov = KLProjectionGradFunctionDiagCovOnly.apply(cov_diag,
-                                                                 old_cov_diag,
-                                                                 eps_cov)
-            proj_chol = proj_cov.sqrt()  # .diag_embed()
-        else:
-            cov = get_cov(p)
-            old_cov = get_cov(q)
-            proj_cov = KLProjectionGradFunctionCovOnly.apply(
-                cov, old_cov, chol, old_chol, eps_cov)
-            proj_chol = th.linalg.cholesky(proj_cov)
-        proj_p = new_dist_like(p, proj_mean, proj_chol)
-        return proj_p
-
-
-class KLProjectionGradFunctionCovOnly(th.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:
-        #std, old_std, eps_cov = args
-        cov, old_cov, chol, old_chol, eps_cov = args
-
-        batch_shape = chol.shape[0]
-        dim = chol.shape[-1]
-
-        cov_np = cov.cpu().detach().numpy()
-        old_cov_np = old_cov.cpu().detach().numpy()
-        chol_np = chol.cpu().detach().numpy()
-        old_chol_np = old_chol.cpu().detach().numpy()
-        # eps = eps_cov.cpu().detach().numpy().astype(old_std_np.dtype) * \
-        eps = eps_cov * \
-            np.ones(batch_shape, dtype=old_chol_np.dtype)
-
-        p_op = KLProjectionGradFunctionCovOnly.get_projection_op(
-            batch_shape, dim)
-        ctx.proj = p_op
-
-        proj_cov = p_op.forward(eps, old_chol_np, chol_np, cov_np)
-
-        return th.Tensor(proj_cov)
-
-    @staticmethod
-    def backward(ctx: Any, *grad_outputs: Any) -> Any:
-        projection_op = ctx.proj
-        d_std, = grad_outputs
-
-        d_std_np = d_std.cpu().detach().numpy()
-        d_std_np = np.atleast_2d(d_std_np)
-        df_stds = projection_op.backward(d_std_np)
-        df_stds = np.atleast_2d(df_stds)
-
-        return d_std.new(df_stds), None, None, None, None
-
-
-class KLProjectionGradFunctionDiagCovOnly(th.autograd.Function):
-    projection_op = None
-
-    @staticmethod
-    def get_projection_op(batch_shape, dim: int, max_eval: int = 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_std_np, eps_cov = args
-
-        batch_shape = cov.shape[0]
-        dim = cov.shape[-1]
-
-        std_np = cov.to('cpu').detach().numpy()
-        old_std_np = old_std_np.to('cpu').detach().numpy()
-        # eps = eps_cov.to('cpu').detach().numpy().astype(old_std_np.dtype) * np.ones(batch_shape, dtype=old_std_np.dtype)
-        eps = eps_cov * np.ones(batch_shape, dtype=old_std_np.dtype)
-
-        p_op = KLProjectionGradFunctionDiagCovOnly.get_projection_op(
-            batch_shape, dim)
-        ctx.proj = p_op
-
-        try:
-            proj_std = p_op.forward(eps, old_std_np, std_np)
-        except:
-            proj_std = std_np
-
-        return cov.new(proj_std)
-
-    @staticmethod
-    def backward(ctx: Any, *grad_outputs: Any) -> Any:
-        projection_op = ctx.proj
-        d_std, = grad_outputs
-
-        d_std_np = d_std.to('cpu').detach().numpy()
-        d_std_np = np.atleast_2d(d_std_np)
-        df_stds = projection_op.backward(d_std_np)
-        df_stds = np.atleast_2d(df_stds)
-
-        return d_std.new(df_stds), None, None
-
-
-class KLProjectionGradFunctionDiagSplit(th.autograd.Function):
-    projection_op = None
-
-    @staticmethod
-    def get_projection_op(batch_shape, dim: int, max_eval: int = MAX_EVAL):
-        if not KLProjectionGradFunctionDiagSplit.projection_op:
-            KLProjectionGradFunctionDiagSplit.projection_op = \
-                cpp_projection.BatchedSplitDiagMoreProjection(
-                    batch_shape, dim, max_eval=max_eval)
-        return KLProjectionGradFunctionDiagSplit.projection_op
-
-    @staticmethod
-    def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
-        mean, cov, old_mean, old_cov, eps_mu, eps_sigma = args
-
-        batch_shape, dim = mean.shape
-
-        mean_np = mean.detach().numpy()
-        cov_np = cov.detach().numpy()
-        old_mean = old_mean.detach().numpy()
-        old_cov = old_cov.detach().numpy()
-        eps_mu = eps_mu * np.ones(batch_shape)
-        eps_sigma = eps_sigma * np.ones(batch_shape)
-
-        # p_op = cpp_projection.BatchedSplitDiagMoreProjection(batch_shape, dim, max_eval=100)
-        p_op = KLProjectionGradFunctionDiagSplit.get_projection_op(
-            batch_shape, dim)
-
-        try:
-            proj_mean, proj_cov = p_op.forward(
-                eps_mu, eps_sigma, old_mean, old_cov, mean_np, cov_np)
-        except Exception:
-            # try a second time
-            proj_mean, proj_cov = p_op.forward(
-                eps_mu, eps_sigma, old_mean, old_cov, mean_np, cov_np)
-        ctx.proj = p_op
-
-        return mean.new(proj_mean), cov.new(proj_cov)
-
-    @staticmethod
-    def backward(ctx: Any, *grad_outputs: Any) -> Any:
-        p_op = ctx.proj
-        d_means, d_std = grad_outputs
-
-        d_std_np = d_std.detach().numpy()
-        d_std_np = np.atleast_2d(d_std_np)
-        d_mean_np = d_means.detach().numpy()
-        dtarget_means, dtarget_covs = p_op.backward(d_mean_np, d_std_np)
-        dtarget_covs = np.atleast_2d(dtarget_covs)
-
-        return d_means.new(dtarget_means), d_std.new(dtarget_covs), None, None, None, None
-
-
-class KLProjectionGradFunctionJoint(th.autograd.Function):
-    projection_op = None
-
-    @staticmethod
-    def get_projection_op(batch_shape, dim: int, max_eval: int = MAX_EVAL):
-        if not KLProjectionGradFunctionJoint.projection_op:
-            KLProjectionGradFunctionJoint.projection_op = \
-                cpp_projection.BatchedProjection(batch_shape, dim, eec=False, constrain_entropy=False,
-                                                 max_eval=max_eval)
-        return KLProjectionGradFunctionJoint.projection_op
-
-    @staticmethod
-    def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
-        mean, cov, old_mean, old_cov, eps, beta = args
-
-        batch_shape, dim = mean.shape
-
-        mean_np = mean.detach().numpy()
-        cov_np = cov.detach().numpy()
-        old_mean = old_mean.detach().numpy()
-        old_cov = old_cov.detach().numpy()
-        eps = eps * np.ones(batch_shape)
-        beta = beta.detach().numpy() * np.ones(batch_shape)
-
-        # projection_op = cpp_projection.BatchedProjection(batch_shape, dim, eec=False, constrain_entropy=False)
-        # ctx.proj = projection_op
-
-        p_op = KLProjectionGradFunctionJoint.get_projection_op(
-            batch_shape, dim)
-        ctx.proj = p_op
-
-        proj_mean, proj_cov = p_op.forward(
-            eps, beta, old_mean, old_cov, mean_np, cov_np)
-
-        return mean.new(proj_mean), cov.new(proj_cov)
-
-    @staticmethod
-    def backward(ctx: Any, *grad_outputs: Any) -> Any:
-        projection_op = ctx.proj
-        d_means, d_covs = grad_outputs
-        df_means, df_covs = projection_op.backward(
-            d_means.detach().numpy(), d_covs.detach().numpy())
-        return d_means.new(df_means), d_means.new(df_covs), None, None, None, None

fancy_rl/projections/w2_projection_layer.py

@@ -1,164 +0,0 @@
-import numpy as np
-import torch as th
-from typing import Tuple, Any
-
-from ..misc.norm import mahalanobis
-
-from .base_projection_layer import BaseProjectionLayer, mean_projection
-
-from ..misc.norm import mahalanobis, _batch_trace
-from ..misc.distTools import get_diag_cov_vec, get_mean_and_chol,  get_mean_and_sqrt, get_cov, has_diag_cov
-
-from stable_baselines3.common.distributions import Distribution
-
-
-class WassersteinProjectionLayer(BaseProjectionLayer):
-    """
-    Stolen from Fabian's Code (Public Version)
-    """
-
-    def _trust_region_projection(self, p, q, eps: th.Tensor, eps_cov: th.Tensor, **kwargs):
-        """
-        Runs commutative Wasserstein projection layer and constructs sqrt of covariance
-        Args:
-            policy: policy instance
-            p: current distribution
-            q: old distribution
-            eps: (modified) kl bound/ kl bound for mean part
-            eps_cov: (modified) kl bound for cov part
-            **kwargs:
-
-        Returns:
-            mean, cov sqrt
-        """
-
-        mean, sqrt = get_mean_and_sqrt(p, expand=True)
-        old_mean, old_sqrt = get_mean_and_sqrt(q, expand=True)
-        batch_shape = mean.shape[:-1]
-
-        ####################################################################################################################
-        # precompute mean and cov part of W2, which are used for the projection.
-        # Both parts differ based on precision scaling.
-        # If activated, the mean part is the maha distance and the cov has a more complex term in the inner parenthesis.
-        mean_part, cov_part = gaussian_wasserstein_commutative(
-            p, q, self.scale_prec)
-
-        ####################################################################################################################
-        # project mean (w/ or w/o precision scaling)
-        proj_mean = mean_projection(mean, old_mean, mean_part, eps)
-
-        ####################################################################################################################
-        # project covariance (w/ or w/o precision scaling)
-
-        cov_mask = cov_part > eps_cov
-
-        if cov_mask.any():
-            # gradient issue with ch.where, it executes both paths and gives NaN gradient.
-            eta = th.ones(batch_shape, dtype=sqrt.dtype, device=sqrt.device)
-            eta[cov_mask] = th.sqrt(cov_part[cov_mask] / eps_cov) - 1.
-            eta = th.max(-eta, eta)
-
-            new_sqrt = (sqrt + th.einsum('i,ijk->ijk', eta, old_sqrt)
-                        ) / (1. + eta + 1e-16)[..., None, None]
-            proj_sqrt = th.where(cov_mask[..., None, None], new_sqrt, sqrt)
-        else:
-            proj_sqrt = sqrt
-
-        if has_diag_cov(p):
-            proj_sqrt = th.diagonal(proj_sqrt, dim1=-2, dim2=-1)
-
-        proj_p = self.new_dist_like(p, proj_mean, proj_sqrt)
-        return proj_p
-
-    def trust_region_value(self, p, q):
-        """
-        Computes the Wasserstein distance between two Gaussian distributions p and q.
-        Args:
-            policy: policy instance
-            p: current distribution
-            q: old distribution
-        Returns:
-            mean and covariance part of Wasserstein distance
-        """
-        mean_part, cov_part = gaussian_wasserstein_commutative(
-            p, q, scale_prec=self.scale_prec)
-        return mean_part + cov_part
-
-    def get_trust_region_loss(self, p, proj_p):
-        # p:
-        #   predicted distribution from network output
-        # proj_p:
-        #   projected distribution
-
-        proj_mean, proj_sqrt = get_mean_and_sqrt(proj_p)
-        p_target = self.new_dist_like(p, proj_mean, proj_sqrt)
-        kl_diff = self.trust_region_value(p, p_target)
-
-        kl_loss = kl_diff.mean()
-
-        return kl_loss * self.trust_region_coeff
-
-    def new_dist_like(self, orig_p, mean, cov_sqrt):
-        assert isinstance(orig_p, Distribution)
-        p = orig_p.distribution
-        if isinstance(p, th.distributions.Normal):
-            p_out = orig_p.__class__(orig_p.action_dim)
-            p_out.distribution = th.distributions.Normal(mean, cov_sqrt)
-        elif isinstance(p, th.distributions.Independent):
-            p_out = orig_p.__class__(orig_p.action_dim)
-            p_out.distribution = th.distributions.Independent(
-                th.distributions.Normal(mean, cov_sqrt), 1)
-        elif isinstance(p, th.distributions.MultivariateNormal):
-            p_out = orig_p.__class__(orig_p.action_dim)
-            p_out.distribution = th.distributions.MultivariateNormal(
-                mean, scale_tril=cov_sqrt, validate_args=False)
-        else:
-            raise Exception('Dist-Type not implemented (of sb3 dist)')
-        p_out.cov_sqrt = cov_sqrt
-        return p_out
-
-
-def gaussian_wasserstein_commutative(p, q, scale_prec=False) -> Tuple[th.Tensor, th.Tensor]:
-    """
-    Compute mean part and cov part of W_2(p || q_values) with p,q_values ~ N(y, SS).
-    This version DOES assume commutativity of both distributions, i.e. covariance matrices.
-    This is less general and assumes both distributions are somewhat close together.
-    When scale_prec is true scale both distributions with old precision matrix.
-    Args:
-        policy: current policy
-        p: mean and sqrt of gaussian p
-        q: mean and sqrt of gaussian q_values
-        scale_prec: scale objective by old precision matrix.
-                    This penalizes directions based on old uncertainty/covariance.
-    Returns: mean part of W2, cov part of W2
-    """
-    mean, sqrt = get_mean_and_sqrt(p, expand=True)
-    mean_other, sqrt_other = get_mean_and_sqrt(q, expand=True)
-
-    if scale_prec:
-        # maha objective for mean
-        mean_part = mahalanobis(mean, mean_other, sqrt_other)
-    else:
-        # euclidean distance for mean
-        # mean_part = ch.norm(mean_other - mean, ord=2, axis=1) ** 2
-        mean_part = ((mean_other - mean) ** 2).sum(1)
-
-    cov = get_cov(p)
-    if scale_prec and False:
-        # cov constraint scaled with precision of old dist
-        batch_dim, dim = mean.shape
-
-        identity = th.eye(dim, dtype=sqrt.dtype, device=sqrt.device)
-        sqrt_inv_other = th.linalg.solve(sqrt_other, identity)
-        c = sqrt_inv_other @ cov @ sqrt_inv_other
-
-        cov_part = _batch_trace(
-            identity + c - 2 * sqrt_inv_other @ sqrt)
-
-    else:
-        # W2 objective for cov assuming normal W2 objective for mean
-        cov_other = get_cov(q)
-        cov_part = _batch_trace(
-            cov_other + cov - 2 * th.bmm(sqrt_other, sqrt))
-
-    return mean_part, cov_part

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,57 +0,0 @@
-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
-)
-
-try:
-    import gym
-    from gym.spaces import Discrete as GymDiscrete, MultiDiscrete as GymMultiDiscrete, MultiBinary as GymMultiBinary, Box as GymBox
-    gym_available = True
-except ImportError:
-    gym_available = False
-
-def is_discrete_space(action_space):
-    discrete_types = (
-        GymnasiumDiscrete, GymnasiumMultiDiscrete, GymnasiumMultiBinary,
-        DiscreteTensorSpec, OneHotDiscreteTensorSpec, MultiDiscreteTensorSpec, BinaryDiscreteTensorSpec
-    )
-    continuous_types = (
-        GymnasiumBox, BoundedTensorSpec, UnboundedContinuousTensorSpec
-    )
-    
-    if gym_available:
-        discrete_types += (GymDiscrete, GymMultiDiscrete, GymMultiBinary)
-        continuous_types += (GymBox,)
-    
-    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_available:
-        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