add simple manual real world gym env example

examples/real_robot_example/gym_real_env/env.py

@@ -0,0 +1,158 @@
+import time
+
+import cv2
+import gymnasium as gym
+import numpy as np
+from gymnasium import spaces
+
+from .dynamixel import pos2pwm, pwm2pos
+from .robot import Robot
+
+FPS = 30
+
+CAMERAS_SHAPES = {
+    "observation.images.high": (480, 640, 3),
+    "observation.images.low": (480, 640, 3),
+}
+
+CAMERAS_PORTS = {
+    "observation.images.high": "/dev/video6",
+    "observation.images.low": "/dev/video0",
+}
+
+LEADER_PORT = "/dev/ttyACM1"
+FOLLOWER_PORT = "/dev/ttyACM0"
+
+
+def capture_image(cam, cam_width, cam_height):
+    # Capture a single frame
+    _, frame = cam.read()
+    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    # # Define your crop coordinates (top left corner and bottom right corner)
+    # x1, y1 = 400, 0  # Example starting coordinates (top left of the crop rectangle)
+    # x2, y2 = 1600, 900  # Example ending coordinates (bottom right of the crop rectangle)
+    # # Crop the image
+    # image = image[y1:y2, x1:x2]
+    # Resize the image
+    image = cv2.resize(image, (cam_width, cam_height), interpolation=cv2.INTER_AREA)
+
+    return image
+
+
+class RealEnv(gym.Env):
+    metadata = {}
+
+    def __init__(
+        self,
+        record: bool = False,
+        num_joints: int = 6,
+        cameras_shapes: dict = CAMERAS_SHAPES,
+        cameras_ports: dict = CAMERAS_PORTS,
+        follower_port: str = FOLLOWER_PORT,
+        leader_port: str = LEADER_PORT,
+        warmup_steps: int = 100,
+        trigger_torque=70,
+    ):
+        self.num_joints = num_joints
+        self.cameras_shapes = cameras_shapes
+        self.cameras_ports = cameras_ports
+        self.warmup_steps = warmup_steps
+        assert len(self.cameras_shapes) == len(self.cameras_ports), "Number of cameras and shapes must match."
+
+        self.follower_port = follower_port
+        self.leader_port = leader_port
+        self.record = record
+
+        # Initialize the robot
+        self.follower = Robot(device_name=self.follower_port)
+        if self.record:
+            self.leader = Robot(device_name=self.leader_port)
+            self.leader.set_trigger_torque(trigger_torque)
+
+        # Initialize the cameras - sorted by camera names
+        self.cameras = {}
+        for cn, p in sorted(self.cameras_ports.items()):
+            assert cn.startswith("observation.images."), "Camera names must start with 'observation.images.'."
+            self.cameras[cn] = cv2.VideoCapture(p)
+        if not all(c.isOpened() for c in self.cameras.values()):
+            raise OSError("Cannot open all camera ports.")
+
+        # Specify gym action and observation spaces
+        observation_space = {}
+
+        if self.num_joints > 0:
+            observation_space["agent_pos"] = spaces.Box(
+                low=-1000.0,
+                high=1000.0,
+                shape=(num_joints,),
+                dtype=np.float64,
+            )
+        if self.record:
+            observation_space["leader_pos"] = spaces.Box(
+                low=-1000.0,
+                high=1000.0,
+                shape=(num_joints,),
+                dtype=np.float64,
+            )
+
+        if self.cameras_shapes:
+            for cn, hwc_shape in self.cameras_shapes.items():
+                # Assumes images are unsigned int8 in [0,255]
+                observation_space[f"images.{cn}"] = spaces.Box(
+                    low=0,
+                    high=255,
+                    # height x width x channels (e.g. 480 x 640 x 3)
+                    shape=hwc_shape,
+                    dtype=np.uint8,
+                )
+
+        self.observation_space = spaces.Dict(observation_space)
+        self.action_space = spaces.Box(low=-1, high=1, shape=(num_joints,), dtype=np.float32)
+
+        self._observation = {}
+        self._terminated = False
+        self._action_time = time.time()
+
+    def _get_obs(self):
+        qpos = self.follower.read_position()
+        self._observation["agent_pos"] = pwm2pos(qpos)
+        for cn, c in self.cameras.items():
+            self._observation[f"images.{cn}"] = capture_image(
+                c, self.cameras_shapes[cn][1], self.cameras_shapes[cn][0]
+            )
+
+        if self.record:
+            leader_pos = self.leader.read_position()
+            self._observation["leader_pos"] = pwm2pos(leader_pos)
+
+    def reset(self, seed: int | None = None):
+        del seed
+        # Reset the robot and sync the leader and follower if we are recording
+        for _ in range(self.warmup_steps):
+            self._get_obs()
+            if self.record:
+                self.follower.set_goal_pos(pos2pwm(self._observation["leader_pos"]))
+        self._terminated = False
+        info = {}
+        return self._observation, info
+
+    def step(self, action: np.ndarray = None):
+        # Reset the observation
+        self._get_obs()
+        if self.record:
+            # Teleoperate the leader
+            self.follower.set_goal_pos(pos2pwm(self._observation["leader_pos"]))
+        else:
+            # Apply the action to the follower
+            self.follower.set_goal_pos(pos2pwm(action))
+        reward = 0
+        terminated = truncated = self._terminated
+        info = {}
+        return self._observation, reward, terminated, truncated, info
+
+    def render(self): ...
+
+    def close(self):
+        self.follower._disable_torque()
+        if self.record:
+            self.leader._disable_torque()