fixed unit tests

Author: pbontrager

apps/toy_rl/main.py

@@ -16,7 +16,7 @@
 import torch
 from forge.actors.collector import Collector
 
-from forge.data.replay_buffer import ReplayBuffer
+from forge.actors.replay_buffer import ReplayBuffer
 from forge.interfaces import Environment, Policy
 from forge.types import Action, Observation, State
 from monarch.actor import endpoint, proc_mesh
@@ -255,7 +255,7 @@ async def replay_buffer_sampler_task():
                             )
 
                             print(
-                                f"  Step {i+1:2d}: State={state_value:6.2f} → Action={action_value:6.2f}"
+                                f"  Step {i + 1:2d}: State={state_value:6.2f} → Action={action_value:6.2f}"
                             )
 
                     if idx < len(trajectories):  # Add spacing between trajectories

src/forge/actors/collector.py

@@ -12,14 +12,14 @@
 
 from typing import Callable
 
-from forge.data.replay_buffer import ReplayBuffer
+from monarch.actor import Actor, endpoint
+
+from forge.actors.replay_buffer import ReplayBuffer
 
 from forge.interfaces import Policy
 
 from forge.types import Trajectory
 
-from monarch.actor import Actor, endpoint
-
 
 class Collector(Actor):
     """Collects trajectories for the training loop."""

tests/unit_tests/rl/test_toy_rl.py

@@ -20,7 +20,7 @@
 
 from apps.toy_rl.main import ToyAction, ToyEnvironment, ToyObservation, ToyPolicy
 from forge.actors.collector import Collector
-from forge.data.replay_buffer import ReplayBuffer
+from forge.actors.replay_buffer import ReplayBuffer
 
 # local_proc_mesh is an implementation of proc_mesh for
 # testing purposes. It lacks some features of the real proc_mesh
@@ -214,6 +214,7 @@ async def test_full_rl_pipeline_simulation(self):
             1,  # batch_size
             1,  # max_policy_age
         )
+        await replay_buffer.setup.call()
         collector = await proc.spawn(
             "collector",
             Collector,

tests/unit_tests/test_replay_buffer.py

@@ -8,7 +8,7 @@
 
 import pytest
 import pytest_asyncio
-from forge.data.replay_buffer import ReplayBuffer
+from forge.actors.replay_buffer import ReplayBuffer
 from forge.types import Trajectory
 
 from monarch.actor import proc_mesh
@@ -21,6 +21,7 @@ async def replay_buffer(self) -> ReplayBuffer:
         replay_buffer = await mesh.spawn(
             "replay_buffer", ReplayBuffer, batch_size=2, max_policy_age=1
         )
+        await replay_buffer.setup.call()
         return replay_buffer
 
     @pytest.mark.asyncio