From a0a9f15d84f40645d68b17814b66f9ba0bd7a4b7 Mon Sep 17 00:00:00 2001 From: Duy Phuong Nguyen Date: Tue, 29 Apr 2025 11:48:40 +0800 Subject: [PATCH] bb iw pushed to separate branch - ignore train_result (left in local storage of Go2) --- .gitignore | 2 + SEP.py | 105 ++++++++++++++++++++ imports_test.py | 19 ++++ ros_viz.py | 49 ++++++++++ safe_walker.py | 220 +++++++++++++++++++++++++++++++++++++++++ safe_walker.py.save | 0 safe_walker_reader.py | 63 ++++++++++++ safe_walker_test.py | 157 ++++++++++++++++++++++++++++++ safety_enforcer_BB.py | 186 +++++++++++++++++++++++++++++++++++ safety_obs.py | 132 +++++++++++++++++++++++++ safety_test2.py | 121 +++++++++++++++++++++++ safety_test3.py | 210 +++++++++++++++++++++++++++++++++++++++ safety_tester.py | 114 ++++++++++++++++++++++ test.py | 221 ++++++++++++++++++++++++++++++++++++++++++ test2.py | 95 ++++++++++++++++++ vec_pub.py | 91 +++++++++++++++++ vec_pub_sem.py | 108 +++++++++++++++++++++ 17 files changed, 1893 insertions(+) create mode 100644 SEP.py create mode 100644 imports_test.py create mode 100644 ros_viz.py create mode 100644 safe_walker.py create mode 100644 safe_walker.py.save create mode 100644 safe_walker_reader.py create mode 100644 safe_walker_test.py create mode 100644 safety_enforcer_BB.py create mode 100644 safety_obs.py create mode 100644 safety_test2.py create mode 100644 safety_test3.py create mode 100644 safety_tester.py create mode 100644 test.py create mode 100644 test2.py create mode 100644 vec_pub.py create mode 100644 vec_pub_sem.py diff --git a/.gitignore b/.gitignore index 82f9275..cde8fec 100644 --- a/.gitignore +++ b/.gitignore @@ -160,3 +160,5 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + +/train_result diff --git a/SEP.py b/SEP.py new file mode 100644 index 0000000..51f80e6 --- /dev/null +++ b/SEP.py @@ -0,0 +1,105 @@ +from typing import Optional +import numpy as np +import os +import torch +from RARL.sac_adv import SAC_adv +from utils.utils import load_config +from RARL.sac_mini import SAC_mini + + + +class SafetyEnforcer: + + def __init__(self, + epsilon: float = 0.0, + imaginary_horizon: int = 100, + shield_type: Optional[str] = "value", + parent_dir: Optional[str] = "") -> None: + """_summary_ + + Args: + epsilon (float, optional): The epsilon value to be used for value shielding, determining the conservativeness of safety enforcer. Defaults to 0.0. + imaginary_horizon (int, optional): The horizon to be used for rollout-based shielding. Defaults to 100. + shield_type (Optional[str], optional): The shielding type to be used, choose from ["value", "rollout"]. Defaults to "value". + """ + #! TODO: Apply rollout-based shielding with the simulator + if shield_type != "value": + raise NotImplementedError + + self.epsilon = epsilon + self.imaginary_horizon = imaginary_horizon + + + training_dir = "train_result/success/" + load_dict = {"ctrl": 300_000, "dstb": 1_000_000} + + + model_path = os.path.join(parent_dir, training_dir, "model") + model_config_path = os.path.join(parent_dir, training_dir, + "config.yaml") + + config_file = os.path.join(parent_dir, model_config_path) + + if not os.path.exists(config_file): + raise ValueError( + "Cannot find config file for the model, terminated") + + config = load_config(config_file) + config_arch = config['arch'] + config_update = config['update'] + + self.policy = SAC_adv(config_update, config_arch) + self.policy.build_network(verbose=True) + print("Loading frozen weights of model at {} with load_dict {}".format( + model_path, load_dict)) + + self.policy.restore_refactor(None, model_path, load_dict=load_dict) + print("-> Done") + + self.critic = self.policy.adv_critic + self.dstb = self.policy.dstb + self.ctrl = self.policy.ctrl + + self.is_shielded = None + self.prev_q = None + + def get_action(self, state: np.ndarray, action: np.ndarray) -> np.ndarray: + assert len(state) == 78 + s_dstb = np.copy(state) + # s_dstb = np.concatenate((state, action), axis=0) + dstb = self.dstb(s_dstb) + + critic_q = max( + self.critic(torch.FloatTensor(state), torch.FloatTensor(action), + torch.FloatTensor(dstb))).detach().numpy() + + # positive is good + if critic_q < self.epsilon: + action = self.ctrl(state) + self.is_shielded = True + else: + self.is_shielded = False + + self.prev_q = critic_q.reshape(-1)[0] + + return action + + def get_q(self, state: np.ndarray, action: np.ndarray): + if state is not None and action is not None: + assert len(state) == 78 + s_dstb = np.copy(state) + # s_dstb = np.concatenate((state, action), axis=0) + dstb = self.dstb(s_dstb) + + critic_q = max( + self.critic(torch.FloatTensor(state), + torch.FloatTensor(action), + torch.FloatTensor(dstb))).detach().numpy() + + self.prev_q = critic_q.reshape(-1)[0] + + return self.prev_q + + + def get_shielding_status(self): + return self.is_shielded diff --git a/imports_test.py b/imports_test.py new file mode 100644 index 0000000..696e397 --- /dev/null +++ b/imports_test.py @@ -0,0 +1,19 @@ +import time +import sys +import math +import numpy as np +from collections import deque + +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import PointCloud2 + + +from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize +from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_ +from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ +from unitree_sdk2py.go2.sport.sport_client import SportClient + +from safety_enforcer import SafetyEnforcer + +print("Succesful imports") diff --git a/ros_viz.py b/ros_viz.py new file mode 100644 index 0000000..59762cd --- /dev/null +++ b/ros_viz.py @@ -0,0 +1,49 @@ +import rclpy +from rclpy.node import Node +from std_msgs.msg import Float32MultiArray +import numpy as np +import os +import time + +VIZ_PATH = "/tmp/state_vec_viz.npy" +FREQ_HZ = 10 # publish frequency + +class VecPublisher(Node): + def __init__(self): + super().__init__('vec77_publisher') + self.publisher_ = self.create_publisher(Float32MultiArray, '/safe_walker/viz_state', 10) + self.timer = self.create_timer(1.0 / FREQ_HZ, self.timer_callback) + self.prev_vec = None + + def timer_callback(self): + try: + vec = np.load(VIZ_PATH) + if vec.shape != (81,): # task(2), yaw (1), full_state (78) + self.get_logger().warn(f"Ignoring unexpected shape: {vec.shape}") + return + + if self.prev_vec is not None and np.array_equal(vec, self.prev_vec): + return # skip duplicate + + msg = Float32MultiArray() + msg.data = vec.astype(np.float32).tolist() + self.publisher_.publish(msg) + self.prev_vec = vec + + except Exception as e: + self.get_logger().warn(f"Failed to read vector: {e}") + + +def main(args=None): + rclpy.init(args=args) + node = VecPublisher() + try: + rclpy.spin(node) + except KeyboardInterrupt: + pass + node.destroy_node() + rclpy.shutdown() + + +if __name__ == '__main__': + main() diff --git a/safe_walker.py b/safe_walker.py new file mode 100644 index 0000000..8ec4c50 --- /dev/null +++ b/safe_walker.py @@ -0,0 +1,220 @@ +import time +import sys +import math +import numpy as np +from collections import deque + +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import PointCloud2 +from geometry_msgs.msg import PoseStamped +import sensor_msgs_py.point_cloud2 as pc2 + +from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize +from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_ +from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ +from unitree_sdk2py.go2.sport.sport_client import SportClient + +from SEP import SafetyEnforcer as safety_enforcer + +DECAY_TIME = 3.0 +ROBOT_RADIUS = 0.85 +INTENSITY_THRESHOLD = 220 + + +def yaw_from_quaternion(x, y, z, w): + siny_cosp = 2.0 * (w * z + x * y) + cosy_cosp = 1.0 - 2.0 * (y * y + z * z) + return np.arctan2(siny_cosp, cosy_cosp) + + +class SportModeSafeWalker(Node): + def __init__(self, client): + super().__init__('safe_walker_node') + + self.dt = 0.01 + self.t = 0 + self.yaw0 = 0 + self.prev_yaw = 0 + self.px0 = 0 + self.py0 = 0 + + self.lidar_max_range = 6.0 + self.obs_sequence_len = 2 + self.angle_sequence_len = 2 + self.obs_sequence = deque(maxlen=2) + self.angle_sequence = deque(maxlen=2) + self.latest_lidar_xy = np.zeros((18, 2), dtype=np.float32) + + self.buffer = [] + self.robot_pos = None + self.robot_yaw = 0.0 + + self.client = client + + self.safety_enforcer = safety_enforcer( + epsilon=0.5, + shield_type="value", + parent_dir="" + ) + + self.timer = self.create_timer(self.dt, self.timer_callback) + self.create_subscription(PointCloud2, "/utlidar/cloud", self.lidar_callback, 10) + self.create_subscription(PoseStamped, "/odom_pose", self.pose_callback, 10) + + self.robot_state = unitree_go_msg_dds__SportModeState_() + self.robot_state_ready = False + + def init_sdk(self): + self.client = SportClient() + self.client.SetTimeout(10.0) + self.client.Init() + + + def pose_callback(self, msg: PoseStamped): + pos = msg.pose.position + ori = msg.pose.orientation + self.robot_pos = np.array([pos.x, pos.y]) + self.robot_yaw = yaw_from_quaternion(ori.x, ori.y, ori.z, ori.w) + + def lidar_callback(self, msg: PointCloud2): + if self.robot_pos is None: + return + + now = self.get_clock().now().nanoseconds * 1e-9 + points = np.array([ + [p[0], p[1], p[3]] for p in pc2.read_points( + msg, field_names=("x", "y", "z", "intensity"), skip_nans=True) + ]) + + xy = points[:, :2] + intensity = points[:, 2] + + shifted = xy - self.robot_pos + c, s = np.cos(-self.robot_yaw), np.sin(-self.robot_yaw) + rot_matrix = np.array([[c, -s], [s, c]]) + aligned = shifted @ rot_matrix.T + aligned_points = np.hstack((aligned, intensity[:, None])) + + self.buffer.append((now, aligned_points)) + self.buffer = [(t, pts) for (t, pts) in self.buffer if now - t <= DECAY_TIME] + + all_points = np.vstack([pts for (_, pts) in self.buffer]) + xy = all_points[:, :2] + intensity = all_points[:, 2] + dists = np.linalg.norm(xy, axis=1) + + mask = (dists > ROBOT_RADIUS) | ((dists <= ROBOT_RADIUS) & (intensity > INTENSITY_THRESHOLD)) + filtered_points = xy[mask] + + if len(filtered_points) >= 18: + closest = np.argsort(np.linalg.norm(filtered_points, axis=1))[:18] + self.latest_lidar_xy = filtered_points[closest].astype(np.float32) + else: + self.latest_lidar_xy = np.zeros((18, 2), dtype=np.float32) + + def GetInitState(self, robot_state: SportModeState_): + self.px0 = robot_state.position[0] + self.py0 = robot_state.position[1] + self.yaw0 = robot_state.imu_state.rpy[2] + self.prev_yaw = self.yaw0 + print(self.yaw0, self.robot_yaw) + + def StandAndStart(self): + self.client.StandUp() + time.sleep(1) + self.client.BalanceStand() + print("Standing complete.") + + def encode_state(self, control: np.ndarray, lidar_obs: np.ndarray) -> np.ndarray: + lidar_xy = lidar_obs.reshape(18, 2) + dists = np.linalg.norm(lidar_xy, axis=1) + closest_idx = np.argmin(dists) + closest_vec = lidar_xy[closest_idx] + escape_vec = -closest_vec + escape_angle = np.arctan2(escape_vec[1], escape_vec[0]) + angle_diff = np.arctan2(math.sin(self.prev_yaw - escape_angle), math.cos(self.prev_yaw - escape_angle)) + + angle_obs = np.array([np.cos(angle_diff), np.sin(angle_diff)]) + self.angle_sequence.appendleft(angle_obs) + while len(self.angle_sequence) < self.angle_sequence_len: + self.angle_sequence.append(angle_obs) + angle_seq = np.concatenate(list(self.angle_sequence), axis=0) + + flat_lidar = lidar_obs.flatten() + self.obs_sequence.appendleft(flat_lidar) + while len(self.obs_sequence) < self.obs_sequence_len: + self.obs_sequence.append(flat_lidar) + lidar_seq = np.concatenate(list(self.obs_sequence), axis=0) + + return np.concatenate([control[:2], angle_seq, lidar_seq], axis=0) + + def safe_forward_control(self, robot_state: SportModeState_): + vx_desired = 0.5 + #No strafing + #vy_desired = 0.0 + + yaw = robot_state.imu_state.rpy[2] + yaw_error = math.atan2(math.sin(self.yaw0 - yaw), math.cos(self.yaw0 - yaw)) + + Kp_yaw = 2.0 + Kd_yaw = 0.5 + yaw_rate = (yaw - self.prev_yaw) / self.dt + vyaw_desired = Kp_yaw * yaw_error - Kd_yaw * yaw_rate + vyaw_desired = max(min(vyaw_desired, 1.0), -1.0) + self.prev_yaw = yaw + + ctrl_input = np.array([vx_desired, vyaw_desired], dtype=np.float32) + #raw_action = np.array([vx_desired, vy_desired, vyaw_desired], dtype=np.float32) + raw_action = ctrl_input + + state = self.encode_state(ctrl_input, self.latest_lidar_xy) + safe_action = self.safety_enforcer.get_action(state, raw_action) + + self.client.Move(*safe_action.tolist()) + print(f"[Shielded: {self.safety_enforcer.get_shielding_status()} | Q: {self.safety_enforcer.prev_q:.3f}]") + + def timer_callback(self): + if self.robot_state_ready: + self.t += self.dt + self.safe_forward_control(self.robot_state) + + +walker = None + +def HighStateHandler(msg: SportModeState_): + global walker + if walker is not None: + walker.robot_state = msg + walker.robot_state_ready = True + + +def main(args=None): + global walker + ChannelFactoryInitialize(0) + sub = ChannelSubscriber("rt/sportmodestate", SportModeState_) + sub.Init(HighStateHandler, 10) + + time.sleep(1) + print("Inside") + + client = SportClient() + client.Init() + + rclpy.init() + walker = SportModeSafeWalker(client) + + walker.GetInitState(walker.robot_state) + walker.StandAndStart() + try: + rclpy.spin(walker) + except KeyBoardInterrupt: + pass + + walker.destroy_node() + rclpy.shutdown() + + +if __name__ == '__main__': + main() + diff --git a/safe_walker.py.save b/safe_walker.py.save new file mode 100644 index 0000000..e69de29 diff --git a/safe_walker_reader.py b/safe_walker_reader.py new file mode 100644 index 0000000..424264f --- /dev/null +++ b/safe_walker_reader.py @@ -0,0 +1,63 @@ +import os +import time +import numpy as np +from unitree_sdk2py.core.channel import ChannelFactoryInitialize +from unitree_sdk2py.go2.sport.sport_client import SportClient +from SEP import SafetyEnforcer as safety_enforcer + +VEC_PATH = "/tmp/state_vec_77.npy" + +def try_load_vec77(): + try: + vec = np.load(VEC_PATH) + if vec.shape == (77,): + return vec + except Exception: + pass + return np.zeros((77,), dtype=np.float32) + +def main(): + ChannelFactoryInitialize(0) + client = SportClient() + client.SetTimeout(10.0) + client.Init() + client.StandUp() + time.sleep(1) + client.BalanceStand() + + sep = safety_enforcer(epsilon=0.5, shield_type="value", parent_dir="") + dt = 0.05 + + yaw0 = None + prev_yaw = None + + while True: + vec77 = try_load_vec77() + vec76, yaw = vec77[:-1], vec77[-1] + + # Initialize yaw offset + if yaw0 is None: + yaw0 = yaw + prev_yaw = yaw + + yaw_error = np.arctan2(np.sin(yaw0 - yaw), np.cos(yaw0 - yaw)) + yaw_rate = (yaw - prev_yaw) / dt + + Kp, Kd = 2.0, 0.5 + vyaw = Kp * yaw_error - Kd * yaw_rate + vyaw = np.clip(vyaw, -1.0, 1.0) + prev_yaw = yaw + + vx = 0.5 # constant forward speed + ctrl = np.array([vx, vyaw], dtype=np.float32) + + input_vec = np.concatenate([ctrl, vec76]) + safe_action = sep.get_action(input_vec, ctrl) + + client.Move(*safe_action.tolist()) + print(f"vx={safe_action[0]:.2f}, vyaw={safe_action[1]:.2f}, raw_yaw={yaw:.2f}, error={yaw_error:.2f}, Q={sep.prev_q:.3f}") + + time.sleep(dt) + +if __name__ == '__main__': + main() diff --git a/safe_walker_test.py b/safe_walker_test.py new file mode 100644 index 0000000..71b780d --- /dev/null +++ b/safe_walker_test.py @@ -0,0 +1,157 @@ +import time +import math +import numpy as np +from collections import deque + +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import PointCloud2 +from geometry_msgs.msg import PoseStamped +import sensor_msgs_py.point_cloud2 as pc2 + +from unitree_sdk2py.go2.sport.sport_client import SportClient +from SEP import SafetyEnforcer as safety_enforcer # your safety class +from unitree_sdk2py.core.channel import ChannelFactoryInitialize + +DECAY_TIME = 3.0 +ROBOT_RADIUS = 0.85 +INTENSITY_THRESHOLD = 220 + +def yaw_from_quaternion(x, y, z, w): + siny_cosp = 2.0 * (w * z + x * y) + cosy_cosp = 1.0 - 2.0 * (y * y + z * z) + return np.arctan2(siny_cosp, cosy_cosp) + +class SportModeSafeWalker(Node): + def __init__(self): + super().__init__('safe_walker_node') + + self.dt = 0.01 + self.t = 0 + + self.yaw0 = None + self.prev_yaw = 0.0 + + self.lidar_max_range = 6.0 + self.obs_sequence = deque(maxlen=2) + self.angle_sequence = deque(maxlen=2) + self.latest_lidar_xy = np.zeros((18, 2), dtype=np.float32) + + self.buffer = [] + self.robot_pos = None + self.robot_yaw = 0.0 + + self.client = SportClient() + self.client.SetTimeout(10.0) + self.client.Init() + + self.safety_enforcer = safety_enforcer( + epsilon=0.5, + shield_type="value", + parent_dir="" + ) + + self.create_subscription(PointCloud2, "/utlidar/cloud", self.lidar_callback, 10) + self.create_subscription(PoseStamped, "/odom_pose", self.pose_callback, 10) + self.timer = self.create_timer(self.dt, self.timer_callback) + + self.ready = False + + def pose_callback(self, msg: PoseStamped): + pos = msg.pose.position + ori = msg.pose.orientation + self.robot_pos = np.array([pos.x, pos.y]) + self.robot_yaw = yaw_from_quaternion(ori.x, ori.y, ori.z, ori.w) + + if self.yaw0 is None: + self.yaw0 = self.robot_yaw + self.prev_yaw = self.robot_yaw + self.client.StandUp() + time.sleep(1) + self.client.BalanceStand() + print("✅ Robot standing — starting walk.") + self.ready = True + + def lidar_callback(self, msg: PointCloud2): + if self.robot_pos is None: + return + + now = self.get_clock().now().nanoseconds * 1e-9 + points = np.array([ + [p[0], p[1], p[3]] for p in pc2.read_points( + msg, field_names=("x", "y", "z", "intensity"), skip_nans=True) + ]) + + xy = points[:, :2] + intensity = points[:, 2] + + shifted = xy - self.robot_pos + c, s = np.cos(-self.robot_yaw), np.sin(-self.robot_yaw) + aligned = shifted @ np.array([[c, -s], [s, c]]).T + aligned_points = np.hstack((aligned, intensity[:, None])) + + self.buffer.append((now, aligned_points)) + self.buffer = [(t, pts) for (t, pts) in self.buffer if now - t <= DECAY_TIME] + + all_points = np.vstack([pts for (_, pts) in self.buffer]) + xy = all_points[:, :2] + intensity = all_points[:, 2] + dists = np.linalg.norm(xy, axis=1) + + mask = (dists > ROBOT_RADIUS) | ((dists <= ROBOT_RADIUS) & (intensity > INTENSITY_THRESHOLD)) + filtered_points = xy[mask] + + if len(filtered_points) >= 18: + closest = np.argsort(np.linalg.norm(filtered_points, axis=1))[:18] + self.latest_lidar_xy = filtered_points[closest].astype(np.float32) + else: + self.latest_lidar_xy = np.zeros((18, 2), dtype=np.float32) + + def encode_state(self, control: np.ndarray, lidar_obs: np.ndarray) -> np.ndarray: + lidar_xy = lidar_obs.reshape(18, 2) + dists = np.linalg.norm(lidar_xy, axis=1) + closest_vec = lidar_xy[np.argmin(dists)] + escape_angle = np.arctan2(-closest_vec[1], -closest_vec[0]) + angle_diff = np.arctan2(math.sin(self.robot_yaw - escape_angle), math.cos(self.robot_yaw - escape_angle)) + + angle_obs = np.array([np.cos(angle_diff), np.sin(angle_diff)]) + self.angle_sequence.appendleft(angle_obs) + while len(self.angle_sequence) < 2: + self.angle_sequence.append(angle_obs) + angle_seq = np.concatenate(list(self.angle_sequence), axis=0) + + flat_lidar = lidar_obs.flatten() + self.obs_sequence.appendleft(flat_lidar) + while len(self.obs_sequence) < 2: + self.obs_sequence.append(flat_lidar) + lidar_seq = np.concatenate(list(self.obs_sequence), axis=0) + + return np.concatenate([control[:2], angle_seq, lidar_seq], axis=0) + + def timer_callback(self): + if not self.ready: + return + + vx = 0.5 + yaw_error = math.atan2(math.sin(self.yaw0 - self.robot_yaw), math.cos(self.yaw0 - self.robot_yaw)) + vyaw = 2.0 * yaw_error - 0.5 * ((self.robot_yaw - self.prev_yaw) / self.dt) + vyaw = max(min(vyaw, 1.0), -1.0) + self.prev_yaw = self.robot_yaw + + ctrl_input = np.array([vx, vyaw], dtype=np.float32) + state = self.encode_state(ctrl_input, self.latest_lidar_xy) + action = self.safety_enforcer.get_action(state, ctrl_input) + self.client.Move(*action.tolist()) + print(f"[vx: {action[0]:.2f}, vyaw: {action[1]:.2f}] | Q = {self.safety_enforcer.prev_q:.3f}") + + +def main(args=None): + ChannelFactoryInitialize(0) + rclpy.init(args=args) + walker = SportModeSafeWalker() + rclpy.spin(walker) + walker.destroy_node() + rclpy.shutdown() + +if __name__ == '__main__': + main() diff --git a/safety_enforcer_BB.py b/safety_enforcer_BB.py new file mode 100644 index 0000000..ddfdebe --- /dev/null +++ b/safety_enforcer_BB.py @@ -0,0 +1,186 @@ +from typing import Optional +import numpy as np +import os +import torch +from RARL.sac_adv import SAC_adv +from utils.utils import load_config +from RARL.sac_mini import SAC_mini + + +class SafetyEnforcerBB: + + def __init__(self, + epsilon: float = 0.0, + imaginary_horizon: int = 100, + shield_type: Optional[str] = "value", + parent_dir: Optional[str] = "") -> None: + """_summary_ + + Args: + epsilon (float, optional): The epsilon value to be used for value shielding, determining the conservativeness of safety enforcer. Defaults to 0.0. + imaginary_horizon (int, optional): The horizon to be used for rollout-based shielding. Defaults to 100. + shield_type (Optional[str], optional): The shielding type to be used, choose from ["value", "rollout"]. Defaults to "value". + """ + #! TODO: Apply rollout-based shielding with the simulator + if shield_type != "value": + raise NotImplementedError + + self.epsilon = epsilon + self.imaginary_horizon = imaginary_horizon + + # training_dir = "train_result/test_go2/test_isaacs_centerSampling" + + # training_dir = "train_result/test_go2/test_isaacs_centerSampling_withContact" + # load_dict = {"ctrl": 7_400_000, "dstb": 7_500_000} + + # training_dir = "train_result/test_go2/test_isaacs_postCoRL_arbitraryGx" + # load_dict = {"ctrl": 7_200_000, "dstb": 8_000_001} + # load_dict = {"ctrl": 100_000, "dstb": 100_000} + + # training_dir = "train_result/test_go2/go2_corldemo_pretrained" + # load_dict = {"ctrl": 12_000_001, "dstb": 12_000_001} + + # training_dir = "train_result/test_go2/go2_corldemo_2" + # load_dict = {"ctrl": 12_000_001, "dstb": 12_000_001} + + # training_dir = "train_result/test_go2/go2_corldemo_tgda_richURDF" + # load_dict = {"ctrl": 6_100_000, "dstb": 8_000_001} + + # training_dir = "train_result/test_go2/go2_corldemo_tgda" + # # load_dict = {"ctrl": 8_000_001, "dstb": 8_000_001} + # load_dict = {"ctrl": 6_500_000, "dstb": 8_000_001} + + training_dir = "train_result/smart_BB/" + load_dict = {"ctrl": 1_400_000, "dstb": 1_500_000} + print("Bryannnn") + # load_dict = {"ctrl": 2_800_000, "dstb": 5_500_000} + #load_dict = {"ctrl": 7_000_000, "dstb": 8_600_000} + # load_dict = {"ctrl": 8_800_000, "dstb": 8_600_000} + # load_dict = {"ctrl": 10_600_000, "dstb": 12_000_001} + + # SMART + # alternate + # training_dir = "train_result/smart/go2_isaacs" + # load_dict = {"ctrl": 2_100_000, "dstb": 2_100_000} + + # tgda + # training_dir = "train_result/smart/go2_tgda" + # load_dict = {"ctrl": 1_900_000, "dstb": 1_700_000} + + model_path = os.path.join(parent_dir, training_dir, "model") + model_config_path = os.path.join(parent_dir, training_dir, + "config.yaml") + + config_file = os.path.join(parent_dir, model_config_path) + + if not os.path.exists(config_file): + raise ValueError( + "Cannot find config file for the model, terminated") + + config = load_config(config_file) + config_arch = config['arch'] + config_update = config['update'] + + self.policy = SAC_adv(config_update, config_arch) + self.policy.build_network(verbose=True) + print("Loading frozen weights of model at {} with load_dict {}".format( + model_path, load_dict)) + + self.policy.restore_refactor(None, model_path, load_dict=load_dict) + print("-> Done") + + self.critic = self.policy.adv_critic + self.dstb = self.policy.dstb + self.ctrl = self.policy.ctrl + + self.is_shielded = None + self.prev_q = None + + def get_action(self, state: np.ndarray, action: np.ndarray) -> np.ndarray: + # state = np.concatenate((state[3:8], state[9:]), axis=0) + assert len(state) == 36 + s_dstb = np.copy(state) + # s_dstb = np.concatenate((state, action), axis=0) + dstb = self.dstb(s_dstb) + + critic_q = max( + self.critic(torch.FloatTensor(state), torch.FloatTensor(action), + torch.FloatTensor(dstb))).detach().numpy() + + # positive is good + if critic_q < self.epsilon: + action = self.ctrl(state) + self.is_shielded = True + else: + self.is_shielded = False + + self.prev_q = critic_q.reshape(-1)[0] + + return action + + def get_q(self, state: np.ndarray, action: np.ndarray): + if state is not None and action is not None: + assert len(state) == 36 + # state = np.concatenate((state[3:8], state[9:]), axis=0) + s_dstb = np.copy(state) + # s_dstb = np.concatenate((state, action), axis=0) + dstb = self.dstb(s_dstb) + + critic_q = max( + self.critic(torch.FloatTensor(state), + torch.FloatTensor(action), + torch.FloatTensor(dstb))).detach().numpy() + + self.prev_q = critic_q.reshape(-1)[0] + + return self.prev_q + + def target_margin(self, state): + """ (36) and 33D state, 32D state omits z + (x, y), z, + x_dot, y_dot, z_dot, + roll, pitch, (yaw) + w_x, w_y, w_z, + joint_pos x 12, + joint_vel x 12 + """ + # this is not the correct target margin, missing corner pos and toe pos, replacing corner pos with height, assuming that toes always touch ground + # l(x) < 0 --> x \in T + # state = np.concatenate((state[3:8], state[9:]), axis=0) + assert len(state) == 36 + return {"roll": 0.2 - abs(state[3]), "pitch": 0.2 - abs(state[4])} + # return { + # "body_ang_x": 0.17444 - abs(state[5]), + # "body_ang_y": 0.17444 - abs(state[6]), + # "body_ang_z": 0.17444 - abs(state[7]), + # "x_dot": 0.2 - abs(state[0]), + # "y_dot": 0.2 - abs(state[1]), + # "z_dot": 0.2 - abs(state[2]) + # } + + def get_safety_action(self, state, target=True, threshold=0.0): + assert len(state) == 36 + + stable_stance = np.array([ + -0.5, 0.7, -2.0, 0.5, 0.7, -2.0, -0.5, 0.7, -2.0, -0.5, 0.7, -2.0 + ]) + + if not target: + return self.ctrl(state) + else: + # switch between fallback and target stable stance, depending on the current state + margin = self.target_margin(state) + lx = min(margin.values()) + current_joint_pos = state[8:20] + + if lx > threshold: # account for sensor noise + # in target set, just output stable stance + #! TODO: enforce stable stance instead of just outputting zero changes to the current stance + return np.clip(stable_stance - current_joint_pos, + -np.ones(12) * 0.5, + np.ones(12) * 0.5) + else: + return self.ctrl(state) + + def get_shielding_status(self): + return self.is_shielded diff --git a/safety_obs.py b/safety_obs.py new file mode 100644 index 0000000..0043e6a --- /dev/null +++ b/safety_obs.py @@ -0,0 +1,132 @@ +import time +import sys +import math +import numpy as np +from collections import deque +import os + +from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize +from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_ +from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ +from unitree_sdk2py.go2.sport.sport_client import SportClient + +from SEP import SafetyEnforcer as safety_enforcer + +DECAY_TIME = 3.0 +ROBOT_RADIUS = 0.85 +INTENSITY_THRESHOLD = 220 +VEC_PATH = "/tmp/state_vec_76.npy" +OBS_PATH = "/tmp/obs.npy" + +def try_load_vec76(): + try: + vec = np.load(VEC_PATH) + if vec.shape == (76,): + return vec + except Exception: + pass + return np.zeros((76,), dtype=np.float32) + + +class SportModeSafeWalker: + def __init__(self): + self.dt = 0.02 + self.t = 0 + self.yaw0 = 0 + self.prev_yaw = 0 + + + self.client = SportClient() + self.client.SetTimeout(10.0) + self.client.Init() + + self.safety_enforcer = safety_enforcer( + epsilon=0.5, + shield_type="value", + parent_dir="" + ) + + self.robot_state = unitree_go_msg_dds__SportModeState_() + self.robot_state_ready = False + + def GetInitState(self, robot_state: SportModeState_): + self.px0 = robot_state.position[0] + self.py0 = robot_state.position[1] + self.yaw0 = robot_state.imu_state.rpy[2] + self.prev_yaw = self.yaw0 + print(self.yaw0) + + def StandAndStart(self): + self.client.StandUp() + time.sleep(1) + self.client.BalanceStand() + print("Standing complete.") + + def safe_forward_control(self, robot_state: SportModeState_): + vx_desired = 0.5 # forward velocity + yaw = robot_state.imu_state.rpy[2] + yaw_error = math.atan2(math.sin(self.yaw0 - yaw), math.cos(self.yaw0 - yaw)) + + Kp_yaw = 2.0 + Kd_yaw = 0.5 + yaw_rate = (yaw - self.prev_yaw) / self.dt + vyaw_desired = Kp_yaw * yaw_error - Kd_yaw * yaw_rate + vyaw_desired = max(min(vyaw_desired, 1.0), -1.0) + self.prev_yaw = yaw + + ctrl_input = np.array([vx_desired, vyaw_desired], dtype=np.float32) + raw_action = ctrl_input + + state_vec76 = try_load_vec76() + full_state = np.concatenate([ctrl_input, state_vec76]) + temp_path = OBS_PATH + ".tmp" + with open(temp_path, 'wb') as f: + np.save(f, full_state) + f.flush() + os.fsync(f.fileno()) + os.replace(temp_path, OBS_PATH) + print("obs sent") + safe_action = self.safety_enforcer.get_action(full_state, raw_action) + assert len(safe_action) == 2 + safe_action = np.array([safe_action[0], 0.0, safe_action[1]], dtype=np.float32) + + #Robot has a separate vy for strating but currently unused + + self.client.Move(*safe_action.tolist()) + print(f"[Shielded: {self.safety_enforcer.get_shielding_status()} | Q: {self.safety_enforcer.prev_q:.3f}]") + + def update(self): + if self.robot_state_ready: + self.t += self.dt + self.safe_forward_control(self.robot_state) + +walker = None + +def HighStateHandler(msg: SportModeState_): + global walker + if walker is not None: + walker.robot_state = msg + walker.robot_state_ready = True + +def main(args=None): + global walker + + if len(sys.argv) > 1: + ChannelFactoryInitialize(0, sys.argv[1]) + else: + ChannelFactoryInitialize(0) + + sub = ChannelSubscriber("rt/sportmodestate", SportModeState_) + sub.Init(HighStateHandler, 10) + time.sleep(1) + + walker = SportModeSafeWalker() + walker.GetInitState(walker.robot_state) + walker.StandAndStart() + + while True: + walker.update() + time.sleep(walker.dt) + +if __name__ == '__main__': + main() diff --git a/safety_test2.py b/safety_test2.py new file mode 100644 index 0000000..779ff91 --- /dev/null +++ b/safety_test2.py @@ -0,0 +1,121 @@ +import time +import sys +import math +import numpy as np +from collections import deque + +from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize +from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_ +from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ +from unitree_sdk2py.go2.sport.sport_client import SportClient + +from SEP import SafetyEnforcer as safety_enforcer + +DECAY_TIME = 4.0 +ROBOT_RADIUS = 0.85 +INTENSITY_THRESHOLD = 220 +VEC_PATH = "/tmp/state_vec_76.npy" + +def try_load_vec76(): + try: + vec = np.load(VEC_PATH) + if vec.shape == (76,): + return vec + except Exception: + pass + return np.zeros((76,), dtype=np.float32) + + +class SportModeSafeWalker: + def __init__(self): + self.dt = 0.001 + self.t = 0 + self.yaw0 = 0 + self.prev_yaw = 0 + + + self.client = SportClient() + self.client.SetTimeout(10.0) + self.client.Init() + + self.safety_enforcer = safety_enforcer( + epsilon=1.4, + shield_type="value", + parent_dir="" + ) + + + def GetInitState(self, robot_state: SportModeState_): + self.px0 = robot_state.position[0] + self.py0 = robot_state.position[1] + self.yaw0 = robot_state.imu_state.rpy[2] + self.prev_yaw = self.yaw0 + print(self.yaw0,"Init") + + def StandAndStart(self): + self.client.StandUp() + time.sleep(1) + self.client.BalanceStand() + print("Standing complete.") + + def safe_forward_control(self, robot_state: SportModeState_): + vx_desired = 0.15 # forward velocity + yaw = robot_state.imu_state.rpy[2] + yaw_error = math.atan2(math.sin(self.yaw0 - yaw), math.cos(self.yaw0 - yaw)) + + Kp_yaw = 2.0 + Kd_yaw = 0.5 + yaw_rate = (yaw - self.prev_yaw) / self.dt + vyaw_desired = Kp_yaw * yaw_error - Kd_yaw * yaw_rate + vyaw_desired = max(min(vyaw_desired, 1.0), -1.0) + self.prev_yaw = yaw + + ctrl_input = np.array([vx_desired, vyaw_desired], dtype=np.float32) + raw_action = ctrl_input + + state_vec76 = try_load_vec76() + full_state = np.concatenate([ctrl_input, state_vec76]) + safe_action = self.safety_enforcer.get_action(full_state, raw_action) + assert len(safe_action) == 2 + safe_action = np.array([safe_action[0], 0.0, safe_action[1]], dtype=np.float32) + + #Robot has a separate vy for strating but currently unused + print(yaw, "yaw", yaw_error, "error", yaw_rate, "rate", vyaw_desired,"vyaw") + self.client.Move(*safe_action.tolist()) + if self.safety_enforcer.get_shielding_status(): + print(safe_action, "safe_action") + print(f"[Shielded: {self.safety_enforcer.get_shielding_status()} | Q: {self.safety_enforcer.prev_q:.3f}]") + + def update(self, robot_state:SportModeState_): + self.t += self.dt + self.safe_forward_control(robot_state) + +# === Robot state subscription === +robot_state = unitree_go_msg_dds__SportModeState_() + +def HighStateHandler(msg: SportModeState_): + global robot_state + robot_state = msg + +def main(args=None): + + + if len(sys.argv) > 1: + ChannelFactoryInitialize(0, sys.argv[1]) + else: + ChannelFactoryInitialize(0) + + sub = ChannelSubscriber("rt/sportmodestate", SportModeState_) + sub.Init(HighStateHandler, 10) + time.sleep(1) + + walker = SportModeSafeWalker() + walker.GetInitState(robot_state) + walker.StandAndStart() + + while True: + walker.update(robot_state) + time.sleep(walker.dt) + +if __name__ == '__main__': + main() diff --git a/safety_test3.py b/safety_test3.py new file mode 100644 index 0000000..bec610e --- /dev/null +++ b/safety_test3.py @@ -0,0 +1,210 @@ +import time +import sys +import math +import numpy as np +from collections import deque +import os + +from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize +from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_ +from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ +from unitree_sdk2py.go2.sport.sport_client import SportClient + +from SEP import SafetyEnforcer as safety_enforcer + +DECAY_TIME = 4.0 +ROBOT_RADIUS = 0.8 +INTENSITY_THRESHOLD = 220 +VEC_PATH = "/tmp/state_vec_72.npy" +VIZ_PATH = "/tmp/state_vec_viz.npy" + +def try_load_vec72(): + try: + vec = np.load(VEC_PATH) + if vec.shape == (72,): + return vec + except Exception: + pass + return np.zeros((72,), dtype=np.float32) + + +def scale_and_shift(x, old_range, new_range): + ratio = (new_range[1] - new_range[0]) / (old_range[1] - old_range[0]) + x_new = (x - old_range[0]) * ratio + new_range[0] + return x_new + +def get_control(action): + # clip before scaling + if action[-1] == 1: + xdot_range_model = [0.2, 0.1] + xdot_range = [.2, .1] + else: + xdot_range_model = [.2, .5] + xdot_range = [0.2, 0.5] + action = np.clip(action, -1, 1) + v = scale_and_shift(action[0], [-1, 1], xdot_range_model) + v = np.clip(v, xdot_range[0], xdot_range[1]) + w = np.clip(action[1], -1, 1) + return v, w + +def get_safe_extreme(lidar, yaw_value) -> float: + """ + Computes the escape angle pointing directly away from the nearest obstacle. + + Returns: + float: Angle difference between current heading and escape direction ∈ (–π, π]. + """ + # Get LiDAR readings (relative XY coordinates of 18 closest obstacles) + + lidar_xy = lidar.reshape(18, 2) + + # Find the closest obstacle + dists = np.linalg.norm(lidar_xy, axis=1) + closest_idx = np.argmin(dists) + closest_vec = lidar_xy[closest_idx] + + # Compute angle of closest vector + angle_safe = np.arctan2(closest_vec[1], closest_vec[0]) + # Return signed angle difference ∈ (–π, π] + clearance = np.arctan2(np.sin(yaw_value - angle_safe), np.cos(yaw_value - angle_safe)) + + return np.array([np.cos(clearance), np.sin(clearance)]) + +class SportModeSafeWalker: + def __init__(self): + self.dt = 1 + self.t = 0 + self.yaw0 = 0 + self.prev_yaw = 0 + self.curr_heading = None + self.prev_heading = None + + + self.client = SportClient() + self.client.SetTimeout(10.0) + self.client.Init() + self.ctrl_list = deque(maxlen=1) + self.ctrl_list.append((0.0, 0.0)) # Append vx, vyaw + + self.safety_enforcer = safety_enforcer( + epsilon=1.0, + shield_type="value", + parent_dir="" + ) + + self.viz_valid = True + + + def GetInitState(self, robot_state: SportModeState_): + self.px0 = robot_state.position[0] + self.py0 = robot_state.position[1] + self.yaw0 = robot_state.imu_state.rpy[2] + self.prev_yaw = self.yaw0 + print(self.yaw0,"Init is set as offset") + self.offset = self.yaw0 + + def StandAndStart(self): + self.client.StandUp() + time.sleep(1) + self.client.BalanceStand() + print("Standing complete.") + + def safe_forward_control(self, robot_state: SportModeState_): + vx_desired = 0.00# forward velocity + yaw = robot_state.imu_state.rpy[2] + yaw_error = math.atan2(math.sin(self.yaw0 - yaw), math.cos(self.yaw0 - yaw)) + + Kp_yaw = 2.0 + Kd_yaw = 0.5 + yaw_rate = (yaw - self.prev_yaw) / self.dt + vyaw_desired = Kp_yaw * yaw_error - Kd_yaw * yaw_rate + vyaw_desired = 0.12 + self.prev_yaw = yaw + + # Unpack current command for observation + ctrl_input = np.array([*self.ctrl_list[0]], dtype=np.float32) + assert ctrl_input.shape == (2,), f"Expected shape (2,), got {ctrl_input.shape}" + + + # Set the desired velocity and yaw rate: next action + raw_action = np.array([vx_desired, vyaw_desired], dtype=np.float32) + + # Load lidar representation and nearest obstacle heading information + state_vec72 = try_load_vec72() + + self.prev_heading = self.curr_heading + calib_yaw = yaw - self.offset #calibrated angle helps us with clearance + self.curr_heading = get_safe_extreme(state_vec72[0:36], calib_yaw) + if self.prev_heading is None: self.prev_heading = self.curr_heading + + # Define full observation state as a concatenation of the current command, curent and past headings and the state vector + full_state = np.concatenate([ctrl_input, self.curr_heading, self.prev_heading, state_vec72]) + # assert 78 + # Monitoring, Intervention, and Fallback: Safety Filter with current observation and action + safe_action = self.safety_enforcer.get_action(full_state, raw_action) + assert len(safe_action) == 2 + safe_action = get_control(safe_action) + + # Imput vy information (0.0) for compatibility with the robot sdk + # Robot training happens in u,v frame u along right axis and v along forward axis both local to robot + # However sdk has x in forward direction and y in right ward direction + # (vx,vy,vyaw){sdk} --> (v, -u, w){robot training frame} + + + # Nevermind success model has correct alignmnet trains in y-direction so we can use first entry for vx + safe_action = (.8) * np.array([safe_action[0], 0.0, safe_action[1]], dtype=np.float32) + + # If visualization is enabled, save the state vector + # Use atomic file writing to avoid data corruption + if self.viz_valid: + viz_tmp_path = VIZ_PATH + ".tmp" + # Plotting in SF pipeline definitely needs calibrated yaw + viz_state = np.concatenate([np.array([calib_yaw]), raw_action, full_state], axis=0) + with open(viz_tmp_path, 'wb') as f: + np.save(f, viz_state) + f.flush() + os.fsync(f.fileno()) + os.replace(viz_tmp_path, VIZ_PATH) + + #Robot has a separate vy for strating but currently unused + print(yaw, "yaw", calib_yaw, "calib_yaw", yaw_error, "error", yaw_rate, "rate", vyaw_desired,"vyaw") + self.client.Move(*safe_action.tolist()) + print("safe_action", safe_action) + # Clear the previous command and append the current command for observation + self.ctrl_list.clear() + self.ctrl_list.append((safe_action[0], safe_action[2])) + print(f"[Shielded: {self.safety_enforcer.get_shielding_status()} | Q: {self.safety_enforcer.prev_q:.3f}]") + + def update(self, robot_state: SportModeState_): + self.t += self.dt + self.safe_forward_control(robot_state) + +# === Robot state subscription === +robot_state = unitree_go_msg_dds__SportModeState_() + +def HighStateHandler(msg: SportModeState_): + global robot_state + robot_state = msg + +def main(args=None): + + + if len(sys.argv) > 1: + ChannelFactoryInitialize(0, sys.argv[1]) + else: + ChannelFactoryInitialize(0) + + sub = ChannelSubscriber("rt/sportmodestate", SportModeState_) + sub.Init(HighStateHandler, 10) + time.sleep(1) + + walker = SportModeSafeWalker() + walker.GetInitState(robot_state) + walker.StandAndStart() + + while True: + walker.update(robot_state) + time.sleep(walker.dt) + if walker.safety_enforcer.prev_q <= .15: raise KeyboardInterrupt +if __name__ == '__main__': + main() diff --git a/safety_tester.py b/safety_tester.py new file mode 100644 index 0000000..79521f7 --- /dev/null +++ b/safety_tester.py @@ -0,0 +1,114 @@ +import time +import sys +import math +import numpy as np +from collections import deque + +from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize +from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_ +from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ +from unitree_sdk2py.go2.sport.sport_client import SportClient + +from SEP import SafetyEnforcer as safety_enforcer + +DECAY_TIME = 3.0 +ROBOT_RADIUS = 0.85 +INTENSITY_THRESHOLD = 220 +VEC_PATH = "/tmp/state_vec_76.npy" + +def try_load_vec76(): + try: + vec = np.load(VEC_PATH) + if vec.shape == (76,): + return vec + except Exception: + pass + return np.zeros((76,), dtype=np.float32) + + +class SportModeSafeWalker: + def __init__(self): + self.dt = 0.001 + self.t = 0 + + + self.client = SportClient() + self.client.SetTimeout(10.0) + self.client.Init() + + self.safety_enforcer = safety_enforcer( + epsilon=0.5, + shield_type="value", + parent_dir="" + ) + + self.robot_state = unitree_go_msg_dds__SportModeState_() + self.robot_state_ready = False + + def GetInitState(self, robot_state: SportModeState_): + self.px0 = robot_state.position[0] + self.py0 = robot_state.position[1] + self.yaw0 = robot_state.imu_state.rpy[2] + self.prev_yaw = self.yaw0 + print(self.yaw0, "Init") + + def StandAndStart(self): + self.client.StandUp() + time.sleep(1) + self.client.BalanceStand() + print("Standing complete.") + + def safe_forward_control(self, robot_state: SportModeState_): + vx_desired = 0.0 # forward velocity + + vyaw_desired = 0.0 + ctrl_input = np.array([vx_desired, vyaw_desired], dtype=np.float32) + raw_action = ctrl_input + + state_vec76 = try_load_vec76() + full_state = np.concatenate([ctrl_input, state_vec76]) + safe_action = self.safety_enforcer.get_action(full_state, raw_action) + assert len(safe_action) == 2 + safe_action = np.array([safe_action[0], 0.0, safe_action[1]], dtype=np.float32) + + #Robot has a separate vy for strating but currently unused + print(safe_action.tolist(), "list") + + self.client.Move(*safe_action.tolist()) + print(f"[Shielded: {self.safety_enforcer.get_shielding_status()} | Q: {self.safety_enforcer.prev_q:.3f}]") + + def update(self): + if self.robot_state_ready: + self.t += self.dt + self.safe_forward_control(self.robot_state) + +walker = None + +def HighStateHandler(msg: SportModeState_): + global walker + if walker is not None: + walker.robot_state = msg + walker.robot_state_ready = True + +def main(args=None): + global walker + + if len(sys.argv) > 1: + ChannelFactoryInitialize(0, sys.argv[1]) + else: + ChannelFactoryInitialize(0) + + sub = ChannelSubscriber("rt/sportmodestate", SportModeState_) + sub.Init(HighStateHandler, 10) + time.sleep(1) + + walker = SportModeSafeWalker() + walker.GetInitState(walker.robot_state) + walker.StandAndStart() + + while True: + walker.update() + time.sleep(walker.dt) + +if __name__ == '__main__': + main() diff --git a/test.py b/test.py new file mode 100644 index 0000000..b43cbab --- /dev/null +++ b/test.py @@ -0,0 +1,221 @@ +import time +import sys +import math +import numpy as np +from collections import deque + +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import PointCloud2 +from geometry_msgs.msg import PoseStamped +import sensor_msgs_py.point_cloud2 as pc2 + +from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize +from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_ +from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ +from unitree_sdk2py.go2.sport.sport_client import SportClient + +from SEP import SafetyEnforcer as safety_enforcer + +DECAY_TIME = 3.0 +ROBOT_RADIUS = 0.85 +INTENSITY_THRESHOLD = 220 + + +def yaw_from_quaternion(x, y, z, w): + siny_cosp = 2.0 * (w * z + x * y) + cosy_cosp = 1.0 - 2.0 * (y * y + z * z) + return np.arctan2(siny_cosp, cosy_cosp) + + +class SportModeSafeWalker(Node): + def __init__(self): + super().__init__('safe_walker_node') + + self.dt = 0.01 + self.t = 0 + self.yaw0 = 0 + self.prev_yaw = 0 + self.px0 = 0 + self.py0 = 0 + + self.lidar_max_range = 6.0 + self.obs_sequence_len = 2 + self.angle_sequence_len = 2 + self.obs_sequence = deque(maxlen=2) + self.angle_sequence = deque(maxlen=2) + self.latest_lidar_xy = np.zeros((18, 2), dtype=np.float32) + + self.buffer = [] + self.robot_pos = None + self.robot_yaw = 0.0 + + self.client = SportClient() + self.client.SetTimeout(10.0) + self.client.Init() + + self.safety_enforcer = safety_enforcer( + epsilon=0.5, + shield_type="value", + parent_dir="" + ) + + self.timer = self.create_timer(self.dt, self.timer_callback) + self.create_subscription(PointCloud2, "/utlidar/cloud", self.lidar_callback, 10) + self.create_subscription(PoseStamped, "/odom_pose", self.pose_callback, 10) + + self.robot_state = unitree_go_msg_dds__SportModeState_() + self.robot_state_ready = False + + def pose_callback(self, msg: PoseStamped): + pos = msg.pose.position + ori = msg.pose.orientation + self.robot_pos = np.array([pos.x, pos.y]) + self.robot_yaw = yaw_from_quaternion(ori.x, ori.y, ori.z, ori.w) + + def lidar_callback(self, msg: PointCloud2): + if self.robot_pos is None: + return + + now = self.get_clock().now().nanoseconds * 1e-9 + points = np.array([ + [p[0], p[1], p[3]] for p in pc2.read_points( + msg, field_names=("x", "y", "z", "intensity"), skip_nans=True) + ]) + + xy = points[:, :2] + intensity = points[:, 2] + + shifted = xy - self.robot_pos + c, s = np.cos(-self.robot_yaw), np.sin(-self.robot_yaw) + rot_matrix = np.array([[c, -s], [s, c]]) + aligned = shifted @ rot_matrix.T + aligned_points = np.hstack((aligned, intensity[:, None])) + + self.buffer.append((now, aligned_points)) + self.buffer = [(t, pts) for (t, pts) in self.buffer if now - t <= DECAY_TIME] + + all_points = np.vstack([pts for (_, pts) in self.buffer]) + xy = all_points[:, :2] + intensity = all_points[:, 2] + dists = np.linalg.norm(xy, axis=1) + + mask = (dists > ROBOT_RADIUS) | ((dists <= ROBOT_RADIUS) & (intensity > INTENSITY_THRESHOLD)) + filtered_points = xy[mask] + + if len(filtered_points) >= 18: + closest = np.argsort(np.linalg.norm(filtered_points, axis=1))[:18] + self.latest_lidar_xy = filtered_points[closest].astype(np.float32) + else: + self.latest_lidar_xy = np.zeros((18, 2), dtype=np.float32) + + def GetInitState(self, robot_state: SportModeState_): + self.px0 = robot_state.position[0] + self.py0 = robot_state.position[1] + self.yaw0 = robot_state.imu_state.rpy[2] + self.prev_yaw = self.yaw0 + print(self.yaw0, self.robot_yaw) + + def StandAndStart(self): + self.client.StandUp() + time.sleep(1) + self.client.BalanceStand() + print("Standing complete.") + + def encode_state(self, control: np.ndarray, lidar_obs: np.ndarray) -> np.ndarray: + lidar_xy = lidar_obs.reshape(18, 2) + dists = np.linalg.norm(lidar_xy, axis=1) + closest_idx = np.argmin(dists) + closest_vec = lidar_xy[closest_idx] + escape_vec = -closest_vec + escape_angle = np.arctan2(escape_vec[1], escape_vec[0]) + angle_diff = np.arctan2(math.sin(self.prev_yaw - escape_angle), math.cos(self.prev_yaw - escape_angle)) + + angle_obs = np.array([np.cos(angle_diff), np.sin(angle_diff)]) + self.angle_sequence.appendleft(angle_obs) + while len(self.angle_sequence) < self.angle_sequence_len: + self.angle_sequence.append(angle_obs) + angle_seq = np.concatenate(list(self.angle_sequence), axis=0) + + flat_lidar = lidar_obs.flatten() + self.obs_sequence.appendleft(flat_lidar) + while len(self.obs_sequence) < self.obs_sequence_len: + self.obs_sequence.append(flat_lidar) + lidar_seq = np.concatenate(list(self.obs_sequence), axis=0) + + return np.concatenate([control[:2], angle_seq, lidar_seq], axis=0) + + def safe_forward_control(self, robot_state: SportModeState_): + vx_desired = 0.5 + #No strafing + #vy_desired = 0.0 + + yaw = robot_state.imu_state.rpy[2] + yaw_error = math.atan2(math.sin(self.yaw0 - yaw), math.cos(self.yaw0 - yaw)) + + Kp_yaw = 2.0 + Kd_yaw = 0.5 + yaw_rate = (yaw - self.prev_yaw) / self.dt + vyaw_desired = Kp_yaw * yaw_error - Kd_yaw * yaw_rate + vyaw_desired = max(min(vyaw_desired, 1.0), -1.0) + self.prev_yaw = yaw + + ctrl_input = np.array([vx_desired, vyaw_desired], dtype=np.float32) + #raw_action = np.array([vx_desired, vy_desired, vyaw_desired], dtype=np.float32) + raw_action = ctrl_input + + state = self.encode_state(ctrl_input, self.latest_lidar_xy) + safe_action = self.safety_enforcer.get_action(state, raw_action) + + self.client.Move(*safe_action.tolist()) + print(f"[Shielded: {self.safety_enforcer.get_shielding_status()} | Q: {self.safety_enforcer.prev_q:.3f}]") + + def timer_callback(self): + if self.robot_state_ready: + self.t += self.dt + self.safe_forward_control(self.robot_state) + + +class SafeWalkerApp: + def __init__(self): + self.robot_state = None + self.robot_state_ready = False + + # DDS must be initialized before anything else + if len(sys.argv) > 1: + ChannelFactoryInitialize(0, sys.argv[1]) + else: + ChannelFactoryInitialize(0) + + # Subscribe to Go2's sportmodestate + self.sub = ChannelSubscriber("rt/sportmodestate", SportModeState_) + self.sub.Init(self.high_state_handler, 10) + + def high_state_handler(self, msg: SportModeState_): + self.robot_state = msg + self.robot_state_ready = True + + def run(self): + # Wait until the first robot_state message is received + print("⏳ Waiting for initial robot state...") + timeout = time.time() + 5.0 + while not self.robot_state_ready: + if time.time() > timeout: + print("❌ Timed out waiting for robot state.") + return + time.sleep(0.05) + + # Now safe to launch ROS 2 node + rclpy.init() + walker = SportModeSafeWalker() + walker.GetInitState(self.robot_state) + walker.StandAndStart() + rclpy.spin(walker) + + walker.destroy_node() + rclpy.shutdown() + + +if __name__ == "__main__": + app = SafeWalkerApp() + app.run() diff --git a/test2.py b/test2.py new file mode 100644 index 0000000..9d0ecff --- /dev/null +++ b/test2.py @@ -0,0 +1,95 @@ +import time +import sys +import math + +from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize +from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_ +from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ +from unitree_sdk2py.go2.sport.sport_client import SportClient + +class SportModeForwardPolicy: + def __init__(self) -> None: + # Timing + self.t = 0 + self.dt = 0.01 + + # Robot state (initial pose) + self.px0 = 0 + self.py0 = 0 + self.yaw0 = 0 + self.prev_yaw = 0 + + # Connect to SportClient + self.client = SportClient() + self.client.SetTimeout(10.0) + self.client.Init() + + def GetInitState(self, robot_state: SportModeState_): + self.px0 = robot_state.position[0] + self.py0 = robot_state.position[1] + self.yaw0 = robot_state.imu_state.rpy[2] + self.prev_yaw = self.yaw0 + + def ForwardFeedbackControl(self, robot_state: SportModeState_): + """ + PD controller to move forward at constant speed and maintain heading. + """ + # Desired forward velocity + vx_desired = 0.5 # m/s + vy_desired = 0.0 + + # Heading error + yaw = robot_state.imu_state.rpy[2] + yaw_error = math.atan2(math.sin(self.yaw0 - yaw), math.cos(self.yaw0 - yaw)) + + # PD control for yaw + Kp_yaw = 2.0 + Kd_yaw = 0.5 + yaw_rate = (yaw - self.prev_yaw) / self.dt + vyaw_command = Kp_yaw * yaw_error - Kd_yaw * yaw_rate + vyaw_command = max(min(vyaw_command, 1.0), -1.0) + + # Update previous yaw + self.prev_yaw = yaw + + # Send command + self.client.Move(vx_desired, vy_desired, vyaw_command) + + def StandAndStart(self): + """ + Stand up first before starting the walk. + """ + print("Standing up...") + self.client.StandUp() + time.sleep(1) + self.client.BalanceStand() + print("Entering feedback control loop") + +# === Robot state subscription === +robot_state = unitree_go_msg_dds__SportModeState_() + +def HighStateHandler(msg: SportModeState_): + global robot_state + robot_state = msg + +# === MAIN === +if __name__ == "__main__": + if len(sys.argv) > 1: + ChannelFactoryInitialize(0, sys.argv[1]) + else: + ChannelFactoryInitialize(0) + + sub = ChannelSubscriber("rt/sportmodestate", SportModeState_) + sub.Init(HighStateHandler, 10) + time.sleep(1) + + test = SportModeForwardPolicy() + test.GetInitState(robot_state) + test.StandAndStart() + + print("Starting forward feedback walk...") + + while True: + test.t += test.dt + test.ForwardFeedbackControl(robot_state) + time.sleep(test.dt) diff --git a/vec_pub.py b/vec_pub.py new file mode 100644 index 0000000..152b80a --- /dev/null +++ b/vec_pub.py @@ -0,0 +1,91 @@ +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import PointCloud2 +from geometry_msgs.msg import PoseStamped +import sensor_msgs_py.point_cloud2 as pc2 +import numpy as np +import os + +DECAY_TIME = 3.0 +ROBOT_RADIUS = 0.80 +INTENSITY_THRESHOLD = 220 +VEC_PATH = "/tmp/state_vec_72.npy" + + +TRANS = np.array([[0, -1], [-1, 0]]) + + +class LidarVec72Writer(Node): + def __init__(self): + super().__init__('lidar_vec72_writer') + + self.create_subscription(PointCloud2, '/lidar_chatter', self.lidar_callback, 10) + + + + self.buffer = [] # time-decayed point cloud + self.prev_lidar = np.zeros((18, 2), dtype=np.float32) + self.curr_lidar = np.zeros((18, 2), dtype=np.float32) + + self.prev_heading = None + self.curr_heading = None + + + def lidar_callback(self, msg: PointCloud2): + + now = self.get_clock().now().nanoseconds * 1e-9 + points = np.array([ + [p[0], p[1], p[3]] for p in pc2.read_points( + msg, field_names=("x", "y", "z", "intensity"), skip_nans=True) + ]) + + xy = points[:, :2] + intensity = points[:, 2] + aligned = xy @ TRANS.T #alignes utlidar with convetional x to right , y forward axis + aligned_points = np.hstack((aligned, intensity[:, None])) + + self.buffer.append((now, aligned_points)) + self.buffer = [(t, pts) for (t, pts) in self.buffer if now - t <= DECAY_TIME] + + all_points = np.vstack([pts for (_, pts) in self.buffer]) + xy = all_points[:, :2] + intensity = all_points[:, 2] + dists = np.linalg.norm(xy, axis=1) + mask = (dists > ROBOT_RADIUS) | ((dists <= ROBOT_RADIUS) & (intensity > INTENSITY_THRESHOLD)) + filtered = xy[mask] + + if len(filtered) < 18: + filtered = np.vstack((filtered, np.zeros((18 - len(filtered), 2)))) + else: + idx = np.argsort(np.linalg.norm(filtered, axis=1))[:18] + filtered = filtered[idx] + + # Update sequences + self.prev_lidar = self.curr_lidar + self.curr_lidar = filtered.astype(np.float32) + + + vec72 = np.concatenate([ + self.curr_lidar.flatten(), + self.prev_lidar.flatten() + ], axis=0) + + + tmp_path = VEC_PATH + ".tmp" + with open(tmp_path, 'wb') as f: + np.save(f, vec72) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, VEC_PATH) + + +def main(): + rclpy.init() + node = LidarVec72Writer() + rclpy.spin(node) + node.destroy_node() + rclpy.shutdown() + + +if __name__ == '__main__': + main() diff --git a/vec_pub_sem.py b/vec_pub_sem.py new file mode 100644 index 0000000..f84f4b2 --- /dev/null +++ b/vec_pub_sem.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 + +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import PointCloud2 +import sensor_msgs_py.point_cloud2 as pc2 +import numpy as np +import os + +DECAY_TIME = 3.0 +ROBOT_RADIUS = 0.85 +INTENSITY_THRESHOLD = 220 +VEC_PATH = "/tmp/state_vec_72.npy" + +TRANS = np.array([[0, -1], [-1, 0]]) + + +class LidarVec72Writer(Node): + def __init__(self): + super().__init__('lidar_vec72_writer') + + self.create_subscription(PointCloud2, '/lidar_chatter', self.lidar_callback, 10) + self.create_subscription(PointCloud2, '/semantic/pointcloud', self.semantic_callback, 10) + + self.buffer = [] # time-decayed point cloud (LiDAR + semantic) + self.prev_lidar = np.zeros((18, 2), dtype=np.float32) + self.curr_lidar = np.zeros((18, 2), dtype=np.float32) + + + def semantic_callback(self, msg: PointCloud2): + self._ingest_pointcloud(msg, use_intensity=False) + + def lidar_callback(self, msg: PointCloud2): + self._ingest_pointcloud(msg, use_intensity=True) + self._process() + + def _ingest_pointcloud(self, msg: PointCloud2, use_intensity=True): + now = self.get_clock().now().nanoseconds * 1e-9 + + fields = ("x", "y", "z", "intensity") if use_intensity else ("x", "y", "z") + try: + raw = list(pc2.read_points(msg, field_names=fields, skip_nans=True)) + if not raw: + return + raw = np.array(raw) + xy = raw[:, :2] + aligned = xy @ TRANS.T + + if use_intensity: + intensity = raw[:, 3] + else: + intensity = np.ones((raw.shape[0],), dtype=np.float32) * 255.0 + + aligned_points = np.hstack((aligned, intensity[:, None])) + self.buffer.append((now, aligned_points)) + except Exception as e: + self.get_logger().error(f"❌ Failed to parse point cloud: {e}") + + def _process(self): + now = self.get_clock().now().nanoseconds * 1e-9 + self.buffer = [(t, pts) for (t, pts) in self.buffer if now - t <= DECAY_TIME] + + if not self.buffer: + return + + all_points = np.vstack([pts for (_, pts) in self.buffer]) + xy = all_points[:, :2] + intensity = all_points[:, 2] + dists = np.linalg.norm(xy, axis=1) + + mask = (dists > ROBOT_RADIUS) | ((dists <= ROBOT_RADIUS) & (intensity > INTENSITY_THRESHOLD)) + filtered = xy[mask] + + if len(filtered) < 18: + filtered = np.vstack((filtered, np.zeros((18 - len(filtered), 2)))) + else: + idx = np.argsort(np.linalg.norm(filtered, axis=1))[:18] + filtered = filtered[idx] + + self.prev_lidar = self.curr_lidar + self.curr_lidar = filtered.astype(np.float32) + + + vec72 = np.concatenate([ + self.curr_lidar.flatten(), + self.prev_lidar.flatten() + ], axis=0) + + tmp_path = VEC_PATH + ".tmp" + with open(tmp_path, 'wb') as f: + np.save(f, vec72) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, VEC_PATH) + + +def main(): + rclpy.init() + node = LidarVec72Writer() + rclpy.spin(node) + node.destroy_node() + rclpy.shutdown() + + +if __name__ == '__main__': + main() + +