Merge pull request #244 from RobotControlStack/juelg/fixes

A set of small fixes

Author: juelg

extensions/rcs_fr3/pyproject.toml

@@ -8,7 +8,7 @@ version = "0.5.2"
 description="RCS libfranka integration"
 dependencies = [
     "rcs>=0.5.2",
-    "frankik @ git+https://github.com/juelg/frankik",
+    "frankik",
 ]
 readme = "README.md"
 maintainers = [

extensions/rcs_fr3/src/rcs_fr3/creators.py

@@ -31,9 +31,9 @@
 
 
 class FrankIK(Kinematics):
-    def __init__(self, allow_elbow_flips: bool = False):
+    def __init__(self, global_solution: bool = False):
         Kinematics.__init__(self)
-        self.allow_elbow_flips = allow_elbow_flips
+        self.global_solution = global_solution
         self.kin = FrankaKinematics(robot_type="fr3")
 
     def forward(self, q0: np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]], tcp_offset: Pose) -> Pose:  # type: ignore
@@ -43,10 +43,7 @@ def forward(self, q0: np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]]
     def inverse(  # type: ignore
         self, pose: Pose, q0: np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]], tcp_offset: Pose
     ) -> np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]] | None:
-        tcp_offset = self.kin.FrankaHandTCPOffset
-        return self.kin.inverse(
-            pose.pose_matrix(), q0, tcp_offset.pose_matrix(), allow_elbow_flips=self.allow_elbow_flips
-        )
+        return self.kin.inverse(pose.pose_matrix(), q0, tcp_offset.pose_matrix(), global_solution=self.global_solution)
 
 
 # FYI: this needs to be in global namespace to avoid auto garbage collection issues

extensions/rcs_fr3/src/rcs_fr3/envs.py

@@ -30,8 +30,8 @@ def step(self, action: Any) -> tuple[dict[str, Any], SupportsFloat, bool, bool,
     def get_obs(self, obs: dict | None = None) -> dict[str, Any]:
         if obs is None:
             obs = dict(self.unwrapped.get_obs())
-        robot_state = cast(hw.FrankaState, self.unwrapped.robot.get_state())
-        obs["robot_state"] = vars(robot_state.robot_state)
+        # robot_state = cast(hw.FrankaState, self.unwrapped.robot.get_state())
+        # obs["robot_state"] = vars(robot_state.robot_state)
         return obs
 
     def reset(

extensions/rcs_tacto/src/rcs_tacto/tacto_wrapper.py

@@ -82,9 +82,9 @@ def reset(
         self.tacto_last_render = -1  # Reset last render time
         colors, depths = self.tacto_sensor.render(self.model, self.data)
         for site, color, depth in zip(self.tacto_sites, colors, depths, strict=False):
-            obs.setdefault("tacto", {}).setdefault(site, {}).setdefault("rgb", {})["data"] = color
+            obs.setdefault("frames", {}).setdefault(f"tactile_{site}", {}).setdefault("rgb", {})["data"] = color
             if self.enable_depth:
-                obs.setdefault("tacto", {}).setdefault(site, {}).setdefault("depth", {})["data"] = depth
+                obs.setdefault("frames", {}).setdefault(f"tactile_{site}", {}).setdefault("depth", {})["data"] = depth
         return obs, info
 
     def step(self, action: dict[str, Any]):
@@ -94,7 +94,9 @@ def step(self, action: dict[str, Any]):
             self.tacto_sensor.updateGUI(colors, depths) if self.visualize else None
             self.tacto_last_render = self.data.time
             for site, color, depth in zip(self.tacto_sites, colors, depths, strict=False):
-                obs.setdefault("tacto", {}).setdefault(site, {}).setdefault("rgb", {})["data"] = color
+                obs.setdefault("frames", {}).setdefault(f"tactile_{site}", {}).setdefault("rgb", {})["data"] = color
                 if self.enable_depth:
-                    obs.setdefault("tacto", {}).setdefault(site, {}).setdefault("depth", {})["data"] = depth
+                    obs.setdefault("frames", {}).setdefault(f"tactile_{site}", {}).setdefault("depth", {})[
+                        "data"
+                    ] = depth
         return obs, reward, done, truncated, info

python/rcs/envs/storage_wrapper.py

@@ -1,3 +1,4 @@
+import io
 import operator
 from concurrent.futures import ThreadPoolExecutor, wait
 from itertools import chain
@@ -10,6 +11,7 @@
 import pyarrow as pa
 import pyarrow.dataset as ds
 import simplejpeg
+from PIL import Image
 
 
 class StorageWrapper(gym.Wrapper):
@@ -22,7 +24,7 @@ def __init__(
         instruction: str,
         batch_size: int = 32,
         schema: Optional[pa.Schema] = None,
-        start_record: bool = False,
+        always_record: bool = False,
         basename_template: Optional[str] = None,
         max_rows_per_group: Optional[int] = None,
         max_rows_per_file: Optional[int] = None,
@@ -72,9 +74,10 @@ def __init__(
         self.max_rows_per_file = max_rows_per_file
         self.buffer: list[dict[str, Any]] = []
         self.step_cnt = 0
-        self._pause = True
+        self._pause = not always_record
+        self.always_record = always_record
         self.instruction = instruction
-        self._success = start_record
+        self._success = False
         self._prev_action = None
         self.thread_pool = ThreadPoolExecutor()
         self.queue: Queue[pa.Table | pa.RecordBatch] = Queue(maxsize=2)
@@ -122,6 +125,7 @@ def _flatten_arrays(self, d: dict[str, Any]):
         d.update(updates)
 
     def _encode_images(self, obs: dict[str, Any]):
+        # images
         _ = [
             *self.thread_pool.map(
                 lambda cam: operator.setitem(
@@ -133,6 +137,31 @@ def _encode_images(self, obs: dict[str, Any]):
             )
         ]
 
+        # depth
+        def to_tiff(depth_data):
+            img_bytes = io.BytesIO()
+            Image.fromarray(
+                depth_data.reshape((depth_data.shape[0], depth_data.shape[1])),
+            ).save(
+                img_bytes, format="TIFF"
+            )  # type: ignore
+            return img_bytes.getvalue()  # type: ignore
+
+        _ = [
+            *self.thread_pool.map(
+                lambda cam: (
+                    operator.setitem(
+                        obs["frames"][cam]["depth"],
+                        "data",
+                        to_tiff(obs["frames"][cam]["depth"]["data"]),
+                    )
+                    if "depth" in obs["frames"][cam]
+                    else None
+                ),
+                obs["frames"],
+            )
+        ]
+
     def step(self, action):
         # NOTE: expects the observation to be a dictionary
         if self._writer_future.done():
@@ -179,7 +208,7 @@ def start_record(self):
     def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None):
         if len(self.buffer) > 0:
             self._flush()
-        self._pause = True
+        self._pause = not self.always_record
         self._success = False
         self._prev_action = None
         obs, info = self.env.reset()