Move contents to src. Use UV instead of setuptools

.python-version

@@ -0,0 +1 @@
+3.13

pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "reinforcelab"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "dill>=0.3.6",
+    "gymnasium>=0.28",
+    "matplotlib>=3.0",
+    "opencv-python>=4.0",
+    "seaborn>=0.12.2",
+    "torch>=2.9.1",
+    "tqdm>=4.64.1",
+]

src/reinforcelab/modules/action_selectors/__init__.py

@@ -0,0 +1,5 @@
+from .action_selector import ActionSelector
+from .discrete import DiscreteActionSelector
+from .continuous import ContinuousActionSelector, NoisyAction, ContinuousEpsilonGreedy
+
+__all__ = ["ActionSelector", "DiscreteActionSelector", "ContinuousActionSelector", "NoisyAction", "ContinuousEpsilonGreedy"]

src/reinforcelab/modules/action_selectors/continuous/__init__.py

@@ -0,0 +1,5 @@
+from .continuous_action_selector import ContinuousActionSelector
+from .noisy_action import NoisyAction
+from .epsilon_greedy import ContinuousEpsilonGreedy
+
+__all__ = ["ContinuousActionSelector", "NoisyAction", "ContinuousEpsilonGreedy"]

reinforcelab/action_selectors/continuous/epsilon_greedy.py → src/reinforcelab/modules/action_selectors/continuous/epsilon_greedy.py

@@ -4,7 +4,7 @@
 
 import gymnasium as gym
 
-from reinforcelab.utils import tie_breaker, get_space_size, space_is_type
+from reinforcelab.modules.utils import tie_breaker, get_space_size, space_is_type
 
 
 class ContinuousEpsilonGreedy(ContinuousActionSelector):

reinforcelab/action_selectors/continuous/noisy_action.py → src/reinforcelab/modules/action_selectors/continuous/noisy_action.py

@@ -4,7 +4,7 @@
 
 import gymnasium as gym
 
-from reinforcelab.utils import tie_breaker, get_space_size, space_is_type
+from reinforcelab.modules.utils import tie_breaker, get_space_size, space_is_type
 
 
 class NoisyAction(ContinuousActionSelector):

reinforcelab/action_selectors/discrete/__init__.py → src/reinforcelab/modules/action_selectors/discrete/__init__.py

@@ -1,2 +1,4 @@
 from .discrete_action_selector import DiscreteActionSelector
 from .epsilon_greedy import EpsilonGreedy
+
+__all__ = ["DiscreteActionSelector", "EpsilonGreedy"]

reinforcelab/action_selectors/discrete/epsilon_greedy.py → src/reinforcelab/modules/action_selectors/discrete/epsilon_greedy.py

@@ -3,7 +3,7 @@
 
 import gymnasium as gym
 
-from reinforcelab.utils import tie_breaker, get_space_size, space_is_type
+from reinforcelab.modules.utils import tie_breaker, get_space_size, space_is_type
 
 
 class EpsilonGreedy(DiscreteActionSelector):

src/reinforcelab/modules/agents/__init__.py

@@ -0,0 +1,5 @@
+from .agent import BaseAgent, Agent
+from .value_optimization import QLearning, SARSA, ExpectedSARSA, DQN, DCQN
+from .policy_gradient import DDPG
+
+__all__ = ["BaseAgent", "Agent", "QLearning", "SARSA", "ExpectedSARSA", "DQN", "DCQN", "DDPG"]

reinforcelab/agents/agent.py → src/reinforcelab/modules/agents/agent.py

@@ -1,24 +1,22 @@
 import os
 import dill
 from torch import Tensor
-from typing import Union
 from abc import ABC, abstractmethod
 
-from reinforcelab.brains import Brain
-from reinforcelab.estimators import Estimator
-from reinforcelab.memory_buffers import MemoryBuffer
-from reinforcelab.action_selectors import ActionSelector
-from reinforcelab.experience import Experience
+from reinforcelab.modules.brains import Brain
+from reinforcelab.modules.memory_buffers import MemoryBuffer
+from reinforcelab.modules.action_selectors import ActionSelector
+from reinforcelab.modules.experience import Experience
 
 
 class BaseAgent(ABC):
     @abstractmethod
-    def act(self, state: Tensor, epsilon: float = 0.0):
+    def act(self, state: Tensor, **kwargs):
         """Choose an action given a state
 
         Args:
-            state (Any): A representation of the state
-            epsilon (float, optional): Probability of taking an exploratory action. Defaults to 0.0.
+            state (Tensor): A representation of the state
+            **kwargs: Additional arguments to the action selector
         """
 
     @abstractmethod

src/reinforcelab/modules/agents/policy_gradient/__init__.py

@@ -0,0 +1,3 @@
+from .ddpg import DDPG
+
+__all__ = ["DDPG"]

reinforcelab/agents/policy_gradient/ddpg.py → src/reinforcelab/modules/agents/policy_gradient/ddpg.py

@@ -2,11 +2,11 @@
 from torch import nn
 
 
-from reinforcelab.agents import Agent
-from reinforcelab.brains import ActorCritic
-from reinforcelab.action_selectors.continuous import *
-from reinforcelab.estimators import SARSEstimator
-from reinforcelab.memory_buffers import ExperienceReplay
+from reinforcelab.modules.agents import Agent
+from reinforcelab.modules.brains import ActorCritic
+from reinforcelab.modules.action_selectors.continuous import ContinuousEpsilonGreedy
+from reinforcelab.modules.estimators import SARSEstimator
+from reinforcelab.modules.memory_buffers import ExperienceReplay
 
 class DDPG(Agent):
     def __init__(self, env: Env, actor_model: nn.Module, critic_model: nn.Module, learning_rate=0.01, discount_factor=0.999, alpha=0.03, batch_size=128, update_every=4, max_buffer_size=2**12):

src/reinforcelab/modules/agents/value_optimization/__init__.py

@@ -0,0 +1,4 @@
+from .tabular import QLearning, SARSA, ExpectedSARSA
+from .function_approximation import DQN, DCQN
+
+__all__ = ["QLearning", "SARSA", "ExpectedSARSA", "DQN", "DCQN"]

reinforcelab/agents/value_optimization/function_approximation/__init__.py → src/reinforcelab/modules/agents/value_optimization/function_approximation/__init__.py

@@ -1,2 +1,4 @@
 from .dqn import DQN
 from .dcqn import DCQN
+
+__all__ = ["DQN", "DCQN"]

reinforcelab/agents/value_optimization/function_approximation/dcqn.py → src/reinforcelab/modules/agents/value_optimization/function_approximation/dcqn.py

@@ -1,13 +1,12 @@
 from gymnasium import Env
 from torch import nn
 
-from reinforcelab.agents.agent import Agent
-from reinforcelab.brains import QNetwork
-from reinforcelab.estimators import MaxQEstimator
-from reinforcelab.transforms.experience import IntrinsicCuriosityModule
-from reinforcelab.action_selectors import EpsilonGreedy
-from reinforcelab.memory_buffers import ExperienceReplay
-from reinforcelab.utils import get_state_action_sizes
+from reinforcelab.modules.agents.agent import Agent
+from reinforcelab.modules.brains import QNetwork
+from reinforcelab.modules.estimators import MaxQEstimator
+from reinforcelab.modules.transforms.experience import IntrinsicCuriosityModule
+from reinforcelab.modules.action_selectors import EpsilonGreedy
+from reinforcelab.modules.memory_buffers import ExperienceReplay
 
 
 class DCQN(Agent):
@@ -17,7 +16,17 @@ class DCQN(Agent):
     procedure.
     """
 
-    def __init__(self, env: Env, model: nn.Module, learning_rate=0.01, discount_factor: float = 0.999, alpha=0.03, batch_size=128, update_every=4, max_buffer_size=2**12):
+    def __init__(
+        self,
+        env: Env,
+        model: nn.Module,
+        learning_rate=0.01,
+        discount_factor: float = 0.999,
+        alpha=0.03,
+        batch_size=128,
+        update_every=4,
+        max_buffer_size=2**12,
+    ):
         action_selector = EpsilonGreedy(env)
         icm = IntrinsicCuriosityModule(
             env, 4, learning_rate=0.0001, state_transform_hidden_layers=[4, 4])
@@ -26,5 +35,9 @@ def __init__(self, env: Env, model: nn.Module, learning_rate=0.01, discount_fact
         estimator = MaxQEstimator(env, discount_factor)
         brain = QNetwork(model, estimator, learning_rate, alpha)
 
-        super().__init__(brain,
-                         action_selector, buffer, update_every=update_every)
+        super().__init__(
+            brain,
+            action_selector,
+            buffer,
+            update_every=update_every
+        )

reinforcelab/agents/value_optimization/function_approximation/dqn.py → src/reinforcelab/modules/agents/value_optimization/function_approximation/dqn.py

@@ -1,12 +1,11 @@
 from gymnasium import Env
-from copy import deepcopy
 from torch import nn
 
-from reinforcelab.agents.agent import Agent
-from reinforcelab.brains import QNetwork
-from reinforcelab.estimators import MaxQEstimator
-from reinforcelab.action_selectors import EpsilonGreedy
-from reinforcelab.memory_buffers import ExperienceReplay
+from reinforcelab.modules.agents.agent import Agent
+from reinforcelab.modules.brains import QNetwork
+from reinforcelab.modules.estimators import MaxQEstimator
+from reinforcelab.modules.action_selectors import EpsilonGreedy
+from reinforcelab.modules.memory_buffers import ExperienceReplay
 
 
 class DQN(Agent):
@@ -16,7 +15,17 @@ class DQN(Agent):
     procedure.
     """
 
-    def __init__(self, env: Env, model: nn.Module, learning_rate=0.01, discount_factor: float = 0.999, alpha=0.03, batch_size=128, update_every=4, max_buffer_size=2**12):
+    def __init__(
+        self,
+        env: Env,
+        model: nn.Module,
+        learning_rate=0.01,
+        discount_factor: float = 0.999,
+        alpha=0.03,
+        batch_size=128,
+        update_every=4,
+        max_buffer_size=2**12,
+    ):
         action_selector = EpsilonGreedy(env)
         estimator = MaxQEstimator(env, discount_factor)
         brain = QNetwork(model, estimator, learning_rate=learning_rate, alpha=alpha)

reinforcelab/agents/value_optimization/tabular/__init__.py → src/reinforcelab/modules/agents/value_optimization/tabular/__init__.py

@@ -1,3 +1,5 @@
 from .q_learning import QLearning
 from .sarsa import SARSA
 from .expected_sarsa import ExpectedSARSA
+
+__all__ = ["QLearning", "SARSA", "ExpectedSARSA"]

reinforcelab/agents/value_optimization/tabular/expected_sarsa.py → src/reinforcelab/modules/agents/value_optimization/tabular/expected_sarsa.py

@@ -1,10 +1,10 @@
 from gymnasium import Env
 
-from reinforcelab.agents.agent import Agent
-from reinforcelab.brains import QTable
-from reinforcelab.estimators import ExpectedSARSAEstimator
-from reinforcelab.action_selectors import EpsilonGreedy
-from reinforcelab.memory_buffers import OrderedBuffer
+from reinforcelab.modules.agents.agent import Agent
+from reinforcelab.modules.brains import QTable
+from reinforcelab.modules.estimators import ExpectedSARSAEstimator
+from reinforcelab.modules.action_selectors import EpsilonGreedy
+from reinforcelab.modules.memory_buffers import OrderedBuffer
 
 
 class ExpectedSARSA(Agent):

reinforcelab/agents/value_optimization/tabular/q_learning.py → src/reinforcelab/modules/agents/value_optimization/tabular/q_learning.py

@@ -1,10 +1,10 @@
 from gymnasium import Env
 
-from reinforcelab.agents.agent import Agent
-from reinforcelab.brains import QTable
-from reinforcelab.estimators import MaxQEstimator
-from reinforcelab.action_selectors import EpsilonGreedy
-from reinforcelab.memory_buffers import OrderedBuffer
+from reinforcelab.modules.agents.agent import Agent
+from reinforcelab.modules.brains import QTable
+from reinforcelab.modules.estimators import MaxQEstimator
+from reinforcelab.modules.action_selectors import EpsilonGreedy
+from reinforcelab.modules.memory_buffers import OrderedBuffer
 
 
 class QLearning(Agent):