update precommit config

.pre-commit-config.yaml

@@ -21,6 +21,6 @@ repos:
         language: node
         pass_filenames: false
         types: [python]
-        additional_dependencies: ["pyright"]
+        additional_dependencies: ["pyright@latest"]
         args:
           - --project=pyproject.toml

wingman/replay_buffer/core.py

@@ -206,7 +206,10 @@ def push(
 
         """
         self.base_buffer.push(
-            data=self.unwrap_data(data, bulk),
+            data=self.unwrap_data(
+                wrapped_data=data,
+                bulk=bulk,
+            ),
             bulk=bulk,
         )
 

wingman/replay_buffer/flat_replay_buffer.py

@@ -107,7 +107,7 @@ def _format_data(
             # cast to the right dtype
             data = np.asarray(
                 thing,
-                dtype=self.mode_dtype,  # pyright: ignore[reportGeneralTypeIssues]
+                dtype=self.mode_dtype,  # pyright: ignore[reportArgumentType, reportCallIssue]
             )
 
             # dim check
@@ -118,7 +118,7 @@ def _format_data(
             data = torch.asarray(
                 thing,
                 device=self.storage_device,
-                dtype=self.mode_dtype,  # pyright: ignore[reportGeneralTypeIssues]
+                dtype=self.mode_dtype,  # pyright: ignore[reportArgumentType]
             )
             data.requires_grad_(False)
 
@@ -182,7 +182,7 @@ def push(
                     [
                         self.mode_caller.zeros(
                             (self.mem_size, *item.shape),
-                            dtype=self.mode_dtype,  # pyright: ignore[reportGeneralTypeIssues]
+                            dtype=self.mode_dtype,  # pyright: ignore[reportArgumentType, reportCallIssue]
                         )
                         for item in array_data
                     ]
@@ -192,7 +192,7 @@ def push(
                     [
                         self.mode_caller.zeros(
                             (self.mem_size, *item.shape[1:]),
-                            dtype=self.mode_dtype,  # pyright: ignore[reportGeneralTypeIssues]
+                            dtype=self.mode_dtype,  # pyright: ignore[reportArgumentType, reportCallIssue]
                         )
                         for item in array_data
                     ]

wingman/utils.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 import numpy as np
 
 try:
@@ -17,37 +19,75 @@
 
 
 def gpuize(
-    input,
+    input: np.ndarray | torch.Tensor,
     device: str | torch.device = __device,
     dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
     """gpuize.
 
     Args:
     ----
-        input: the array that we want to gpuize
-        device: a string of the device we want to move the thing to
-        dtype: the datatype that the returned tensor should be
+        input (np.ndarray | torch.Tensor): the array that we want to gpuize
+        device (str | torch.device): a string of the device we want to move the thing to
+        dtype (torch.dtype): the datatype that the returned tensor should be
 
     """
     if torch.is_tensor(input):
-        return input.to(device=device, dtype=dtype)
+        return input.to(device=device, dtype=dtype)  # pyright: ignore[reportAttributeAccessIssue]
     else:
         return torch.tensor(input, device=device, dtype=dtype)
 
 
-def cpuize(input) -> np.ndarray:
+def nested_gpuize(
+    input: dict[str, Any],
+    device: str | torch.device = __device,
+    dtype: torch.dtype = torch.float32,
+) -> dict[str, Any]:
+    """Gpuize but for nested dictionaries of elements.
+
+    Args:
+    ----
+        input (dict[str, Any]): the array that we want to gpuize
+        device (str | torch.device): a string of the device we want to move the thing to
+        dtype (torch.dtype): the datatype that the returned tensor should be
+
+    """
+    for key, value in input.items():
+        if isinstance(value, dict):
+            input[key] = nested_gpuize(value, device=device, dtype=dtype)
+        else:
+            input[key] = gpuize(value)
+    return input
+
+
+def cpuize(input: np.ndarray | torch.Tensor) -> np.ndarray:
     """cpuize.
 
     Args:
     ----
-        input: the array of the thing we want to put on the cpu
+        input (np.ndarray | torch.Tensor): the array of the thing we want to put on the cpu
 
     """
     if torch.is_tensor(input):
-        return input.detach().cpu().numpy()
+        return input.detach().cpu().numpy()  # pyright: ignore[reportAttributeAccessIssue]
     else:
-        return input
+        return input  # pyright: ignore[reportReturnType]
+
+
+def nested_cpuize(input: dict[str, Any]) -> dict[str, Any]:
+    """Gpuize but for nested dictionaries of elements.
+
+    Args:
+    ----
+        input (dict[str, Any]): the array of the thing we want to put on the cpu
+
+    """
+    for key, value in input.items():
+        if isinstance(value, dict):
+            input[key] = nested_cpuize(value)
+        else:
+            input[key] = cpuize(value)
+    return input
 
 
 def shutdown_handler(*_):