Merge branch 'main' of github.com:Rooholla-KhorramBakht/Go2Py into main

Go2Py/control/walk_these_ways.py

@@ -131,8 +131,7 @@ class CommandInterface:
         self.stance_width_cmd = 0.25
 
     def setGaitType(self, gait_type):
-        assert gait_type in [key for key in self.gaits.keys()], f'The gain type should be in {
-            [key for key in self.gaits.keys()]}'
+        assert gait_type in [key for key in self.gaits.keys()], f'The gain type should be in {self.gaits.keys()}'
         self.gait = torch.tensor(self.gaits[gait_type])
 
     def get_command(self):

Go2Py/sim/mujoco.py

@@ -6,13 +6,26 @@ from Go2Py import ASSETS_PATH
 import os
 from scipy.spatial.transform import Rotation
 
+pnt = np.array([-0.2, 0, 0.05])
+lidar_angles = np.linspace(0.0, 2 * np.pi, 1024).reshape(-1, 1)
+x_vec = np.cos(lidar_angles)
+y_vec = np.sin(lidar_angles)
+z_vec = np.zeros_like(x_vec)
+vec = np.concatenate([x_vec, y_vec, z_vec], axis=1)
+nray = vec.shape[0]
+geomid = np.zeros(nray, np.int32)
+dist = np.zeros(nray, np.float64)
+
 
 class Go2Sim:
-    def __init__(self, render=True, dt=0.002):
+    def __init__(self, mode='lowlevel', render=True, dt=0.002, xml_path=None):
 
-        self.model = mujoco.MjModel.from_xml_path(
-            os.path.join(ASSETS_PATH, 'mujoco/go2.xml')
-        )
+        if xml_path is None:
+            self.model = mujoco.MjModel.from_xml_path(
+                os.path.join(ASSETS_PATH, 'mujoco/go2.xml')
+            )
+        else:
+            self.model = mujoco.MjModel.from_xml_path(xml_path)
         self.simulated = True
         self.data = mujoco.MjData(self.model)
         self.dt = dt
@@ -68,6 +81,25 @@ class Go2Sim:
         self.kv = np.zeros(12)
         self.latest_command_stamp = time.time()
         self.actuator_tau = np.zeros(12)
+        self.mode = mode
+        if self.mode == 'highlevel':
+            from Go2Py.control.walk_these_ways import CommandInterface, loadParameters, Policy, WalkTheseWaysAgent, HistoryWrapper
+            checkpoint_path = os.path.join(ASSETS_PATH,'checkpoints/walk_these_ways')
+            self.cfg = loadParameters(checkpoint_path)
+            self.policy = Policy(checkpoint_path)
+            self.command_profile = CommandInterface()
+            self.agent = WalkTheseWaysAgent(self.cfg, self.command_profile, robot=self)
+            self.agent = HistoryWrapper(self.agent)
+            self.control_dt = self.cfg["control"]["decimation"] * self.cfg["sim"]["dt"]
+            self.obs = self.agent.reset()
+            self.standUpReset()
+            self.step_counter = 0
+            self.step = self.stepHighlevel
+            self.ex_sum=0
+            self.ey_sum=0
+            self.e_omega_sum=0
+        else:
+            self.step = self.stepLowlevel
 
     def reset(self):
         self.q_nominal = np.hstack(
@@ -75,6 +107,9 @@ class Go2Sim:
         )
         self.data.qpos = self.q_nominal
         self.data.qvel = np.zeros(18)
+        self.ex_sum=0
+        self.ey_sum=0
+        self.e_omega_sum=0
 
     def standUpReset(self):
         self.q0 = self.standing_q
@@ -117,7 +152,7 @@ class Go2Sim:
         self.tau_ff = tau_ff
         self.latest_command_stamp = time.time()
 
-    def step(self):
+    def stepLowlevel(self):
         state = self.getJointStates()
         q, dq = state['q'], state['dq']
         tau = np.diag(self.kp) @ (self.q_des - q).reshape(12, 1) + \
@@ -131,6 +166,30 @@ class Go2Sim:
         if self.render and (self.step_counter % self.render_ds_ratio) == 0:
             self.viewer.sync()
 
+    def stepHighlevel(self, vx, vy, omega_z, body_z_offset=0, step_height = 0.08, kp=[2, 0.5, 0.5], ki=[0.02, 0.01, 0.01]):
+        policy_info = {}
+        if self.step_counter % (self.control_dt // self.dt) == 0:
+            action = self.policy(self.obs, policy_info)
+            self.obs, ret, done, info = self.agent.step(action)
+        #Body velocity tracker PI controller
+        _, q = self.getPose()
+        world_R_body = Rotation.from_quat([q[1], q[2], q[3], q[0]]).as_matrix()
+        body_v = world_R_body.T@self.data.qvel[0:3].reshape(3,1)
+        ex = (vx-body_v[0])
+        ey = (vy-body_v[1])
+        e_omega = (omega_z-self.data.qvel[5])
+        self.ex_sum+=ex
+        self.ey_sum+=ey
+        self.e_omega_sum+=e_omega
+        self.command_profile.yaw_vel_cmd = np.clip(kp[2]*e_omega+ki[2]*self.e_omega_sum + omega_z, -2*np.pi, 2*np.pi)
+        self.command_profile.x_vel_cmd = np.clip(kp[0]*ex+ki[0]*self.ex_sum + vx, -2.5, 2.5)
+        self.command_profile.y_vel_cmd = np.clip(kp[1]*ey+ki[1]*self.ey_sum + vy,-1.5, 1.5)
+        self.command_profile.body_height_cmd = body_z_offset
+        self.command_profile.footswing_height_cmd = step_height
+
+        self.step_counter+=1
+        self.stepLowlevel()
+
     def getSiteJacobian(self, site_name):
         id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, site_name)
         assert id > 0, 'The requested site could not be found'
@@ -151,6 +210,29 @@ class Go2Sim:
         g_in_body = R.T @ np.array([0.0, 0.0, -1.0]).reshape(3, 1)
         return g_in_body
 
+    def getLaserScan(self, max_range=30):
+        t, q = self.getPose()
+        world_R_body = Rotation.from_quat([q[1], q[2], q[3], q[0]]).as_matrix()
+        pnt = t.copy()
+        pnt[2]+=0.25
+        vec_in_w = (world_R_body@vec.T).T
+        mujoco.mj_multiRay(
+            m=self.model,
+            d=self.data,
+            pnt=pnt,
+            vec=vec_in_w.flatten(),
+            geomgroup=None,
+            flg_static=1,
+            bodyexclude=-1,
+            geomid=geomid,
+            dist=dist,
+            nray=nray,
+            cutoff=max_range#mujoco.mjMAXVAL,
+        )
+        pcd = dist.reshape(-1, 1) * vec
+        idx = np.where(np.logical_and(dist!=-1, dist<max_range))[0]
+        return {"pcd": pcd[idx,...], "geomid": geomid[idx,...], "dist": dist[idx,...]}
+
     def overheat(self):
         return False
 

Go2Py/utils/ros2.py

@@ -0,0 +1,90 @@
+import struct
+import threading
+import time
+import numpy as np
+import rclpy
+import tf2_ros
+from rclpy.node import Node
+from rclpy.qos import QoSProfile
+from rclpy.executors import MultiThreadedExecutor
+from geometry_msgs.msg import TransformStamped
+from scipy.spatial.transform import Rotation as R
+def ros2_init(args=None):
+    rclpy.init(args=args)
+
+
+def ros2_close():
+    rclpy.shutdown()
+
+class ROS2ExecutorManager:
+    """A class to manage the ROS2 executor. It allows to add nodes and start the executor in a separate thread."""
+    def __init__(self):
+        self.executor = MultiThreadedExecutor()
+        self.nodes = []
+        self.executor_thread = None
+
+    def add_node(self, node: Node):
+        """Add a new node to the executor."""
+        self.nodes.append(node)
+        self.executor.add_node(node)
+
+    def _run_executor(self):
+        try:
+            self.executor.spin()
+        except KeyboardInterrupt:
+            pass
+        finally:
+            self.terminate()
+
+    def start(self):
+        """Start spinning the nodes in a separate thread."""
+        self.executor_thread = threading.Thread(target=self._run_executor)
+        self.executor_thread.start()
+
+    def terminate(self):
+        """Terminate all nodes and shutdown rclpy."""
+        for node in self.nodes:
+            node.destroy_node()
+        rclpy.shutdown()
+        if self.executor_thread:
+            self.executor_thread.join()
+
+class ROS2TFInterface(Node):
+
+    def __init__(self, parent_name, child_name, node_name):
+        super().__init__(f'{node_name}_tf2_listener')
+        self.parent_name = parent_name
+        self.child_name = child_name
+        self.tfBuffer = tf2_ros.Buffer()
+        self.listener = tf2_ros.TransformListener(self.tfBuffer, self)
+        self.T = None
+        self.stamp = None
+        self.running = True
+        self.thread = threading.Thread(target=self.update_loop)
+        self.thread.start()
+        self.trans = None
+
+    def update_loop(self):
+        while self.running:
+            try:
+                self.trans = self.tfBuffer.lookup_transform(self.parent_name, self.child_name, rclpy.time.Time(), rclpy.time.Duration(seconds=0.1))
+            except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
+                print(e)
+            time.sleep(0.01)    
+
+    def get_pose(self):
+        if self.trans is None:
+            return None
+        else:
+            translation = [self.trans.transform.translation.x, self.trans.transform.translation.y, self.trans.transform.translation.z]
+            rotation = [self.trans.transform.rotation.x, self.trans.transform.rotation.y, self.trans.transform.rotation.z, self.trans.transform.rotation.w]
+            self.T = np.eye(4)
+            self.T[0:3, 0:3] = R.from_quat(rotation).as_matrix()
+            self.T[:3, 3] = translation
+            self.stamp = self.trans.header.stamp.nanosec * 1e-9 + self.trans.header.stamp.sec
+            return self.T
+
+    def close(self):
+        self.running = False
+        self.thread.join()  
+        self.destroy_node()

deploy/launch_files/realsense.launch.py

@@ -0,0 +1,31 @@
+from launch import LaunchDescription
+from launch.actions import IncludeLaunchDescription
+from launch.launch_description_sources import PythonLaunchDescriptionSource
+from launch.substitutions import ThisLaunchFileDir
+from launch_ros.actions import Node
+
+
+def generate_launch_description():
+    return LaunchDescription([
+        Node(
+            name='go2_cam',
+            namespace='go2/cam',
+            package='realsense2_camera',
+            executable='realsense2_camera_node',
+            parameters=[{
+                    'enable_infra1': True,
+                    'enable_infra2': True,
+                    'enable_color': True,
+                    'enable_depth': False,
+                    'depth_module.emitter_enabled': 0,
+                    'rgb_camera.profile':'640x480x30',
+                    'depth_module.profile': '640x480x30',
+                    'enable_gyro': True,
+                    'enable_accel': True,
+                    'gyro_fps': 400,
+                    'accel_fps': 200,
+                    'unite_imu_method': 2,
+                    # 'tf_publish_rate': 0.0
+            }]
+        )
+    ])

examples/02-MuJoCo-sim.ipynb

@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Lowlevel Simulation"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -16,7 +23,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "robot = Go2Sim()\n",
+    "robot = Go2Sim(mode='lowlevel')\n",
     "robot.standUpReset()"
    ]
   },
@@ -44,6 +51,77 @@
     "    robot.setCommands(np.zeros(12), np.zeros(12), np.zeros(12), np.zeros(12), tau)\n",
     "    robot.step()"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Highlevel Simulation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from Go2Py.sim.mujoco import Go2Sim\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "robot = Go2Sim(mode='highlevel')\n",
+    "robot.standUpReset()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "robot.standUpReset\n",
+    "for i in range(10000):\n",
+    "    robot.step(0,0,0., step_height=0,kp=[2, 0.5, 0.5], ki=[0.02, 0.01, 0.01])\n",
+    "    time.sleep(robot.dt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "robot.getLaserScan()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import mujoco\n",
+    "mujoco.mjMAXVAL"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "lidar = robot.getLaserScan(max_range=3.)\n",
+    "idx = np.where(lidar['dist']!=-1)[0]\n",
+    "plt.plot(lidar['dist'][idx])"
+   ]
   }
  ],
  "metadata": {
@@ -62,7 +140,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.18"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,