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,20 +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,10 +1,3 @@
-import gymnasium
-try:
-    import fancy_gym
-except ImportError:
-    pass
+from fancy_rl.algos import PPO, TRPL #, VLEARN
 
-from fancy_rl.algos import PPO, TRPL, VLEARN
-from fancy_rl.projections import get_projection
-
-__all__ = ["PPO", "TRPL", "VLEARN", "get_projection"]
+__all__ = ['PPO', 'TRPL']

fancy_rl/algos/__init__.py

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

fancy_rl/algos/algo.py

@@ -1,8 +1,13 @@
 import torch
 import gymnasium as gym
-from torchrl.envs.libs.gym import GymWrapper
+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
 
@@ -46,19 +51,37 @@ class Algo(ABC):
         self.eval_episodes = eval_episodes
 
     def make_env(self, eval=False):
-        """Creates an environment and wraps it if necessary."""
         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 = gym.make(env_spec)
-            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)
+            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("env_spec must be a string or a callable that returns an environment.")
+            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):
@@ -70,6 +93,23 @@ class Algo(ABC):
     def evaluate(self, epoch):
         raise NotImplementedError("evaluate method must be implemented in subclass.")
 
-def dump_video(module):
-    if isinstance(module, VideoRecorder):
-        module.dump()
+    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

@@ -55,7 +55,7 @@ class OnPolicy(Algo):
         # 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,

fancy_rl/algos/ppo.py

@@ -2,9 +2,9 @@ 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
-from fancy_rl.projections import get_projection  # Updated import
 
 class PPO(OnPolicy):
     def __init__(
@@ -31,25 +31,50 @@ 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
+        
+        # Get spaces from specs for parallel env
+        self.obs_space = env.observation_spec
+        self.act_space = env.action_spec
+        
+        self.discrete = isinstance(self.act_space, DiscreteTensorSpec)
 
-        self.critic = Critic(obs_space, critic_hidden_sizes, critic_activation_fn, device)
-        actor_net = Actor(obs_space, act_space, actor_hidden_sizes, actor_activation_fn, device)
-        self.actor = ProbabilisticActor(
-            module=actor_net,
-            in_keys=["loc", "scale"],
-            out_keys=["action"],
-            distribution_class=torch.distributions.Normal,
-            return_log_prob=True
-        )
+        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"],
+            )
+        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 = {
             "actor": torch.optim.Adam(self.actor.parameters(), lr=learning_rate),
@@ -86,9 +111,9 @@ 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,
             normalize_advantage=self.normalize_advantage,
-        )
+        )

fancy_rl/algos/trpl.py

@@ -1,6 +1,7 @@
 import torch
 from torch import nn
 from typing import Dict, Any, Optional
+from torchrl.data.tensor_specs import DiscreteTensorSpec
 from torchrl.modules import ProbabilisticActor, ValueOperator
 from torchrl.objectives import ClipPPOLoss
 from torchrl.collectors import SyncDataCollector
@@ -11,6 +12,67 @@ 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__(
@@ -40,6 +102,7 @@ class TRPL(OnPolicy):
         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
@@ -47,11 +110,16 @@ class TRPL(OnPolicy):
         # 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)
-        actor_net = Actor(obs_space, act_space, actor_hidden_sizes, actor_activation_fn, device)
+        # Get spaces from specs for parallel env
+        self.obs_space = env.observation_spec
+        self.act_space = env.action_spec
+
+        assert not isinstance(self.act_space, DiscreteTensorSpec), "TRPL does not support discrete action spaces"
+
+        self.critic = Critic(self.obs_space, critic_hidden_sizes, critic_activation_fn, device)
+        self.raw_actor = Actor(self.obs_space, self.act_space, actor_hidden_sizes, actor_activation_fn, device, full_covariance=full_covariance)
+        self.old_actor = Actor(self.obs_space, self.act_space, actor_hidden_sizes, actor_activation_fn, device, full_covariance=full_covariance)
 
         # Handle projection_class
         if isinstance(projection_class, str):
@@ -60,20 +128,27 @@ class TRPL(OnPolicy):
             raise ValueError("projection_class must be a string or a subclass of BaseProjection")
 
         self.projection = projection_class(
-            in_keys=["loc", "scale"],
-            out_keys=["loc", "scale"],
-            trust_region_bound_mean=trust_region_bound_mean,
-            trust_region_bound_cov=trust_region_bound_cov
+            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 = ProbabilisticActor(
-            module=actor_net,
-            in_keys=["observation"],
-            out_keys=["loc", "scale"],
-            distribution_class=torch.distributions.Normal,
-            return_log_prob=True
+        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.old_actor = deepcopy(self.actor)
 
         self.trust_region_coef = trust_region_coef
         self.loss_module = TRPLLoss(
@@ -88,7 +163,7 @@ class TRPL(OnPolicy):
         )
 
         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)
         }
 
@@ -119,23 +194,7 @@ class TRPL(OnPolicy):
         )
 
     def update_old_policy(self):
-        self.old_actor.load_state_dict(self.actor.state_dict())
-
-    def project_policy(self, obs):
-        with torch.no_grad():
-            old_dist = self.old_actor(obs)
-        new_dist = self.actor(obs)
-        projected_params = self.projection.project(new_dist, old_dist)
-        return projected_params
-
-    def pre_update(self, tensordict):
-        obs = tensordict["observation"]
-        projected_dist = self.project_policy(obs)
-        
-        # Update tensordict with projected distribution parameters
-        tensordict["projected_loc"] = projected_dist[0]
-        tensordict["projected_scale"] = projected_dist[1]
-        return tensordict
+        self.old_actor.load_state_dict(self.raw_actor.state_dict())
 
     def post_update(self):
-        self.update_old_policy()
+        self.update_old_policy()

fancy_rl/algos/vlearn.py

@@ -5,10 +5,9 @@ from torchrl.modules import ProbabilisticActor, ValueOperator
 from torchrl.collectors import SyncDataCollector
 from torchrl.data import TensorDictReplayBuffer, LazyMemmapStorage
 
-from fancy_rl.utils import get_env, get_actor, get_critic
-from fancy_rl.modules.vlearn_loss import VLEARNLoss
-from fancy_rl.modules.projection import get_vlearn_projection
-from fancy_rl.modules.squashed_normal import get_squashed_normal
+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):

fancy_rl/objectives/trpl.py

@@ -83,8 +83,8 @@ class TRPLLoss(PPOLoss):
 
     def _trust_region_loss(self, tensordict):
         old_distribution = self.old_actor_network(tensordict)
-        raw_distribution = self.actor_network(tensordict)
-        return self.projection(self.actor_network, raw_distribution, old_distribution)
+        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)

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):
-        act_space_shape = get_space_shape(act_space)
-        if is_discrete_space(act_space):
-            out_features = act_space_shape[-1]
-        else:
-            out_features = act_space_shape[-1] * 2
+    def __init__(self, obs_space, act_space, hidden_sizes, activation_fn, device, full_covariance=False):
+        self.discrete = isinstance(act_space, DiscreteTensorSpec)
+        
+        obs_space = obs_space["observation"]
+        act_space_shape = act_space.shape[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
 
-        actor_module = nn.Sequential(
-            MLP(
-                in_features=get_space_shape(obs_space)[-1],
-                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(),
+        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
         ).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/base_projection.py

@@ -1,16 +1,92 @@
 from abc import ABC, abstractmethod
 import torch
-from typing import Dict
+from torch import nn
+from typing import Dict, List, Tuple
 
-class BaseProjection(ABC, torch.nn.Module):
-    def __init__(self, in_keys: list[str], out_keys: list[str]):
+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
 
-    def forward(self, policy_params: Dict[str, torch.Tensor], old_policy_params: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        return self.project(policy_params, old_policy_params)
+    @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/frob_projection.py

@@ -1,67 +0,0 @@
-import torch
-from .base_projection import BaseProjection
-from typing import Dict
-
-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):
-        super().__init__(in_keys=in_keys, out_keys=out_keys, trust_region_coeff=trust_region_coeff, mean_bound=mean_bound, cov_bound=cov_bound)
-        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, chol = policy_params["loc"], policy_params["scale_tril"]
-        old_mean, old_chol = old_policy_params["loc"], old_policy_params["scale_tril"]
-
-        cov = torch.matmul(chol, chol.transpose(-1, -2))
-        old_cov = torch.matmul(old_chol, old_chol.transpose(-1, -2))
-
-        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)
-
-        proj_chol = torch.linalg.cholesky(proj_cov)
-        return {"loc": proj_mean, "scale_tril": proj_chol}
-
-    def get_trust_region_loss(self, policy_params: Dict[str, torch.Tensor], proj_policy_params: Dict[str, torch.Tensor]) -> torch.Tensor:
-        mean, chol = policy_params["loc"], policy_params["scale_tril"]
-        proj_mean, proj_chol = proj_policy_params["loc"], proj_policy_params["scale_tril"]
-
-        cov = torch.matmul(chol, chol.transpose(-1, -2))
-        proj_cov = torch.matmul(proj_chol, proj_chol.transpose(-1, -2))
-
-        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, q):
-        mean, cov = p
-        old_mean, old_cov = q
-
-        if self.scale_prec:
-            prec_old = torch.inverse(old_cov)
-            mean_part = torch.sum(torch.matmul(mean - old_mean, prec_old) * (mean - old_mean), dim=-1)
-            cov_part = torch.sum(prec_old * cov, dim=(-2, -1)) - torch.logdet(torch.matmul(prec_old, cov)) - mean.shape[-1]
-        else:
-            mean_part = torch.sum(torch.square(mean - old_mean), dim=-1)
-            cov_part = torch.sum(torch.square(cov - old_cov), dim=(-2, -1))
-
-        return mean_part, cov_part
-
-    def _mean_projection(self, mean: torch.Tensor, old_mean: torch.Tensor, mean_part: torch.Tensor) -> torch.Tensor:
-        diff = mean - old_mean
-        norm = torch.sqrt(mean_part)
-        return torch.where(norm > self.mean_bound, old_mean + diff * self.mean_bound / norm.unsqueeze(-1), mean)
-
-    def _cov_projection(self, cov: torch.Tensor, old_cov: torch.Tensor, cov_part: torch.Tensor) -> torch.Tensor:
-        batch_shape = cov.shape[:-2]
-        cov_mask = cov_part > self.cov_bound
-
-        eta = torch.ones(batch_shape, dtype=cov.dtype, device=cov.device)
-        eta[cov_mask] = torch.sqrt(cov_part[cov_mask] / self.cov_bound) - 1.
-        eta = torch.max(-eta, eta)
-
-        new_cov = (cov + torch.einsum('i,ijk->ijk', eta, old_cov)) / (1. + eta + 1e-16)[..., None, None]
-        proj_cov = torch.where(cov_mask[..., None, None], new_cov, cov)
-
-        return proj_cov

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

@@ -3,8 +3,8 @@ 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):
-        super().__init__(in_keys=in_keys, out_keys=out_keys, trust_region_coeff=trust_region_coeff, mean_bound=mean_bound, cov_bound=cov_bound)
+    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

fancy_rl/projections/kl_projection.py

@@ -1,7 +1,12 @@
 import torch
-import cpp_projection
+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
@@ -10,107 +15,131 @@ 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, is_diag: bool = True, 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)
-        self.is_diag = is_diag
-        self.contextual_std = contextual_std
+    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]:
-        mean, std = policy_params["loc"], policy_params["scale_tril"]
-        old_mean, old_std = old_policy_params["loc"], old_policy_params["scale_tril"]
+        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, std), (old_mean, old_std))
+        mean_part, cov_part = self._gaussian_kl((mean, scale_or_tril), (old_mean, old_scale_or_tril))
 
         if not self.contextual_std:
-            std = std[:1]
-            old_std = old_std[:1]
+            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_std = self._cov_projection(std, old_std, cov_part)
+        proj_scale_or_tril = self._cov_projection(scale_or_tril, old_scale_or_tril, cov_part)
 
         if not self.contextual_std:
-            proj_std = proj_std.expand(mean.shape[0], -1, -1)
+            proj_scale_or_tril = proj_scale_or_tril.expand(mean.shape[0], *proj_scale_or_tril.shape[1:])
 
-        return {"loc": proj_mean, "scale_tril": proj_std}
+        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, std = policy_params["loc"], policy_params["scale_tril"]
-        proj_mean, proj_std = proj_policy_params["loc"], proj_policy_params["scale_tril"]
-        kl = sum(self._gaussian_kl((mean, std), (proj_mean, proj_std)))
+        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, std = p
-        mean_other, std_other = q
+        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, std_other)
+        maha_part = 0.5 * self._maha(mean, mean_other, scale_or_tril_other)
 
-        det_term = self._log_determinant(std)
-        det_term_other = self._log_determinant(std_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)
 
-        trace_part = self._torch_batched_trace_square(torch.linalg.solve_triangular(std_other, std, upper=False))
         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, std: torch.Tensor) -> torch.Tensor:
+    def _maha(self, x: torch.Tensor, y: torch.Tensor, scale_or_tril: torch.Tensor) -> torch.Tensor:
         diff = x - y
-        return torch.sum(torch.square(torch.triangular_solve(diff.unsqueeze(-1), std, upper=False)[0].squeeze(-1)), dim=-1)
-
-    def _log_determinant(self, std: torch.Tensor) -> torch.Tensor:
-        return 2 * torch.log(std.diagonal(dim1=-2, dim2=-1)).sum(-1)
-
-    def _torch_batched_trace_square(self, x: torch.Tensor) -> torch.Tensor:
-        return torch.sum(x.pow(2), dim=(-2, -1))
-
-    def _mean_projection(self, mean: torch.Tensor, old_mean: torch.Tensor, mean_part: torch.Tensor) -> torch.Tensor:
-        return old_mean + (mean - old_mean) * torch.sqrt(self.mean_bound / (mean_part + 1e-8)).unsqueeze(-1)
-
-    def _cov_projection(self, std: torch.Tensor, old_std: torch.Tensor, cov_part: torch.Tensor) -> torch.Tensor:
-        cov = torch.matmul(std, std.transpose(-1, -2))
-        old_cov = torch.matmul(old_std, old_std.transpose(-1, -2))
-
-        if self.is_diag:
-            mask = cov_part > self.cov_bound
-            proj_std = torch.zeros_like(std)
-            proj_std[~mask] = std[~mask]
-            try:
-                if mask.any():
-                    proj_cov = KLProjectionGradFunctionDiagCovOnly.apply(cov.diagonal(dim1=-2, dim2=-1),
-                                                                         old_cov.diagonal(dim1=-2, dim2=-1),
-                                                                         self.cov_bound)
-                    is_invalid = (proj_cov.mean(dim=-1).isnan() | proj_cov.mean(dim=-1).isinf() | (proj_cov.min(dim=-1).values < 0)) & mask
-                    if is_invalid.any():
-                        proj_std[is_invalid] = old_std[is_invalid]
-                        mask &= ~is_invalid
-                    proj_std[mask] = proj_cov[mask].sqrt().diag_embed()
-            except Exception as e:
-                proj_std = old_std
+        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:
-            try:
-                mask = cov_part > self.cov_bound
-                proj_std = torch.zeros_like(std)
-                proj_std[~mask] = std[~mask]
-                if mask.any():
-                    proj_cov = KLProjectionGradFunctionCovOnly.apply(cov, std.detach(), old_std, self.cov_bound)
-                    is_invalid = proj_cov.mean([-2, -1]).isnan() & mask
-                    if is_invalid.any():
-                        proj_std[is_invalid] = old_std[is_invalid]
-                        mask &= ~is_invalid
-                    proj_std[mask], failed_mask = torch.linalg.cholesky_ex(proj_cov[mask])
-                    failed_mask = failed_mask.bool()
-                    if failed_mask.any():
-                        proj_std[failed_mask] = old_std[failed_mask]
-            except Exception as e:
-                import logging
-                logging.error('Projection failed, taking old cholesky for projection.')
-                print("Projection failed, taking old cholesky for projection.")
-                proj_std = old_std
-                raise e
+            return torch.sum(torch.square(diff / scale_or_tril), dim=-1)
 
-        return proj_std
+    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):

fancy_rl/projections/wasserstein_projection.py

@@ -2,55 +2,160 @@ import torch
 from .base_projection import BaseProjection
 from typing import Dict, Tuple
 
-def gaussian_wasserstein_commutative(policy, p: Tuple[torch.Tensor, torch.Tensor],
-                                     q: Tuple[torch.Tensor, torch.Tensor], scale_prec=False) -> Tuple[torch.Tensor, torch.Tensor]:
-    mean, sqrt = p
-    mean_other, sqrt_other = q
+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)
 
-    cov = torch.matmul(sqrt, sqrt.transpose(-1, -2))
-    cov_other = torch.matmul(sqrt_other, sqrt_other.transpose(-1, -2))
+    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)
+            )
 
-    if scale_prec:
-        identity = torch.eye(mean.shape[-1], dtype=sqrt.dtype, device=sqrt.device)
-        sqrt_inv_other = torch.linalg.solve(sqrt_other, identity)
-        c = sqrt_inv_other @ cov @ sqrt_inv_other
-        cov_part = torch.trace(identity + c - 2 * sqrt_inv_other @ sqrt)
-    else:
-        cov_part = torch.trace(cov_other + cov - 2 * sqrt_other @ sqrt)
-
-    return mean_part, cov_part
+    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):
-        super().__init__(in_keys=in_keys, out_keys=out_keys, trust_region_coeff=trust_region_coeff, mean_bound=mean_bound, cov_bound=cov_bound)
+    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, sqrt = policy_params["loc"], policy_params["scale_tril"]
-        old_mean, old_sqrt = old_policy_params["loc"], old_policy_params["scale_tril"]
+        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]]
 
-        mean_part, cov_part = gaussian_wasserstein_commutative(None, (mean, sqrt), (old_mean, old_sqrt), self.scale_prec)
+        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_sqrt = self._cov_projection(sqrt, old_sqrt, cov_part)
+        proj_scale_sqrt = self._scale_projection(scale_sqrt, old_scale_sqrt, scale_part)
 
-        return {"loc": proj_mean, "scale_tril": proj_sqrt}
+        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, sqrt = policy_params["loc"], policy_params["scale_tril"]
-        proj_mean, proj_sqrt = proj_policy_params["loc"], proj_policy_params["scale_tril"]
-        mean_part, cov_part = gaussian_wasserstein_commutative(None, (mean, sqrt), (proj_mean, proj_sqrt), self.scale_prec)
-        w2 = mean_part + cov_part
+        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 _mean_projection(self, mean: torch.Tensor, old_mean: torch.Tensor, mean_part: torch.Tensor) -> torch.Tensor:
-        diff = mean - old_mean
-        norm = torch.norm(diff, dim=-1, keepdim=True)
-        return torch.where(norm > self.mean_bound, old_mean + diff * self.mean_bound / norm, mean)
+    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 _cov_projection(self, sqrt: torch.Tensor, old_sqrt: torch.Tensor, cov_part: torch.Tensor) -> torch.Tensor:
-        diff = sqrt - old_sqrt
-        norm = torch.norm(diff, dim=(-2, -1), keepdim=True)
-        return torch.where(norm > self.cov_bound, old_sqrt + diff * self.cov_bound / norm, sqrt)
+    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

@@ -8,7 +8,7 @@
    description = "Minimalistic and efficient implementations of PPO and TRPL for torchrl"
    authors = [{name = "Dominik Roth", email = "mail@dominik-roth.eu"}]
    readme = "README.md"
-   requires-python = ">=3.7,<3.12"
+   requires-python = ">=3.7"
    classifiers = [
        "Development Status :: 3 - Alpha",
        "Intended Audience :: Developers",
@@ -23,13 +23,15 @@
    dependencies = [
        "numpy",
        "torch",
-       "gymnasium",
+       "gymnasium<1.0",
        "tensordict",
        "torchrl",
+       "pytest",
    ]
 
    [project.urls]
    Homepage = "https://git.dominik-roth.eu/dodox/fancy_rl"
 
    [project.optional-dependencies]
-   dev = ["pytest"]
+   dev = ["pytest"]
+   box2d = ["swig", "gymnasium[box2d]"]

test/test_ppo.py

@@ -2,58 +2,53 @@ 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
 
-@pytest.fixture
 def simple_env():
-    return gym.make('CartPole-v1')
+    return gym.make('LunarLander-v2')
 
 def test_ppo_instantiation():
-    ppo = PPO("CartPole-v1")
+    ppo = PPO(simple_env)
     assert isinstance(ppo, PPO)
 
-@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("n_epochs", [5, 10])
-@pytest.mark.parametrize("gamma", [0.95, 0.99])
-@pytest.mark.parametrize("clip_range", [0.1, 0.2, 0.3])
-def test_ppo_initialization_with_different_hps(learning_rate, n_steps, batch_size, n_epochs, gamma, clip_range):
-    ppo = PPO(
-        "CartPole-v1",
-        learning_rate=learning_rate,
-        n_steps=n_steps,
-        batch_size=batch_size,
-        n_epochs=n_epochs,
-        gamma=gamma,
-        clip_range=clip_range
-    )
-    assert ppo.learning_rate == learning_rate
-    assert ppo.n_steps == n_steps
-    assert ppo.batch_size == batch_size
-    assert ppo.n_epochs == n_epochs
-    assert ppo.gamma == gamma
-    assert ppo.clip_range == clip_range
+def test_ppo_instantiation_from_str():
+    ppo = PPO('CartPole-v1')
+    assert isinstance(ppo, PPO)
 
-def test_ppo_predict(simple_env):
-    ppo = PPO("CartPole-v1")
-    obs, _ = simple_env.reset()
-    action, _ = ppo.predict(obs)
-    assert isinstance(action, np.ndarray)
-    assert action.shape == simple_env.action_space.shape
-
-def test_ppo_learn():
-    ppo = PPO("CartPole-v1", n_steps=64, batch_size=32)
-    env = gym.make("CartPole-v1")
-    obs, _ = env.reset()
-    for _ in range(64):
-        action, _ = ppo.predict(obs)
-        next_obs, reward, done, truncated, _ = env.step(action)
-        ppo.store_transition(obs, action, reward, done, next_obs)
-        obs = next_obs
-        if done or truncated:
-            obs, _ = env.reset()
+def test_ppo_predict():
+    ppo = PPO(simple_env)
+    env = ppo.make_env()
+    obs = env.reset()
+    action = ppo.predict(obs)
+    assert isinstance(action, TensorDict)
     
-    loss = ppo.learn()
-    assert isinstance(loss, dict)
-    assert "policy_loss" in loss
-    assert "value_loss" in loss
+    # 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

@@ -2,76 +2,74 @@ import pytest
 import numpy as np
 from fancy_rl import TRPL
 import gymnasium as gym
+from tensordict import TensorDict
 
-@pytest.fixture
 def simple_env():
-    return gym.make('CartPole-v1')
+    return gym.make('Pendulum-v1')
 
 def test_trpl_instantiation():
-    trpl = TRPL("CartPole-v1")
+    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("max_kl", [0.01, 0.05])
-def test_trpl_initialization_with_different_hps(learning_rate, n_steps, batch_size, gamma, max_kl):
+@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(
-        "CartPole-v1",
+        simple_env,
         learning_rate=learning_rate,
         n_steps=n_steps,
         batch_size=batch_size,
         gamma=gamma,
-        max_kl=max_kl
+        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.max_kl == max_kl
+    assert trpl.projection.mean_bound == trust_region_bound_mean
+    assert trpl.projection.cov_bound == trust_region_bound_cov
 
-def test_trpl_predict(simple_env):
-    trpl = TRPL("CartPole-v1")
-    obs, _ = simple_env.reset()
-    action, _ = trpl.predict(obs)
-    assert isinstance(action, np.ndarray)
-    assert action.shape == simple_env.action_space.shape
-
-def test_trpl_learn():
-    trpl = TRPL("CartPole-v1", n_steps=64, batch_size=32)
-    env = gym.make("CartPole-v1")
-    obs, _ = env.reset()
-    for _ in range(64):
-        action, _ = trpl.predict(obs)
-        next_obs, reward, done, truncated, _ = env.step(action)
-        trpl.store_transition(obs, action, reward, done, next_obs)
-        obs = next_obs
-        if done or truncated:
-            obs, _ = env.reset()
+def test_trpl_predict():
+    trpl = TRPL(simple_env)
+    env = trpl.make_env()
+    obs = env.reset()
+    action = trpl.predict(obs)
+    assert isinstance(action, TensorDict)
     
-    loss = trpl.learn()
-    assert isinstance(loss, dict)
-    assert "policy_loss" in loss
-    assert "value_loss" in loss
+    # 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(simple_env):
-    trpl = TRPL("CartPole-v1", total_timesteps=10000)
+def test_trpl_training():
+    trpl = TRPL(simple_env, total_timesteps=100)
+    env = trpl.make_env()
     
-    initial_performance = evaluate_policy(trpl, simple_env)
+    initial_performance = evaluate_policy(trpl, env)
     trpl.train()
-    final_performance = evaluate_policy(trpl, simple_env)
-    
-    assert final_performance > initial_performance, "TRPL should improve performance after training"
+    final_performance = evaluate_policy(trpl, env)
 
-def evaluate_policy(policy, env, n_eval_episodes=10):
+def evaluate_policy(policy, env, n_eval_episodes=3):
     total_reward = 0
     for _ in range(n_eval_episodes):
-        obs, _ = env.reset()
+        tensordict = 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
+            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