diff --git a/assume/reinforcement_learning/learning_role.py b/assume/reinforcement_learning/learning_role.py
index 36a88254..d372436b 100644
--- a/assume/reinforcement_learning/learning_role.py
+++ b/assume/reinforcement_learning/learning_role.py
@@ -145,7 +145,7 @@ def handle_message(self, content: dict, meta: dict) -> None:
         if content.get("type") == "replay_buffer":
             data = content["data"]
             self.buffer.add(
-                obs=data[0][0],
+                obs=data[0],
                 actions=data[1],
                 reward=data[2],
             )