README.md

@@ -23,11 +23,11 @@ Fancy RL provides two main components:
 1. **Ready-to-use Classes for PPO / TRPL**: These classes allow you to quickly get started with reinforcement learning algorithms, enjoying the performance and hackability that comes with using TorchRL.
 
    ```python
-   from fancy_rl import PPO, TRPL
+   from fancy_rl import PPO
 
-   model = TRPL("CartPole-v1")
+   ppo = PPO("CartPole-v1")
 
-   model.train()
+   ppo.train()
    ```
 
    For environments, you can pass any [gymnasium](https://gymnasium.farama.org/) or [Fancy Gym](https://alrhub.github.io/fancy_gym/) environment ID as a string, a function returning a gymnasium environment, or an already instantiated gymnasium environment. Future plans include supporting other torchrl environments.

fancy_rl/__init__.py

@@ -1,9 +1,6 @@
-import gymnasium
-try:
-    import fancy_gym
-except ImportError:
-    pass
-
 from fancy_rl.ppo import PPO
+from fancy_rl.policy import MLPPolicy
+from fancy_rl.loggers import TerminalLogger, WandbLogger
+from fancy_rl.utils import make_env
 
-__all__ = ["PPO"]
+__all__ = ["PPO", "MLPPolicy", "TerminalLogger", "WandbLogger", "make_env"]

fancy_rl/on_policy.py

@@ -1,13 +1,18 @@
 import torch
-import gymnasium as gym
+from abc import ABC, abstractmethod
+from torchrl.record.loggers import Logger
+from torch.optim import Adam
 from torchrl.collectors import SyncDataCollector
 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.envs.libs.gym import GymWrapper
 from torchrl.record import VideoRecorder
-from abc import ABC, abstractmethod
-
+import gymnasium as gym
+try:
+    import fancy_gym
+except ImportError:
+    pass
 
 class OnPolicy(ABC):
     def __init__(
@@ -15,7 +20,6 @@ class OnPolicy(ABC):
         policy,
         env_spec,
         loggers,
-        optimizers,
         learning_rate,
         n_steps,
         batch_size,
@@ -37,7 +41,6 @@ class OnPolicy(ABC):
         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
-        self.optimizers = optimizers
         self.learning_rate = learning_rate
         self.n_steps = n_steps
         self.batch_size = batch_size
@@ -87,15 +90,6 @@ class OnPolicy(ABC):
             raise ValueError("env_spec must be a string or a callable that returns an environment.")
         return env
 
-    def train_step(self, batch):
-        for optimizer in self.optimizers.values():
-            optimizer.zero_grad()
-        loss = self.loss_module(batch)
-        loss.backward()
-        for optimizer in self.optimizers.values():
-            optimizer.step()
-        return loss
-
     def train(self):
         collected_frames = 0
 
@@ -142,6 +136,10 @@ class OnPolicy(ABC):
         for logger in self.loggers:
             logger.log_scalar({"eval_avg_return": avg_return}, step=epoch)
 
+    @abstractmethod
+    def train_step(self, batch):
+        pass
+
 def dump_video(module):
     if isinstance(module, VideoRecorder):
         module.dump()

fancy_rl/policy.py

@@ -1,59 +1,71 @@
-import torch.nn as nn
-from tensordict.nn import TensorDictModule
-from torchrl.modules import MLP
-from tensordict.nn.distributions import NormalParamExtractor
+import torch
+from torch import nn
+from torch.distributions import Categorical, Normal
+import gymnasium as gym
 
-class SharedModule(TensorDictModule):
-    def __init__(self, obs_space, hidden_sizes, activation_fn, device):
-        if hidden_sizes:
-            shared_module = MLP(
-                in_features=obs_space.shape[-1],
-                out_features=hidden_sizes[-1],
-                num_cells=hidden_sizes,
-                activation_class=getattr(nn, activation_fn),
-                device=device
-            )
-            out_features = hidden_sizes[-1]
+class Actor(nn.Module):
+    def __init__(self, observation_space, action_space, hidden_sizes=[64, 64], activation_fn=nn.ReLU):
+        super().__init__()
+        self.continuous = isinstance(action_space, gym.spaces.Box)
+        input_dim = observation_space.shape[-1]
+        if self.continuous:
+            output_dim = action_space.shape[-1]
         else:
-            shared_module = nn.Identity()
-            out_features = obs_space.shape[-1]
+            output_dim = action_space.n
         
-        super().__init__(
-            module=shared_module,
-            in_keys=["observation"],
-            out_keys=["shared"],
-        )
-        self.out_features = out_features
+        layers = []
+        last_dim = input_dim
+        for size in hidden_sizes:
+            layers.append(nn.Linear(last_dim, size))
+            layers.append(activation_fn())
+            last_dim = size
+        
+        if self.continuous:
+            self.mu_layer = nn.Linear(last_dim, output_dim)
+            self.log_std_layer = nn.Linear(last_dim, output_dim)
+        else:
+            layers.append(nn.Linear(last_dim, output_dim))
+            self.model = nn.Sequential(*layers)
 
-class Actor(TensorDictModule):
-    def __init__(self, shared_module, act_space, hidden_sizes, activation_fn, device):
-        actor_module = nn.Sequential(
-            MLP(
-                in_features=shared_module.out_features,
-                out_features=act_space.shape[-1] * 2,
-                num_cells=hidden_sizes,
-                activation_class=getattr(nn, activation_fn),
-                device=device
-            ),
-            NormalParamExtractor(),
-        ).to(device)
-        super().__init__(
-            module=actor_module,
-            in_keys=["shared"],
-            out_keys=["loc", "scale"],
-        )
+    def forward(self, x):
+        if self.continuous:
+            mu = self.mu_layer(x)
+            log_std = self.log_std_layer(x)
+            return mu, log_std.exp()
+        else:
+            return self.model(x)
 
-class Critic(TensorDictModule):
-    def __init__(self, shared_module, hidden_sizes, activation_fn, device):
-        critic_module = MLP(
-            in_features=shared_module.out_features,
-            out_features=1,
-            num_cells=hidden_sizes,
-            activation_class=getattr(nn, activation_fn),
-            device=device
-        ).to(device)
-        super().__init__(
-            module=critic_module,
-            in_keys=["shared"],
-            out_keys=["state_value"],
-        )
+    def act(self, observation, deterministic=False):
+        with torch.no_grad():
+            if self.continuous:
+                mu, std = self.forward(observation)
+                if deterministic:
+                    action = mu
+                else:
+                    action_dist = Normal(mu, std)
+                    action = action_dist.sample()
+            else:
+                logits = self.forward(observation)
+                if deterministic:
+                    action = logits.argmax(dim=-1)
+                else:
+                    action_dist = Categorical(logits=logits)
+                    action = action_dist.sample()
+        return action
+
+class Critic(nn.Module):
+    def __init__(self, observation_space, hidden_sizes=[64, 64], activation_fn=nn.ReLU):
+        super().__init__()
+        input_dim = observation_space.shape[-1]
+        
+        layers = []
+        last_dim = input_dim
+        for size in hidden_sizes:
+            layers.append(nn.Linear(last_dim, size))
+            layers.append(activation_fn())
+            last_dim = size
+        layers.append(nn.Linear(last_dim, 1))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x).squeeze(-1)

fancy_rl/ppo.py

@@ -1,9 +1,11 @@
 import torch
-from torchrl.modules import ActorValueOperator, ProbabilisticActor
+import torch.nn as nn
 from torchrl.objectives import ClipPPOLoss
 from torchrl.objectives.value.advantages import GAE
-from fancy_rl.on_policy import OnPolicy
-from fancy_rl.policy import Actor, Critic, SharedModule
+from torchrl.record.loggers import get_logger
+from on_policy import OnPolicy
+from policy import Actor, Critic
+import gymnasium as gym
 
 class PPO(OnPolicy):
     def __init__(
@@ -12,9 +14,8 @@ class PPO(OnPolicy):
         loggers=None,
         actor_hidden_sizes=[64, 64],
         critic_hidden_sizes=[64, 64],
-        actor_activation_fn="Tanh",
-        critic_activation_fn="Tanh",
-        shared_stem_sizes=[64],
+        actor_activation_fn="ReLU",
+        critic_activation_fn="ReLU",
         learning_rate=3e-4,
         n_steps=2048,
         batch_size=64,
@@ -32,45 +33,21 @@ class PPO(OnPolicy):
         env_spec_eval=None,
         eval_episodes=10,
     ):
-        device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
         # Initialize environment to get observation and action space sizes
-        self.env_spec = env_spec
-        env = self.make_env()
+        env = self.make_env(env_spec)
         obs_space = env.observation_space
         act_space = env.action_space
 
-        # Define the shared, actor, and critic modules
-        self.shared_module = SharedModule(obs_space, shared_stem_sizes, actor_activation_fn, device)
-        self.actor = Actor(self.shared_module, act_space, actor_hidden_sizes, actor_activation_fn, device)
-        self.critic = Critic(self.shared_module, critic_hidden_sizes, critic_activation_fn, device)
+        actor_activation_fn = getattr(nn, actor_activation_fn)
+        critic_activation_fn = getattr(nn, critic_activation_fn)
 
-        # Combine into an ActorValueOperator
-        self.ac_module = ActorValueOperator(
-            self.shared_module,
-            self.actor,
-            self.critic
-        )
-
-        # Define the policy as a ProbabilisticActor
-        self.policy = ProbabilisticActor(
-            module=self.ac_module.get_policy_operator(),
-            in_keys=["loc", "scale"],
-            out_keys=["action"],
-            distribution_class=torch.distributions.Normal,
-            return_log_prob=True
-        )
-
-        optimizers = {
-            "actor": torch.optim.Adam(self.actor.parameters(), lr=learning_rate),
-            "critic": torch.optim.Adam(self.critic.parameters(), lr=learning_rate)
-        }
+        self.actor = Actor(obs_space, act_space, hidden_sizes=actor_hidden_sizes, activation_fn=actor_activation_fn)
+        self.critic = Critic(obs_space, hidden_sizes=critic_hidden_sizes, activation_fn=critic_activation_fn)
 
         super().__init__(
-            policy=self.policy,
+            policy=self.actor,
             env_spec=env_spec,
             loggers=loggers,
-            optimizers=optimizers,
             learning_rate=learning_rate,
             n_steps=n_steps,
             batch_size=batch_size,
@@ -105,3 +82,15 @@ class PPO(OnPolicy):
             critic_coef=self.critic_coef,
             normalize_advantage=self.normalize_advantage,
         )
+
+        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=self.learning_rate)
+        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=self.learning_rate)
+
+    def train_step(self, batch):
+        self.actor_optimizer.zero_grad()
+        self.critic_optimizer.zero_grad()
+        loss = self.loss_module(batch)
+        loss.backward()
+        self.actor_optimizer.step()
+        self.critic_optimizer.step()
+        return loss

pyproject.toml

@@ -1,13 +0,0 @@
-[build-system]
-requires = ["setuptools", "wheel"]
-build-backend = "setuptools.build_meta"
-
-[project]
-name = "fancy_rl"
-version = "0.1.0"
-dependencies = [
-    "gymnasium",
-    "pyyaml",
-    "torch",
-    "torchrl"
-]

setup.py

@@ -0,0 +1,19 @@
+
+from setuptools import setup, find_packages
+
+setup(
+    name="fancy_rl",
+    version="0.1",
+    packages=find_packages(),
+    install_requires=[
+        "torch",
+        "torchrl",
+        "gymnasium",
+        "pyyaml",
+    ],
+    entry_points={
+        "console_scripts": [
+            "fancy_rl=fancy_rl.example:main",
+        ],
+    },
+)

test/test_ppo.py

@@ -1 +1,54 @@
-# TODO
+import pytest
+import torch
+from fancy_rl.ppo import PPO
+from fancy_rl.policy import Policy
+from fancy_rl.loggers import TerminalLogger
+from fancy_rl.utils import make_env
+
+@pytest.fixture
+def policy():
+    return Policy(input_dim=4, output_dim=2, hidden_sizes=[64, 64])
+
+@pytest.fixture
+def loggers():
+    return [TerminalLogger()]
+
+@pytest.fixture
+def env_fn():
+    return make_env("CartPole-v1")
+
+def test_ppo_train(policy, loggers, env_fn):
+    ppo = PPO(policy=policy,
+              env_fn=env_fn,
+              loggers=loggers,
+              learning_rate=3e-4,
+              n_steps=2048,
+              batch_size=64,
+              n_epochs=10,
+              gamma=0.99,
+              gae_lambda=0.95,
+              clip_range=0.2,
+              total_timesteps=10000,
+              eval_interval=2048,
+              eval_deterministic=True,
+              eval_episodes=5,
+              seed=42)
+    ppo.train()
+
+def test_ppo_evaluate(policy, loggers, env_fn):
+    ppo = PPO(policy=policy,
+              env_fn=env_fn,
+              loggers=loggers,
+              learning_rate=3e-4,
+              n_steps=2048,
+              batch_size=64,
+              n_epochs=10,
+              gamma=0.99,
+              gae_lambda=0.95,
+              clip_range=0.2,
+              total_timesteps=10000,
+              eval_interval=2048,
+              eval_deterministic=True,
+              eval_episodes=5,
+              seed=42)
+    ppo.evaluate(epoch=0)