fix: Reset was calling States(States(...)) causing invalid shaped arr…

…ays)

src/jdrones/envs/position.py

@@ -51,18 +51,6 @@ def __init__(
             low=np.array([0, 0, 1]), high=np.array([10, 10, 10])
         )
 
-    def reset(
-        self,
-        *,
-        seed: int | None = None,
-        options: dict[str, Any] | None = None,
-    ) -> tuple[States, dict[str, Any]]:
-        super().reset(seed=seed, options=options)
-
-        obs, _ = self.env.reset(seed=seed, options=options)
-
-        return States([np.copy(obs)]), {}
-
     @staticmethod
     def get_reward(states: States) -> float:
         """
@@ -161,6 +149,18 @@ def _validate_action_input(
             )
         return action_as_state
 
+    def reset(
+        self,
+        *,
+        seed: int | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> tuple[States, dict[str, Any]]:
+        super().reset(seed=seed, options=options)
+
+        obs, _ = self.env.reset(seed=seed, options=options)
+
+        return States([np.copy(obs)]), {}
+
     def step(
         self, action: PositionAction | PositionVelocityAction
     ) -> tuple[States, float, bool, bool, dict[str, Any]]:
@@ -365,6 +365,18 @@ def __init__(
             low=np.array([[0, 0, 1], [0, 0, 1]]), high=10, shape=(2, 3)
         )
 
+    def reset(
+        self,
+        *,
+        seed: int | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> tuple[States, dict[str, Any]]:
+        super().reset(seed=seed, options=options)
+
+        obs, _ = self.env.reset(seed=seed, options=options)
+
+        return obs, {}
+
     @staticmethod
     def calc_v_at_B(A: VEC3, B: VEC3, C: VEC3, *, V: float, N: float = 3):
         """

tests/envs/position/test_stress.py

@@ -6,6 +6,7 @@
 def run(T, dt, env):
     trunc = False
     t, c = 0, 0
+    env.reset()
     while not trunc and t <= T:
         setpoint = env.action_space.sample()
         obs, _, term, trunc, _ = env.step(setpoint)