wip

Author: ubgk

go2_simulation/simulation_node.py

@@ -1,4 +1,5 @@
 import rclpy
+import typing
 from rclpy.node import Node
 from sensor_msgs.msg import Image
 from unitree_go.msg import LowState, LowCmd
@@ -42,67 +43,8 @@ def euler_from_quaternion(quat_angle):
 
     return roll_x, pitch_y, yaw_z # in radians
 
-
-class Go2Simulation(Node):
+class Actor:
     def __init__(self):
-        super().__init__("go2_simulation")
-        simulator_name = self.declare_parameter("simulator", rclpy.Parameter.Type.STRING).value
-        simulator_name = "pybullet" if simulator_name is None else simulator_name
-
-        ########################### State publisher
-        self.lowstate_publisher = self.create_publisher(LowState, "/lowstate", 10)
-        self.odometry_publisher = self.create_publisher(Odometry, "/odometry/filtered", 10)
-        self.tf_broadcaster = TransformBroadcaster(self)
-        self.depth_publisher = self.create_publisher(Image, "/camera/depth", 10)
-        self.clock_publisher = self.create_publisher(Clock, "/clock", 10)
-
-        # Timer to publish periodically
-        self.high_level_period = 1.0 / 50  # seconds
-        self.low_level_sub_step = 24
-        self.timer = self.create_timer(self.high_level_period, self.update)
-
-        ########################## Camera
-        self.camera_period = 1.0 / 10 # seconds
-        self.camera_decimation = int(self.camera_period / self.high_level_period)
-        breakpoint()
-
-        ########################## Cmd listener
-        self.create_subscription(LowCmd, "/lowcmd", self.receive_cmd_cb, 10)
-        self.last_cmd_msg = LowCmd()
-
-        ########################## Simulator
-        self.get_logger().info("go2_simulator::loading simulator")
-        timestep = self.high_level_period / self.low_level_sub_step
-
-        self.simulator: AbstractSimulatorWrapper = None
-        if simulator_name == "simple":
-            from go2_simulation.simple_wrapper import SimpleWrapper
-
-            self.simulator = SimpleWrapper(self, timestep)
-        elif simulator_name == "pybullet":
-            from go2_simulation.bullet_wrapper import BulletWrapper
-
-            self.simulator = BulletWrapper(self, timestep)
-            self.bridge = CvBridge()
-        else:
-            self.get_logger().error("Simulation tool not recognized")
-
-        self.simulator_name = simulator_name
-        self.get_logger().info(f"go2_simulator::simulator {simulator_name} loaded")
-
-        ########################## Initial state
-        self.q_current = np.zeros(7 + 12)
-        self.v_current = np.zeros(6 + 12)
-        self.a_current = np.zeros(6 + 12)
-        self.f_current = np.zeros(4)
-
-        self.i = 0
-        self.sim_time = Time(seconds=0, nanoseconds=0)
-        self.time_delta = Duration(seconds=0, nanoseconds=int(self.high_level_period * 1e9))
-
-        self.init_onnx()
-
-    def init_onnx(self):
         onnx_path = "./models/wall.onnx"
         onnx_path = "/home/hamlet/Workspace/reinforcement-learning/inference/" + onnx_path
         self.onnx_session = rt.InferenceSession(onnx_path)
@@ -134,37 +76,35 @@ def init_onnx(self):
         self.step_counter = np.zeros((1,), dtype=np.float32)
 
         self.actions = np.zeros((1, 12), dtype=np.float32)
-
-
-    def forward(self, camera: bool = False):
-        if self.i == 0:
-            return np.zeros(12)
-
-        robot_id = self.simulator.robot
-
-        if camera:
-            im = self.simulator.get_camera_image().astype(np.float32)
+        self.policy_step = 0
+        
+        # LowCmd
+        self.lowcmd = LowCmd()
+
+    def forward(self, lowstate: LowState, im: typing.Optional[Image] = None):
+        if im:
+            im = np.array(im.data).reshape(im.height,im.width)
             self.vobs[:] = (im / 255.) - 0.5
             self.update_yaw[:] = 1.0
             self.update_depth[:] = 1.0
         else:
             self.update_yaw[:] = 0.0
             self.update_depth[:] = 0.0
 
-        contact_states = self.low_msg.foot_force > 20
+        contact_states = lowstate.foot_force > 20
 
-        quat = self.low_msg.imu_state.quaternion
+        quat = lowstate.imu_state.quaternion
         roll, pitch, yaw = euler_from_quaternion(quat)
         imu_obs = np.array([roll, pitch])
 
-        q = np.array([ms.q for ms in self.low_msg.motor_state])[:12]
+        q = np.array([ms.q for ms in lowstate.motor_state])[:12]
         q -= self.q0
 
         self.joint_vel[:] = (q - self.joint_pos) * 50. 
         self.joint_pos[:] = q
 
         obs_data = [
-            1 * self.low_msg.imu_state.gyroscope * 0.25, # 3
+            1 * lowstate.imu_state.gyroscope * 0.25, # 3
             1 * imu_obs, # 2
             [0.0],
             1 * self.yaws.squeeze(),
@@ -181,7 +121,8 @@ def forward(self, camera: bool = False):
             np.concatenate(obs_data).reshape(1, 53).astype(np.float32)
         )
         self.obs[:] = clip(self.obs)
-        self.step_counter[:] = self.i - 1
+        self.step_counter[:] = self.policy_step
+        self.policy_step += 1
 
         # Policy module
         inputs = {
@@ -208,6 +149,65 @@ def forward(self, camera: bool = False):
         return self.q0 + (np.clip(self.actions.squeeze(), -4.8, 4.8) * .25)
 
 
+class Go2Simulation(Node):
+    def __init__(self):
+        super().__init__("go2_simulation")
+        simulator_name = self.declare_parameter("simulator", rclpy.Parameter.Type.STRING).value
+        simulator_name = "pybullet" if simulator_name is None else simulator_name
+
+        ########################### State publisher
+        self.lowstate_publisher = self.create_publisher(LowState, "/lowstate", 10)
+        self.odometry_publisher = self.create_publisher(Odometry, "/odometry/filtered", 10)
+        self.tf_broadcaster = TransformBroadcaster(self)
+        self.depth_publisher = self.create_publisher(Image, "/camera/depth", 10)
+        self.clock_publisher = self.create_publisher(Clock, "/clock", 10)
+
+        # Timer to publish periodically
+        self.high_level_period = 1.0 / 50  # seconds
+        self.low_level_sub_step = 24
+        self.timer = self.create_timer(self.high_level_period, self.update)
+
+        ########################## Camera
+        self.camera_period = 1.0 / 10 # seconds
+        self.camera_decimation = int(self.camera_period / self.high_level_period)
+
+        ########################## Cmd listener
+        self.create_subscription(LowCmd, "/lowcmd", self.receive_cmd_cb, 10)
+        self.last_cmd_msg = LowCmd()
+
+        ########################## Simulator
+        self.get_logger().info("go2_simulator::loading simulator")
+        timestep = self.high_level_period / self.low_level_sub_step
+
+        self.simulator: AbstractSimulatorWrapper = None
+        if simulator_name == "simple":
+            from go2_simulation.simple_wrapper import SimpleWrapper
+
+            self.simulator = SimpleWrapper(self, timestep)
+        elif simulator_name == "pybullet":
+            from go2_simulation.bullet_wrapper import BulletWrapper
+
+            self.simulator = BulletWrapper(self, timestep)
+            self.bridge = CvBridge()
+        else:
+            self.get_logger().error("Simulation tool not recognized")
+
+        self.simulator_name = simulator_name
+        self.get_logger().info(f"go2_simulator::simulator {simulator_name} loaded")
+
+        ########################## Initial state
+        self.q_current = np.zeros(7 + 12)
+        self.v_current = np.zeros(6 + 12)
+        self.a_current = np.zeros(6 + 12)
+        self.f_current = np.zeros(4)
+
+        self.i = 0
+        self.sim_time = Time(seconds=0, nanoseconds=0)
+        self.time_delta = Duration(seconds=0, nanoseconds=int(self.high_level_period * 1e9))
+
+        self.actor = Actor()
+
+
     def update(self):
         ## Control robot
         if False:
@@ -218,10 +218,13 @@ def update(self):
             kd_des = np.array([self.last_cmd_msg.motor_cmd[i].kd for i in range(12)])
         else:
             # Camera update
-            if self.i % self.camera_decimation == 0:
-                q_des = self.forward(camera=True)
+            if self.i == 0:
+                q_des = np.zeros(12)
+            elif self.i % self.camera_decimation == 0:
+                im = self.camera_update() 
+                q_des = self.actor.forward(self.low_msg, im=im)
             else:
-                q_des = self.forward(camera=False)
+                q_des = self.actor.forward(self.low_msg, im=None)
 
             v_des = np.zeros(12)
             tau_des = np.zeros(12)
@@ -330,12 +333,11 @@ def update(self):
     def camera_update(self):
         if self.simulator_name == "pybullet":
             im = self.simulator.get_camera_image()
-        else:
-            self.get_logger().warn(f"Camera not implemented for this simulator: {self.simulator_name}")
-
-        if im is not None:
             img_msg = self.bridge.cv2_to_imgmsg(im, encoding="mono8")
             self.depth_publisher.publish(img_msg)
+            return img_msg
+        else:
+            self.get_logger().warn(f"Camera not implemented for this simulator: {self.simulator_name}")
 
     def receive_cmd_cb(self, msg):
         self.last_cmd_msg = msg