import torch
import math
import genesis as gs # type: ignore
from genesis.utils.geom import quat_to_xyz, transform_by_quat, inv_quat, transform_quat_by_quat # type: ignore
from genesis.engine.solvers.rigid.rigid_solver_decomp import RigidSolver
import numpy as np
import cv2

def gs_rand_float(lower, upper, shape, device):
    return (upper - lower) * torch.rand(size=shape, device=device) + lower

def gs_rand_normal(mean, std, shape, device):
    return mean + std * torch.randn(size=shape, device=device)
    
class WheelLeggedEnv:
    def __init__(self, num_envs, env_cfg, obs_cfg, reward_cfg, command_cfg, curriculum_cfg, 
                 domain_rand_cfg, terrain_cfg, robot_morphs="urdf", show_viewer=False, num_view = 1, device="cuda", train_mode=True):
        self.device = torch.device(device)

        self.mode = train_mode   #True训练模式开启
        
        self.num_envs = num_envs
        self.num_obs = obs_cfg["num_obs"]
        self.num_slice_obs = obs_cfg["num_slice_obs"]
        self.history_length = obs_cfg["history_length"]
        self.num_privileged_obs = None
        self.num_actions = env_cfg["num_actions"]
        self.num_commands = command_cfg["num_commands"]
        self.curriculum_cfg = curriculum_cfg
        self.domain_rand_cfg = domain_rand_cfg
        self.terrain_cfg = terrain_cfg
        respawn_points = self.terrain_cfg["respawn_points"]

        self.simulate_action_latency = True  # there is a 1 step latency on real robot
        self.dt = 0.01  # control frequency on real robot is 100hz
        self.max_episode_length = math.ceil(env_cfg["episode_length_s"] / self.dt)

        self.env_cfg = env_cfg
        self.obs_cfg = obs_cfg
        self.reward_cfg = reward_cfg
        self.command_cfg = command_cfg  

        self.obs_scales = obs_cfg["obs_scales"]
        self.reward_scales = reward_cfg["reward_scales"]
        self.noise = obs_cfg["noise"]

        # create scene
        self.scene = gs.Scene(
            sim_options=gs.options.SimOptions(dt=self.dt, substeps=5),
            viewer_options=gs.options.ViewerOptions(
                run_in_thread = True,
                max_FPS=int(1.0 / self.dt),
                refresh_rate = int(1.0 / self.dt),
                camera_pos=(2.0, 0.0, 2.5),
                camera_lookat=(0.0, 0.0, 0.5),
                camera_fov=40,
            ),
            vis_options=gs.options.VisOptions(rendered_envs_idx=list(range(num_view))),
            rigid_options=gs.options.RigidOptions(
                dt=self.dt,
                constraint_solver=gs.constraint_solver.Newton,
                enable_collision=True,
                enable_joint_limit=True,
                batch_dofs_info=True,
                max_collision_pairs=64,
                enable_self_collision=False,
                # batch_links_info=True,
            ),
            show_viewer=show_viewer,
        )

        # add plane
        self.scene.add_entity(gs.morphs.URDF(file="assets/terrain/plane/plane.urdf", fixed=True))
        # init roboot quat and pos
        match robot_morphs:
            case "urdf":
                self.base_init_pos = torch.tensor(self.env_cfg["base_init_pos"]["urdf"], device=self.device)
            case "mjcf":
                self.base_init_pos = torch.tensor(self.env_cfg["base_init_pos"]["mjcf"], device=self.device)
            case _:
                self.base_init_pos = torch.tensor(self.env_cfg["base_init_pos"]["urdf"], device=self.device)
        self.base_init_quat = torch.tensor(self.env_cfg["base_init_quat"], device=self.device)
        self.inv_base_init_quat = inv_quat(self.base_init_quat)
        # add terrain 只能有一个Terrain(genesis v0.2.1)
        self.horizontal_scale = self.terrain_cfg["horizontal_scale"]
        self.vertical_scale = self.terrain_cfg["vertical_scale"]
        self.height_field = cv2.imread("assets/terrain/png/"+self.terrain_cfg["train"]+".png", cv2.IMREAD_GRAYSCALE)
        self.terrain_height = torch.tensor(self.height_field, device=self.device) * self.vertical_scale
        terrain_pos = []
        if self.terrain_cfg["terrain"]:
            print("\033[1;35m open terrain\033[0m")
            if self.mode:
                self.terrain = self.scene.add_entity(
                morph=gs.morphs.Terrain(
                height_field = self.height_field,
                horizontal_scale=self.horizontal_scale, 
                vertical_scale=self.vertical_scale,
                ),)
                for i in range(len(respawn_points)):
                    terrain_pos.append(respawn_points[i])
            else:
                height_field = cv2.imread("assets/terrain/png/"+self.terrain_cfg["eval"]+".png", cv2.IMREAD_GRAYSCALE)
                self.terrain = self.scene.add_entity(
                morph=gs.morphs.Terrain(
                pos = (1.0,1.0,0.0),
                height_field = height_field,
                horizontal_scale=self.horizontal_scale, 
                vertical_scale=self.vertical_scale,
                ),)     
                pos = self.base_init_pos.cpu().numpy()
                terrain_pos.append([pos[0], pos[1], 0])
                # print("\033[1;34m respawn_points: \033[0m",self.base_init_pos)
                
        # 楼梯地形
        if self.terrain_cfg["vertical_stairs"]:
            self.v_stairs_num = self.terrain_cfg["v_stairs_num"]
            self.v_stairs_height = self.terrain_cfg["v_stairs_height"]
            self.v_stairs_width = self.terrain_cfg["v_stairs_width"]
            self.v_plane_size = self.terrain_cfg["v_plane_size"]
            
            inverted_pyramid_size = self.v_plane_size + self.v_stairs_width * 2 * (self.v_stairs_num*2) 
            inverted_pyramid_point = [inverted_pyramid_size,-inverted_pyramid_size/2,0]
            self.add_inverted_pyramid(inverted_pyramid_point)
            pyramid_size = self.v_plane_size + self.v_stairs_width*2* (self.v_stairs_num-1) 
            pyramid_point = [pyramid_size,-inverted_pyramid_size-pyramid_size/2,0]
            self.add_pyramid(pyramid_point)
            
            inverted_pyramid_point[2] += self.v_stairs_height
            terrain_pos.append(inverted_pyramid_point)
            pyramid_point[2] += self.v_stairs_height * self.v_stairs_num
            terrain_pos.append(pyramid_point)
            
        # 构建复活点
        self.num_respawn_points = len(terrain_pos)
        self.base_terrain_pos = torch.tensor(terrain_pos,device=self.device)
        self.base_terrain_pos[:,2] += self.base_init_pos[2]
        print("\033[1;34m respawn_points: \033[0m",self.base_terrain_pos)

        # add robot
        base_init_pos = self.base_init_pos.cpu().numpy()
        if self.terrain_cfg["terrain"]:
            if self.mode:
                base_init_pos = self.base_terrain_pos[0].cpu().numpy()

        match robot_morphs:
            case "urdf":
                self.robot = self.scene.add_entity(
                    gs.morphs.URDF(
                        file = self.env_cfg["urdf"],
                        pos = base_init_pos,
                        quat=self.base_init_quat.cpu().numpy(),
                        convexify=self.env_cfg["convexify"],
                        decimate_aggressiveness=self.env_cfg["decimate_aggressiveness"],
                    ),
                )
            case "mjcf":
                self.robot = self.scene.add_entity(
                    gs.morphs.MJCF(
                        file = self.env_cfg["mjcf"],
                        pos=base_init_pos,
                        quat=self.base_init_quat.cpu().numpy(),
                        convexify=self.env_cfg["convexify"],
                        decimate_aggressiveness=self.env_cfg["decimate_aggressiveness"],
                    ),
                    vis_mode='collision'
                )
            case _:
                raise Exception("what robot morphs?(shoud urdf/mjcf)")
            
        # build
        self.scene.build(n_envs=num_envs)

        # names to indices
        self.motors_dof_idx = [self.robot.get_joint(name).dof_start for name in self.env_cfg["joint_names"]]
        joint_dof_idx = []
        wheel_dof_idx = []
        self.joint_dof_idx = []
        self.wheel_dof_idx = []
        for i in range(len(self.env_cfg["joint_names"])):
            if self.env_cfg["joint_type"][self.env_cfg["joint_names"][i]] == "joint":
                joint_dof_idx.append(i)
                self.joint_dof_idx.append(self.motors_dof_idx[i])
            elif self.env_cfg["joint_type"][self.env_cfg["joint_names"][i]] == "wheel":
                wheel_dof_idx.append(i)
                self.wheel_dof_idx.append(self.motors_dof_idx[i])
        self.joint_dof_idx_np = np.array(joint_dof_idx)
        self.wheel_dof_idx_np = np.array(wheel_dof_idx)

        # PD control parameters
        self.kp = np.full((self.num_envs, self.num_actions), self.env_cfg["joint_kp"])
        self.kv = np.full((self.num_envs, self.num_actions), self.env_cfg["joint_kv"])
        self.kp[:,self.wheel_dof_idx_np] = 0.0
        self.kv[:,self.wheel_dof_idx_np] = self.env_cfg["wheel_kv"]
        self.robot.set_dofs_kp(self.kp, self.motors_dof_idx)
        self.robot.set_dofs_kv(self.kv, self.motors_dof_idx)
        
        damping = np.full((self.num_envs, self.robot.n_dofs), self.env_cfg["damping"])
        damping[:,:6] = 0
        self.robot.set_dofs_damping(damping, np.arange(0,self.robot.n_dofs))
        
        # from IPython import embed; embed()
        # print(self.scene.sim.rigid_solver.dofs_info.damping.to_numpy())
        # print(self.scene.sim.rigid_solver.dofs_info.stiffness.to_numpy())
        armature = np.full((self.num_envs, self.robot.n_dofs), self.env_cfg["armature"])
        armature[:,:6] = 0
        self.robot.set_dofs_armature(armature, np.arange(0, self.robot.n_dofs))
        
        #dof limits
        lower = [self.env_cfg["dof_limit"][name][0] for name in self.env_cfg["joint_names"]]
        upper = [self.env_cfg["dof_limit"][name][1] for name in self.env_cfg["joint_names"]]
        self.dof_pos_lower = torch.tensor(lower).to(self.device)
        self.dof_pos_upper = torch.tensor(upper).to(self.device)

        # set safe force
        lower = np.array([[-self.env_cfg["safe_force"][name] for name in self.env_cfg["joint_names"]] for _ in range(num_envs)])
        upper = np.array([[self.env_cfg["safe_force"][name] for name in self.env_cfg["joint_names"]] for _ in range(num_envs)])
        self.robot.set_dofs_force_range(
            lower          = torch.tensor(lower, device=self.device, dtype=torch.float32),
            upper          = torch.tensor(upper, device=self.device, dtype=torch.float32),
            dofs_idx_local = self.motors_dof_idx,
        )

        # prepare reward functions and multiply reward scales by dt
        self.reward_functions, self.episode_sums = dict(), dict()
        for name in self.reward_scales.keys():
            self.reward_scales[name] *= self.dt
            self.reward_functions[name] = getattr(self, "_reward_" + name)
            self.episode_sums[name] = torch.zeros((self.num_envs,), device=self.device, dtype=gs.tc_float)
        # 不要zeros初始化
        self.rew_survive = torch.ones(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False)

        # 存活比例
        self.survive_ratio = torch.zeros((self.num_envs,), device=self.device, dtype=gs.tc_float)

        # prepare command_ranges lin_vel_x lin_vel_y ang_vel height_target
        self.command_ranges = torch.zeros((self.num_envs, self.num_commands,2),device=self.device,dtype=gs.tc_float)
        self.command_ranges[:,0,0] = self.command_cfg["lin_vel_x_range"][0] * self.curriculum_cfg["curriculum_lin_vel_min_range"]
        self.command_ranges[:,0,1] = self.command_cfg["lin_vel_x_range"][1] * self.curriculum_cfg["curriculum_lin_vel_min_range"]
        self.command_ranges[:,1,0] = self.command_cfg["lin_vel_y_range"][0] * self.curriculum_cfg["curriculum_lin_vel_min_range"]
        self.command_ranges[:,1,1] = self.command_cfg["lin_vel_y_range"][1] * self.curriculum_cfg["curriculum_lin_vel_min_range"]
        self.command_ranges[:,2,0] = self.command_cfg["ang_vel_range"][0] * self.curriculum_cfg["curriculum_ang_vel_min_range"]
        self.command_ranges[:,2,1] = self.command_cfg["ang_vel_range"][1] * self.curriculum_cfg["curriculum_ang_vel_min_range"]
        self.command_ranges[:,3,0] = self.command_cfg["leg_length_range"][0]
        self.command_ranges[:,3,1] = self.command_cfg["leg_length_range"][1]
        self.command_ranges[:,4,0] = self.command_cfg["leg_length_range"][0]
        self.command_ranges[:,4,1] = self.command_cfg["leg_length_range"][1]
        self.command_ranges[:,5,0] = self.command_cfg["tsk_range"][0]
        self.command_ranges[:,5,1] = self.command_cfg["tsk_range"][1]
        self.lin_vel_error = torch.zeros(self.num_envs, device=self.device, dtype=gs.tc_float)
        self.ang_vel_error = torch.zeros(self.num_envs, device=self.device, dtype=gs.tc_float)
        # self.height_error = torch.zeros((self.num_envs,1), device=self.device, dtype=gs.tc_float)
        self.linx_vel_err_range_err = self.curriculum_cfg["lin_vel_err_range"][1]-self.curriculum_cfg["lin_vel_err_range"][0]
        self.ang_vel_err_range_err = self.curriculum_cfg["ang_vel_err_range"][1]-self.curriculum_cfg["ang_vel_err_range"][0]
        self.linx_range_up_threshold = 0.0
        self.angv_range_up_threshold = 0.0
        # curriculum
        self.curriculum_step = 0
        
        # initialize buffers
        self.base_lin_vel = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.last_base_lin_vel = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.base_lin_acc = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.base_ang_vel = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.last_base_ang_vel = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.base_ang_acc = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.projected_gravity = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.global_gravity = torch.tensor([0.0, 0.0, -1.0], device=self.device, dtype=gs.tc_float).repeat(self.num_envs, 1)

        self.slice_obs_buf = torch.zeros((self.num_envs, self.num_slice_obs), device=self.device, dtype=gs.tc_float)
        self.history_obs_buf = torch.zeros((self.num_envs, self.history_length, self.num_slice_obs), device=self.device, dtype=gs.tc_float)
        self.obs_buf = torch.zeros((self.num_envs, self.num_obs), device=self.device, dtype=gs.tc_float)
        self.history_idx = torch.zeros(self.num_envs, dtype=torch.long, device=self.device)
        self.rew_buf = torch.zeros((self.num_envs,), device=self.device, dtype=gs.tc_float)
        self.curriculum_rew_buf = torch.zeros_like(self.rew_buf)
        self.reset_buf = torch.ones((self.num_envs,), device=self.device, dtype=gs.tc_int)
        self.episode_length_buf = torch.zeros((self.num_envs,), device=self.device, dtype=gs.tc_int)
        self.commands = torch.zeros((self.num_envs, self.num_commands), device=self.device, dtype=gs.tc_float)
        self.commands_scale = torch.tensor(
            [self.obs_scales["lin_vel"], self.obs_scales["lin_vel"], self.obs_scales["ang_vel"], 
             self.obs_scales["dof_pos"], self.obs_scales["dof_pos"], self.obs_scales["dof_pos"]], 
            device=self.device,
            dtype=gs.tc_float,
        )

        self.actions = torch.zeros((self.num_envs, self.num_actions), device=self.device, dtype=gs.tc_float)
        self.last_actions = torch.zeros_like(self.actions)
        self.dof_pos = torch.zeros_like(self.actions)
        self.dof_vel = torch.zeros_like(self.actions)
        self.dof_force = torch.zeros_like(self.actions)
        self.last_dof_vel = torch.zeros_like(self.actions)
        self.base_pos = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.base_quat = torch.zeros((self.num_envs, 4), device=self.device, dtype=gs.tc_float)
        self.basic_default_dof_pos = torch.tensor(
            [self.env_cfg["default_joint_angles"][name] for name in self.env_cfg["joint_names"]],
            device=self.device,
            dtype=gs.tc_float,
        )
        
        default_dof_pos_list = [[self.env_cfg["default_joint_angles"][name] for name in self.env_cfg["joint_names"]]] * self.num_envs
        self.default_dof_pos = torch.tensor(default_dof_pos_list,device=self.device,dtype=gs.tc_float,)
        init_dof_pos_list = [[self.env_cfg["joint_init_angles"][name] for name in self.env_cfg["joint_names"]]] * self.num_envs
        self.init_dof_pos = torch.tensor(init_dof_pos_list,device=self.device,dtype=gs.tc_float,)
        
        self.connect_force = torch.zeros((self.num_envs,self.robot.n_links, 3), device=self.device, dtype=gs.tc_float)
        self.extras = dict()  # extra information for logging
        
        #跪地重启   注意是idx_local不需要减去base_idx
        if(self.env_cfg["termination_if_base_connect_plane_than"]&self.mode):
            self.reset_links = [(self.robot.get_link(name).idx_local) for name in self.env_cfg["connect_plane_links"]]
        #足端位置
        self.left_foot = self.robot.get_link(self.env_cfg["foot_link"][0])
        self.right_foot = self.robot.get_link(self.env_cfg["foot_link"][1])
        self.left_foot_pos = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.right_foot_pos = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.left_foot_pos[:] = self.left_foot.get_pos()
        self.right_foot_pos[:] = self.right_foot.get_pos()
        self.left_foot_base_pos = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
        self.right_foot_base_pos = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)
            
        #域随机化 domain_rand_cfg
        self.friction_ratio_low = self.domain_rand_cfg["friction_ratio_range"][0]
        self.friction_ratio_range = self.domain_rand_cfg["friction_ratio_range"][1] - self.friction_ratio_low
        self.base_mass_low = self.domain_rand_cfg["random_base_mass_shift_range"][0]
        self.base_mass_range = self.domain_rand_cfg["random_base_mass_shift_range"][1] - self.base_mass_low  
        self.other_mass_low = self.domain_rand_cfg["random_other_mass_shift_range"][0]
        self.other_mass_range = self.domain_rand_cfg["random_other_mass_shift_range"][1] - self.other_mass_low            
        self.dof_damping_low = self.domain_rand_cfg["damping_range"][0]
        self.dof_damping_range = self.domain_rand_cfg["damping_range"][1] - self.dof_damping_low
        self.dof_armature_low = self.domain_rand_cfg["dof_armature_range"][0]
        self.dof_armature_range = self.domain_rand_cfg["dof_armature_range"][1] - self.dof_armature_low
        self.kp_low = self.domain_rand_cfg["random_KP"][0]
        self.kp_range = self.domain_rand_cfg["random_KP"][1] - self.kp_low
        self.kv_low = self.domain_rand_cfg["random_KV"][0]
        self.kv_range = self.domain_rand_cfg["random_KV"][1] - self.kv_low
        self.joint_angle_low = self.domain_rand_cfg["random_default_joint_angles"][0]
        self.joint_angle_range = self.domain_rand_cfg["random_default_joint_angles"][1] - self.joint_angle_low
        #地形训练索引
        self.terrain_buf = torch.ones((self.num_envs,), device=self.device, dtype=gs.tc_int)
        # print("self.obs_buf.size(): ",self.obs_buf.size())
        self.episode_lengths = torch.zeros(self.num_envs, dtype=torch.float32, device=self.device)
        # 跟踪误差
        self.mean_lin_vel_error = 0
        self.mean_ang_vel_error = 0
        #外部力 TODO
        for solver in self.scene.sim.solvers:
            if not isinstance(solver, RigidSolver):
                continue
            rigid_solver = solver

        # print("self.init_dof_pos",self.init_dof_pos)
        #初始化角度
        self.reset()
        
    def _resample_commands(self, envs_idx): 
        if self.command_cfg["high_speed"]:  
            for command_idx in (0,1):
                low = self.command_ranges[envs_idx, command_idx, 0]
                high = self.command_ranges[envs_idx, command_idx, 1]
                self.commands[envs_idx, command_idx] = gs_rand_float(low, high, envs_idx.shape, self.device)
            # 根据线速度求角速度max 因为角速度范围大,线速度范围小所以用线速度采样决定角速度范围比较均匀
            safe_linv_x = torch.clip(torch.abs(self.commands[envs_idx, 0]), min=1e-4)
            angv_limit = self.command_cfg["inverse_linx_angv"]/safe_linv_x
            angv_low = torch.clip(self.command_ranges[envs_idx, 2, 0],-angv_limit)
            angv_high = torch.clip(self.command_ranges[envs_idx, 2, 1],torch.zeros(envs_idx.shape, dtype=torch.float32, device=self.device),angv_limit)
            self.commands[envs_idx, 2] = gs_rand_float(angv_low, angv_high, envs_idx.shape, self.device)
            safe_angv = torch.clip(torch.abs(self.commands[envs_idx, 2]), min=1e-4)
            #角速度命令高就要限制tsk 趋近0
            tsk_std = self.command_cfg["inverse_tsk"]/safe_angv
            tsk_std = torch.clip(tsk_std, 1e-5, 2.0)
            self.commands[envs_idx, 5] = gs_rand_normal(0, tsk_std, envs_idx.shape, self.device)
            self.commands[envs_idx, 5] = torch.clip(self.commands[envs_idx, 5], self.command_cfg["tsk_range"][0], self.command_cfg["tsk_range"][1])
            #角速度命令高就要限制腿部相似 随机一个mean然后高斯采样
            leg_length_mean = gs_rand_float(self.command_ranges[envs_idx, 3, 0],
                                           self.command_ranges[envs_idx, 3, 1],
                                           envs_idx.shape, self.device)
            leg_length_std = self.command_cfg["inverse_leg_length"]/safe_angv
            leg_length_std = torch.clip(leg_length_std, 1e-5, 2.0)
            self.commands[envs_idx, 3] = gs_rand_normal(leg_length_mean, leg_length_std, envs_idx.shape, self.device)
            self.commands[envs_idx, 4] = gs_rand_normal(leg_length_mean, leg_length_std, envs_idx.shape, self.device)
            self.commands[envs_idx, 3] = torch.clip(self.commands[envs_idx, 3], self.command_cfg["leg_length_range"][0], self.command_cfg["leg_length_range"][1])
            self.commands[envs_idx, 4] = torch.clip(self.commands[envs_idx, 4], self.command_cfg["leg_length_range"][0], self.command_cfg["leg_length_range"][1])
        else:
            for command_idx in range(self.num_commands):
                low = self.command_ranges[envs_idx, command_idx, 0]
                high = self.command_ranges[envs_idx, command_idx, 1]
                self.commands[envs_idx, command_idx] = gs_rand_float(low, high, envs_idx.shape, self.device)
        
    def set_commands(self,envs_idx,commands):
        self.commands[envs_idx]=torch.tensor(commands,device=self.device, dtype=gs.tc_float)

    def step(self, actions):
        self.actions = torch.clip(actions, -self.env_cfg["clip_actions"], self.env_cfg["clip_actions"])
        # self.actions[:,0] = 0
        # self.actions[:,3] = 0
        exec_actions = self.last_actions if self.simulate_action_latency else self.actions
        target_dof_pos = exec_actions[:,self.joint_dof_idx_np] * self.env_cfg["joint_action_scale"] + self.default_dof_pos[:,self.joint_dof_idx_np]
        target_dof_vel = exec_actions[:,self.wheel_dof_idx_np] * self.env_cfg["wheel_action_scale"]
        #dof limits
        target_dof_pos = torch.clamp(target_dof_pos, min=self.dof_pos_lower[self.joint_dof_idx_np], max=self.dof_pos_upper[self.joint_dof_idx_np])
        self.robot.control_dofs_position(target_dof_pos, self.joint_dof_idx)
        self.robot.control_dofs_velocity(target_dof_vel, self.wheel_dof_idx)

        self.scene.step()
        # update buffers
        self.episode_length_buf += 1
        self.base_pos[:] = self.get_relative_terrain_pos(self.robot.get_pos())
        # print("base z",self.base_pos[:2])
        self.base_quat[:] = self.robot.get_quat()
        self.base_euler = quat_to_xyz(
            transform_quat_by_quat(torch.ones_like(self.base_quat) * self.inv_base_init_quat, self.base_quat),
            rpy=True,
            degrees=True,
        )
        inv_base_quat = inv_quat(self.base_quat)
        self.base_lin_vel[:] = transform_by_quat(self.robot.get_vel(), inv_base_quat)
        self.base_lin_acc[:] = (self.base_lin_vel[:] - self.last_base_lin_vel[:])/ self.dt
        self.base_ang_vel[:] = transform_by_quat(self.robot.get_ang(), inv_base_quat) 
        self.base_ang_acc[:] = (self.base_ang_vel[:] - self.last_base_ang_vel[:]) / self.dt
        self.projected_gravity = transform_by_quat(self.global_gravity, inv_base_quat) 
        self.dof_pos[:] = self.robot.get_dofs_position(self.motors_dof_idx) 
        self.dof_vel[:] = self.robot.get_dofs_velocity(self.motors_dof_idx) 
        self.dof_force[:] = self.robot.get_dofs_force(self.motors_dof_idx)
        self.left_foot_pos[:] = self.left_foot.get_pos()
        self.right_foot_pos[:] = self.right_foot.get_pos()
        self.left_foot_base_pos[:] = transform_by_quat(self.left_foot_pos, inv_base_quat) + self.base_pos
        self.right_foot_base_pos[:] = transform_by_quat(self.right_foot_pos, inv_base_quat) + self.base_pos

        # print("dof_force:",self.dof_force)
        if self.noise["use"]:
            self.base_ang_vel[:] += torch.randn_like(self.base_ang_vel) * self.noise["ang_vel"][0] + (torch.rand_like(self.base_ang_vel)*2-1) * self.noise["ang_vel"][1]
            self.projected_gravity += torch.randn_like(self.projected_gravity) * self.noise["gravity"][0] + (torch.rand_like(self.projected_gravity)*2-1) * self.noise["gravity"][1]
            self.dof_pos[:] += torch.randn_like(self.dof_pos) * self.noise["dof_pos"][0] + (torch.rand_like(self.dof_pos)*2-1) * self.noise["dof_pos"][1]
            self.dof_vel[:] += torch.randn_like(self.dof_vel) * self.noise["dof_vel"][0] + (torch.rand_like(self.dof_vel)*2-1) * self.noise["dof_vel"][1]
        
        #碰撞力
        self.connect_force = self.robot.get_links_net_contact_force()

        # update last
        self.last_base_lin_vel[:] = self.base_lin_vel[:]
        self.last_base_ang_vel[:] = self.base_ang_vel[:]
        
        #步数
        self.episode_lengths += 1

        # resample commands
        envs_idx = (
            (self.episode_length_buf % int(self.env_cfg["resampling_time_s"] / self.dt) == 0)
            .nonzero(as_tuple=False)
            .flatten()
        )

        # check terrain_buf
        # 线速度达到预设的90%范围，角速度达到90%以上去其他地形(建议高一点)
        self.terrain_buf[:] = 0
        # self.terrain_buf = self.command_ranges[:, 0, 1] > self.command_cfg["lin_vel_x_range"][1] * 0.9
        # self.terrain_buf &= self.command_ranges[:, 2, 1] > self.command_cfg["ang_vel_range"][1] * 0.9
        #固定一部分去地形
        if(self.mode):
            if self.survive_ratio > 0.9:
                self.terrain_buf[:int(self.num_envs*0.15)] = 1
        
        # compute curriculum
        self.lin_vel_error += torch.abs(self.commands[:, 0] - self.base_lin_vel[:, 0])
        self.ang_vel_error += torch.abs(self.commands[:, 2] - self.base_ang_vel[:, 2])
            
        # check termination and reset
        if(self.mode):
            self.check_termination()
            self._resample_commands(envs_idx)
            self.curriculum_commands()

        time_out_idx = (self.episode_length_buf > self.max_episode_length).nonzero(as_tuple=False).flatten()
        self.extras["time_outs"] = torch.zeros_like(self.reset_buf, device=self.device, dtype=gs.tc_float)
        self.extras["time_outs"][time_out_idx] = 1.0
        if(self.mode):
            self.reset_idx(self.reset_buf.nonzero(as_tuple=False).flatten())

        # compute reward
        if(self.mode):
            self.rew_buf[:] = 0.0
            for name, reward_func in self.reward_functions.items():
                rew = reward_func() * self.reward_scales[name]
                self.rew_buf += rew
                self.episode_sums[name] += rew
        
        if self.reward_cfg["only_positive_rewards"]:
            self.rew_buf[:] = torch.clip(self.rew_buf[:], min=0.0)
            
        # else:
        #     print("base_lin_vel: ",self.base_lin_vel[0,:])
            
        # compute observations
        self.slice_obs_buf = torch.cat(
            [
                # self.base_lin_vel * self.obs_scales["lin_vel"],  # 3
                self.base_ang_vel * self.obs_scales["ang_vel"],  # 3
                self.projected_gravity,  # 3
                # self.commands * self.commands_scale,  # 4
                (self.dof_pos[:,self.joint_dof_idx_np] - self.default_dof_pos[:,self.joint_dof_idx_np]) * self.obs_scales["dof_pos"],  # 6
                self.dof_vel * self.obs_scales["dof_vel"],  # 8
                self.actions,  # 8
            ],
            axis=-1,
        )
        self.last_actions[:] = self.actions[:]
        self.last_dof_vel[:] = self.dof_vel[:]

        # print("slice_obs_buf: ",self.slice_obs_buf)
        
        # Combine the current observation with historical observations (e.g., along the time axis)
        self.obs_buf = torch.cat([self.history_obs_buf, self.slice_obs_buf.unsqueeze(1)], dim=1).view(self.num_envs, -1)
        # Update history buffer
        if self.history_length > 1:
            self.history_obs_buf[:, :-1, :] = self.history_obs_buf[:, 1:, :].clone() # 移位操作
        self.history_obs_buf[:, -1, :] = self.slice_obs_buf 
        
        self.obs_buf = torch.cat([self.obs_buf, self.commands * self.commands_scale], axis=-1)
        
        return self.obs_buf, None, self.rew_buf, self.reset_buf, self.extras

    def get_observations(self):
        return self.obs_buf

    def get_privileged_observations(self):
        return None

    def reset_idx(self, envs_idx):
        if len(envs_idx) == 0:
            return
        # 计算要重置的环境步数
        self.survive_ratio = self.episode_length_buf[envs_idx].float().mean() / self.max_episode_length
        # print("\033[1;32m Survive Ratio: \033[0m", self.survive_ratio)
        # reset dofs
        self.dof_pos[envs_idx] = self.init_dof_pos[envs_idx]
        self.dof_vel[envs_idx] = 0.0
        self.robot.set_dofs_position(
            position=self.dof_pos[envs_idx],
            dofs_idx_local=self.motors_dof_idx,
            zero_velocity=True,
            envs_idx=envs_idx,
        )

        # reset base
        self.base_quat[envs_idx] = self.base_init_quat.reshape(1, -1)
        if self.terrain_cfg["terrain"]:
            if self.mode:
                terrain_buf = self.terrain_buf[envs_idx]
                terrain_idx = envs_idx[terrain_buf.nonzero(as_tuple=False).flatten()] # 获取 envs_idx 中满足条件的索引
                non_terrain_idx = envs_idx[(terrain_buf<1).nonzero(as_tuple=False).flatten()] # 获取 envs_idx 中不满足条件的索引
                # 设置地形位置
                if len(terrain_idx) > 0: # 只有当有满足地形重置条件的环境时才执行
                    #目前认为坡路和崎岖路面是相同难度，所以reset随机选取一个环境去复活
                    n = len(terrain_idx)
                    random_idx = torch.randint(1, self.num_respawn_points, (n,)) # 注意从 1 开始，避免使用 base_terrain_pos[0] 作为随机位置
                    selected_pos = self.base_terrain_pos[random_idx]
                    self.base_pos[terrain_idx] = selected_pos
                # 设置非地形位置 (默认位置)
                if len(non_terrain_idx) > 0:
                    self.base_pos[non_terrain_idx] = self.base_terrain_pos[0]
            else:
                num_terrains = len(self.base_terrain_pos)
                num_per_terrain = self.num_envs // num_terrains 
                remainder = self.num_envs % num_terrains
                env_id = 0
                for i in range(num_terrains):
                    for j in range(num_per_terrain):
                        self.base_pos[env_id] = self.base_terrain_pos[i]
                        env_id += 1
                self.base_pos[-remainder:] = self.base_terrain_pos[0]
        else:
            self.base_pos[envs_idx] = self.base_init_pos   #没开地形就基础
            
        self.robot.set_pos(self.base_pos[envs_idx], zero_velocity=False, envs_idx=envs_idx)
        self.robot.set_quat(self.base_quat[envs_idx], zero_velocity=False, envs_idx=envs_idx)
        self.base_lin_vel[envs_idx] = 0
        self.base_ang_vel[envs_idx] = 0
        self.robot.zero_all_dofs_velocity(envs_idx)

        # reset buffers
        self.last_actions[envs_idx] = 0.0
        self.last_dof_vel[envs_idx] = 0.0
        self.episode_length_buf[envs_idx] = 0
        self.reset_buf[envs_idx] = True

        # fill extras
        self.extras["episode"] = {}
        for key in self.episode_sums.keys():
            self.extras["episode"]["rew_" + key] = (
                torch.mean(self.episode_sums[key][envs_idx]).item() / self.env_cfg["episode_length_s"]
            )
            self.episode_sums[key][envs_idx] = 0.0

        if self.mode:
            self._resample_commands(envs_idx)
            # if self.survive_ratio > 0.7:
        self.domain_rand(envs_idx)
        self.episode_lengths[envs_idx] = 0.0
        # 重置误差
        self.lin_vel_error[envs_idx] = 0
        self.ang_vel_error[envs_idx] = 0

    def reset(self):
        self.reset_buf[:] = True
        self.reset_idx(torch.arange(self.num_envs, device=self.device))
        return self.obs_buf, None

    def check_termination(self):
        self.reset_buf = self.episode_length_buf > self.max_episode_length
        if self.survive_ratio > 0.8:
            self.reset_buf |= torch.abs(self.base_euler[:, 1]) > self.env_cfg["termination_if_pitch_greater_than"][1]
            self.reset_buf |= torch.abs(self.base_euler[:, 0]) > self.env_cfg["termination_if_roll_greater_than"][1]
        else:
            self.reset_buf |= torch.abs(self.base_euler[:, 1]) > self.env_cfg["termination_if_pitch_greater_than"][0]
            self.reset_buf |= torch.abs(self.base_euler[:, 0]) > self.env_cfg["termination_if_roll_greater_than"][0]
        # self.reset_buf |= torch.abs(self.base_pos[:, 2]) < self.env_cfg["termination_if_base_height_greater_than"]
        if(self.env_cfg["termination_if_base_connect_plane_than"]):
            for idx in self.reset_links:
                self.reset_buf |= torch.abs(self.connect_force[:,idx,:]).sum(dim=1) > 0
        
    def domain_rand(self, envs_idx):
        friction_ratio = self.friction_ratio_low + self.friction_ratio_range * torch.rand(len(envs_idx), self.robot.n_links)
        self.robot.set_friction_ratio(friction_ratio=friction_ratio,
                                      links_idx_local=np.arange(0, self.robot.n_links),
                                      envs_idx = envs_idx)

        base_mass_shift = self.base_mass_low + self.base_mass_range * torch.rand(len(envs_idx), 1, device=self.device)
        other_mass_shift =-self.other_mass_low + self.other_mass_range * torch.rand(len(envs_idx), self.robot.n_links - 1, device=self.device)
        mass_shift = torch.cat((base_mass_shift, other_mass_shift), dim=1)
        self.robot.set_mass_shift(mass_shift=mass_shift,
                                  links_idx_local=np.arange(0, self.robot.n_links),
                                  envs_idx = envs_idx)

        base_com_shift = -self.domain_rand_cfg["random_base_com_shift"] / 2 + self.domain_rand_cfg["random_base_com_shift"] * torch.rand(len(envs_idx), 1, 3, device=self.device)
        other_com_shift = -self.domain_rand_cfg["random_other_com_shift"] / 2 + self.domain_rand_cfg["random_other_com_shift"] * torch.rand(len(envs_idx), self.robot.n_links - 1, 3, device=self.device)
        com_shift = torch.cat((base_com_shift, other_com_shift), dim=1)
        self.robot.set_COM_shift(com_shift=com_shift,
                                 links_idx_local=np.arange(0, self.robot.n_links),
                                 envs_idx = envs_idx)

        kp_shift = (self.kp_low + self.kp_range * torch.rand(len(envs_idx), self.num_actions, device="cpu")) * self.kp[0]
        self.robot.set_dofs_kp(kp_shift, self.motors_dof_idx, envs_idx=envs_idx)

        kv_shift = (self.kv_low + self.kv_range * torch.rand(len(envs_idx), self.num_actions, device="cpu")) * self.kv[0]
        self.robot.set_dofs_kv(kv_shift, self.motors_dof_idx, envs_idx = envs_idx)

        #random_default_joint_angles
        dof_pos_shift = self.joint_angle_low + self.joint_angle_range * torch.rand(len(envs_idx),self.num_actions,device=self.device,dtype=gs.tc_float)
        self.default_dof_pos[envs_idx] = dof_pos_shift + self.basic_default_dof_pos

        damping = (self.dof_damping_low+self.dof_damping_range * torch.rand(len(envs_idx), self.robot.n_dofs))
        damping[:,:6] = 0
        self.robot.set_dofs_damping(damping=damping, 
                                   dofs_idx_local=np.arange(0, self.robot.n_dofs), 
                                   envs_idx=envs_idx)

        armature = (self.dof_armature_low+self.dof_armature_range * torch.rand(len(envs_idx), self.robot.n_dofs))
        armature[:,:6] = 0
        self.robot.set_dofs_armature(armature=armature, 
                                   dofs_idx_local=np.arange(0, self.robot.n_dofs), 
                                   envs_idx=envs_idx)

    def curriculum_commands(self):
        self.curriculum_step += 1
        if self.curriculum_step >= self.curriculum_cfg["curriculum_step"]:
            self.curriculum_step = 0
            # 更新误差
            self.mean_lin_vel_error = (self.lin_vel_error/self.episode_length_buf).mean().item()
            self.mean_ang_vel_error = (self.ang_vel_error/self.episode_length_buf).mean().item()
            # print("lin_vel_error: ",lin_vel_error)
            # print("ang_vel_error: ",ang_vel_error)
            # 调整线速度
            lin_err_high = 999
            if self.curriculum_cfg["err_mode"]:
                self.linx_range_up_threshold = self.curriculum_cfg["lin_vel_err_range"][0]
                lin_err_high = self.curriculum_cfg["lin_vel_err_range"][1]
            else:
                mean_linx_range = self.command_ranges[:, 0, 1].mean()
                ratio = mean_linx_range / self.command_cfg["lin_vel_x_range"][1]
                self.linx_range_up_threshold = self.curriculum_cfg["lin_vel_err_range"][0]+(self.curriculum_cfg["lin_vel_err_range"][1]-self.curriculum_cfg["lin_vel_err_range"][0])*ratio
                lin_err_high = self.curriculum_cfg["lin_vel_err_range"][2]
                
            if self.survive_ratio > 0.5:  
                if self.mean_lin_vel_error < self.linx_range_up_threshold:
                    self.command_ranges[:, 0, 0] += self.curriculum_cfg["curriculum_lin_vel_step"]*self.command_cfg["lin_vel_x_range"][0]
                    self.command_ranges[:, 0, 1] += self.curriculum_cfg["curriculum_lin_vel_step"]*self.command_cfg["lin_vel_x_range"][1]
                elif self.mean_lin_vel_error > lin_err_high:
                    self.command_ranges[:, 0, 0] -= self.curriculum_cfg["curriculum_lin_vel_step"]*self.command_cfg["lin_vel_x_range"][0]
                    self.command_ranges[:, 0, 1] -= self.curriculum_cfg["curriculum_lin_vel_step"]*self.command_cfg["lin_vel_x_range"][1]
                self.command_ranges[:,0,0] = torch.clamp(self.command_ranges[:,0,0],
                                                        self.command_cfg["lin_vel_x_range"][0],
                                                        self.curriculum_cfg["curriculum_lin_vel_min_range"] * self.command_cfg["lin_vel_x_range"][0])
                self.command_ranges[:,0,1] = torch.clamp(self.command_ranges[:,0,1],
                                                        self.curriculum_cfg["curriculum_lin_vel_min_range"] * self.command_cfg["lin_vel_x_range"][1],
                                                        self.command_cfg["lin_vel_x_range"][1])
            #角度
            angv_err_high = 999
            if self.curriculum_cfg["err_mode"]:
                self.angv_range_up_threshold = self.curriculum_cfg["ang_vel_err_range"][0]
                angv_err_high = self.curriculum_cfg["ang_vel_err_range"][1]
            else:
                mean_angv_range = self.command_ranges[:, 2, 1].mean()  
                ratio = mean_angv_range / self.command_cfg["ang_vel_range"][1]
                self.angv_range_up_threshold = self.curriculum_cfg["ang_vel_err_range"][0]+(self.curriculum_cfg["ang_vel_err_range"][1]-self.curriculum_cfg["ang_vel_err_range"][0])*ratio
                angv_err_high = self.curriculum_cfg["ang_vel_err_range"][2]

            if self.survive_ratio > 0.5:
                if self.mean_ang_vel_error < self.angv_range_up_threshold:
                    self.command_ranges[:, 2, 0] += self.curriculum_cfg["curriculum_ang_vel_step"]*self.command_cfg["ang_vel_range"][0]
                    self.command_ranges[:, 2, 1] += self.curriculum_cfg["curriculum_ang_vel_step"]*self.command_cfg["ang_vel_range"][1]
                elif self.mean_ang_vel_error > angv_err_high:
                    self.command_ranges[:, 2, 0] -= self.curriculum_cfg["curriculum_ang_vel_step"]*self.command_cfg["ang_vel_range"][0]
                    self.command_ranges[:, 2, 1] -= self.curriculum_cfg["curriculum_ang_vel_step"]*self.command_cfg["ang_vel_range"][1]
                self.command_ranges[:,2,0] = torch.clamp(self.command_ranges[:,2,0],
                                                        self.command_cfg["ang_vel_range"][0],
                                                        self.curriculum_cfg["curriculum_ang_vel_min_range"] * self.command_cfg["ang_vel_range"][0])
                self.command_ranges[:,2,1] = torch.clamp(self.command_ranges[:,2,1],
                                                        self.curriculum_cfg["curriculum_ang_vel_min_range"] * self.command_cfg["ang_vel_range"][1],
                                                        self.command_cfg["ang_vel_range"][1])

    '''正金字塔楼梯'''
    def add_pyramid(self,point):
        max_z_pos = self.v_stairs_num * self.v_stairs_height - self.v_stairs_height / 2 + point[2]
        current_z = max_z_pos
        box_size = self.v_plane_size
        box_pos = point
        for _ in range(self.v_stairs_num):
            box_pos[2] = current_z
            box = self.scene.add_entity(
                morph=gs.morphs.Box(
                    pos=tuple(box_pos),
                    size=(box_size, box_size, self.v_stairs_height),
                    fixed=True
                )
            )
            box_size += self.v_stairs_width*2
            current_z -= self.v_stairs_height
            size = self.v_plane_size + self.v_stairs_width*2* (self.v_stairs_num-1) 
        return size

    '''倒金字塔楼梯'''
    def add_inverted_pyramid(self,point):
        min_z_pos = self.v_stairs_height / 2 + point[2]
        box_offset = (self.v_plane_size + self.v_stairs_width)/2
        box_length = self.v_plane_size + self.v_stairs_width*2
        box_pos = [[point[0]+box_offset, point[1], point[2]+self.v_stairs_height/2],
                [point[0]-box_offset, point[1], point[2]+self.v_stairs_height/2],
                [point[0], point[1]+box_offset, point[2]+self.v_stairs_height/2],
                [point[0], point[1]-box_offset, point[2]+self.v_stairs_height/2]]
        box_size = [[self.v_stairs_width, box_length, self.v_stairs_height],
                    [self.v_stairs_width, box_length, self.v_stairs_height],
                    [box_length, self.v_stairs_width, self.v_stairs_height],
                    [box_length, self.v_stairs_width, self.v_stairs_height]]
        box = self.scene.add_entity(
                    morph=gs.morphs.Box(
                        pos=(point[0], point[1], min_z_pos),
                        size=(self.v_plane_size,self.v_plane_size,self.v_stairs_height),
                        fixed=True
                    )
                )
        for num_stairs in range(self.v_stairs_num*2):
            for i in range(4):
                if num_stairs<=self.v_stairs_num-1:
                    box_pos[i][2] += self.v_stairs_height
                else:
                    box_pos[i][2] -= self.v_stairs_height
                box = self.scene.add_entity(
                    morph=gs.morphs.Box(
                        pos=tuple(box_pos[i]),
                        size=tuple(box_size[i]),
                        fixed=True
                    )
                )
            box_pos[0][0] += self.v_stairs_width
            box_pos[1][0] -= self.v_stairs_width
            box_pos[2][1] += self.v_stairs_width
            box_pos[3][1] -= self.v_stairs_width
            box_size[3][0] += self.v_stairs_width * 2
            box_size[0][1] = box_size[1][1] =box_size[2][0]= box_size[3][0]
        size = self.v_plane_size + self.v_stairs_width * 2 * (self.v_stairs_num*2) 
        return size

    def get_relative_terrain_pos(self, base_pos):
        if not self.terrain_cfg["terrain"]:
            return base_pos
        #对多个 (x, y) 坐标进行双线性插值计算地形高度
        # 提取x和y坐标
        x = base_pos[:, 0]
        y = base_pos[:, 1]
        # 转换为浮点数索引
        fx = x / self.horizontal_scale
        fy = y / self.horizontal_scale
        # 获取四个最近的整数网格点，确保在有效范围内
        x0 = torch.floor(fx).int()
        x1 = torch.min(x0 + 1, torch.full_like(x0, self.terrain_height.shape[1] - 1))
        y0 = torch.floor(fy).int()
        y1 = torch.min(y0 + 1, torch.full_like(y0, self.terrain_height.shape[0] - 1))
        # 确保x0, x1, y0, y1在有效范围内
        x0 = torch.clamp(x0, 0, self.terrain_height.shape[1] - 1)
        x1 = torch.clamp(x1, 0, self.terrain_height.shape[1] - 1)
        y0 = torch.clamp(y0, 0, self.terrain_height.shape[0] - 1)
        y1 = torch.clamp(y1, 0, self.terrain_height.shape[0] - 1)
        # 获取四个点的高度值
        # 使用广播机制处理批量数据
        Q11 = self.terrain_height[y0, x0]
        Q21 = self.terrain_height[y0, x1]
        Q12 = self.terrain_height[y1, x0]
        Q22 = self.terrain_height[y1, x1]
        # 计算双线性插值
        wx = fx - x0
        wy = fy - y0
        height = (
            (1 - wx) * (1 - wy) * Q11 +
            wx * (1 - wy) * Q21 +
            (1 - wx) * wy * Q12 +
            wx * wy * Q22
        )
        base_pos[:,2] -= height
        return base_pos


    # ------------ reward functions----------------
    def _reward_tracking_lin_x_vel(self):
        # Tracking of linear velocity commands (x axes)
        # 二次函数跟踪速度和角速度会导致正奖励太低，提高梯度也很难跟进
        lin_vel_error = torch.square(self.commands[:, 0] - self.base_lin_vel[:, 0])
        lin_vel_reward = torch.exp(-lin_vel_error / self.reward_cfg["tracking_linx_sigma"])
        if self.command_cfg["zero_stable"]:
            near_zero_mask = (self.commands[:, 0] >= -0.01) & (self.commands[:, 0] <= 0.01)
            if torch.any(near_zero_mask):
                second_error = torch.square(self.commands[near_zero_mask, 0] - self.base_lin_vel[near_zero_mask, 0])
                second_reward = torch.exp(-second_error / self.reward_cfg["tracking_linx_sigma"])
                lin_vel_reward[near_zero_mask] += second_reward
        return lin_vel_reward

    def _reward_tracking_lin_y_vel(self):
        # Tracking of linear velocity commands (y axes)
        # 当存在横向移动命令时鼓励横向髋关节变化
        lin_vel_error = torch.square(self.commands[:, 1] - self.base_lin_vel[:, 1])
        return torch.exp(-lin_vel_error / self.reward_cfg["tracking_liny_sigma"])

    def _reward_tracking_ang_vel(self):
        # Tracking of angular velocity commands (yaw)
        ang_vel_error = torch.square(self.commands[:, 2] - self.base_ang_vel[:, 2])
        ang_vel_reward = torch.exp(-ang_vel_error / self.reward_cfg["tracking_ang_sigma"])
        if self.command_cfg["zero_stable"]:
            near_zero_mask = (self.commands[:, 0] >= -0.01) & (self.commands[:, 0] <= 0.01)
            if torch.any(near_zero_mask):
                second_error = torch.square(self.commands[near_zero_mask, 2] - self.base_lin_vel[near_zero_mask, 2])
                second_reward = torch.exp(-second_error / self.reward_cfg["tracking_ang_sigma"])
                ang_vel_reward[near_zero_mask] += second_reward
        return ang_vel_reward

    def _reward_tracking_leg_length(self):
        # 身高跟踪 建议用高斯函数
        # base_height_error = torch.square(self.base_pos[:, 2] - self.commands[:, 3])
        # return torch.exp(-base_height_error / self.reward_cfg["tracking_height_sigma"])
        # 膝关节跟踪 建议用二次函数
        # knee_error = torch.square(self.dof_pos[:, [2, 5]] - self.commands[:, 3].unsqueeze(1)).sum(dim=1)
        # return torch.exp(-knee_error / self.reward_cfg["tracking_height_sigma"])
        # return knee_error
        # 髋关节跟踪 建议用二次函数
        knee_error = torch.square(self.dof_pos[:, 1] - self.commands[:, 3])
        knee_error += torch.square(self.dof_pos[:, 4] - self.commands[:, 4])
        return knee_error
        # 脚据base距离
        # base_height_error = torch.abs(self.left_foot_base_pos[:, 2] - self.commands[:, 3]) 
        # + torch.abs(self.right_foot_base_pos[:, 2] - self.commands[:, 3])
        # base_height_error = torch.square(base_height_error)
        # print("base_height_error: ", base_height_error)
        # return torch.exp(-base_height_error / self.reward_cfg["tracking_height_sigma"])
        # return base_height_error

    def _reward_lin_vel_z(self):
        # Penalize z axis base linear velocity
        return torch.square(self.base_lin_vel[:, 2])
    
    def _reward_joint_action_rate(self):
        # Penalize changes in actions
        joint_action_rate = self.last_actions[:,self.joint_dof_idx_np] - self.actions[:,self.joint_dof_idx_np]
        return torch.sum(torch.square(joint_action_rate), dim=1)

    def _reward_wheel_action_rate(self):
        # Penalize changes in actions
        wheel_action_rate = self.last_actions[:,self.wheel_dof_idx_np] - self.actions[:,self.wheel_dof_idx_np]
        return torch.sum(torch.square(wheel_action_rate), dim=1)
        # return torch.log(torch.square(wheel_action_rate)+1).sum(dim=1)

    # def _reward_similar_to_default(self):
    #     # Penalize joint poses far away from default pose
    #     #个人认为为了灵活性这个作用性不大
    #     return torch.sum(torch.abs(self.dof_pos[:,self.joint_dof_idx_np] - self.default_dof_pos[:,self.joint_dof_idx_np]), dim=1)

    def _reward_projected_gravity(self):
        #保持水平奖励使用重力投影 0 0 -1
        #使用e^(-x^2)效果不是很好
        # projected_gravity_error = 1 + self.projected_gravity[:, 2] #[0, 0.2]
        # projected_gravity_error = torch.square(projected_gravity_error)
        # return torch.exp(-projected_gravity_error / self.reward_cfg["tracking_gravity_sigma"])
        #二次函数
        reward = torch.sum(torch.square(self.projected_gravity[:, :2]), dim=1)
        return reward
    
    def _reward_similar_calf(self):
        rew = torch.square(self.dof_pos[:, 2] - self.dof_pos[:, 5])
        return rew

    def _reward_joint_vel(self):
        # Penalize dof velocities
        return torch.sum(torch.square(self.dof_vel[:, self.joint_dof_idx_np]), dim=1)

    def _reward_dof_acc(self):
        dof_acc = (self.dof_vel - self.last_dof_vel)/self.dt
        return torch.sum(torch.square(dof_acc), dim=1)

    def _reward_dof_force(self):
        return torch.sum(torch.square(self.dof_force), dim=1)

    def _reward_ang_vel_xy(self):
        # Penalize xy axes base angular velocity
        return torch.sum(torch.square(self.base_ang_vel[:, :2]), dim=1)

    def _reward_collision(self):
        # 接触地面惩罚 力越大惩罚越大
        collision = torch.zeros(self.num_envs,device=self.device,dtype=gs.tc_float)
        for idx in self.reset_links:
            collision += torch.square(self.connect_force[:,idx,:]).sum(dim=1)
        return collision
    
    def _reward_feet_distance(self):
        # 两腿间距
        feet_distance = torch.norm(self.left_foot_pos - self.right_foot_pos, dim=-1)
        reward = torch.clip(self.reward_cfg["feet_distance"][0] - feet_distance, 0, 1) + \
                 torch.clip(feet_distance - self.reward_cfg["feet_distance"][1], 0, 1)
        return reward
    
    def _reward_survive(self):
        # 存活的给奖励 return / episode_length = reward(显示)
        return torch.ones(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False)
    
    def _reward_tsk(self):
        # 铁山靠 hip_pos
        tsk_err = self.dof_pos[:,0] - self.commands[:, 5]
        tsk_err += self.dof_pos[:,3] - self.commands[:, 5]
        return torch.square(tsk_err)
    