import sys
from legged_gym import LEGGED_GYM_ROOT_DIR
import os

import isaacgym
from isaacgym import gymapi
from legged_gym.envs import *
from legged_gym.utils import get_args, export_policy_as_jit, export_policy_as_onnx, task_registry, Logger

import numpy as np
import torch
from termcolor import colored
import hydra
from omegaconf import DictConfig, OmegaConf
from easydict import EasyDict
from legged_gym.utils.helpers import class_to_dict
import xml.etree.ElementTree as ET

NOROSPY = False
try:
import rospy
except:
NOROSPY = True

from std_msgs.msg import String, Header, Float64MultiArray

command_state = {
'vel_forward': 0.0,
'vel_side': 0.0,
'orientation': 0.0,
}

override = False

EXPORT_ONNX = True

def dict_compare(d1, d2):
d1_keys = set(d1.keys())
d2_keys = set(d2.keys())
shared_keys = d1_keys.intersection(d2_keys)
added = d1_keys - d2_keys
removed = d2_keys - d1_keys
modified = {o : (d1[o], d2[o]) for o in shared_keys if d1[o] != d2[o]}
same = set(o for o in shared_keys if d1[o] == d2[o])
return added, removed, modified, same
def change_joint_types_to_fixed(urdf_path, joint_names, tmp_urdf_path):
# 读取并解析URDF文件
tree = ET.parse(urdf_path)
root = tree.getroot()

text
joint_names_set = set(joint_names)
found_joints = set()

# 查找所有joint元素
for joint in root.findall('joint'):
    joint_name = joint.get('name')
    if joint_name in joint_names_set:
        joint.set('type', 'fixed')
        found_joints.add(joint_name)
        print(f"Joint '{joint_name}' changed to fixed.")

# 检查是否有未找到的joint
not_found_joints = joint_names_set - found_joints
for missing_joint in not_found_joints:
    print(f"Joint '{missing_joint}' not found!")

# 保存修改后的URDF文件
tree.write(tmp_urdf_path, encoding='utf-8', xml_declaration=True)

@hydra.main(
version_base=None,
config_path="../cfg",
config_name="config_base",
)
def play(cfg_hydra: DictConfig) -> None:
cfg_hydra = EasyDict(OmegaConf.to_container(cfg_hydra, resolve=True))
cfg_hydra.physics_engine = gymapi.SIM_PHYSX

text
if "fixed_joints" in cfg_hydra:  #todo xmh 只有teacher的情况
    fixed_joints = cfg_hydra.fixed_joints
    
    num_fixed_joints = len(fixed_joints)
    cfg_hydra.num_fixed = num_fixed_joints
    
    asset_urdf_path = cfg_hydra.asset.file
    tmp_output_path  = os.path.dirname(asset_urdf_path)
    tmp_urdf_path = os.path.join(tmp_output_path, "tmp.urdf")
    change_joint_types_to_fixed(asset_urdf_path, fixed_joints, tmp_urdf_path)
    
    #cfg_hydra.asset.file = "/home/sjtu_smartcar/Desktop/xmh/robotics/human2humanoid/resources/robots/ym_humanoid/urdf/ym_new_origin_fixed.urdf"
    #cfg_hydra.asset.file = "/home/sjtu_smartcar/Desktop/xmh/robotics/human2humanoid/resources/robots/ym_humanoid/urdf/ym_new_origin_fixed_body.urdf"
    cfg_hydra.asset.file = tmp_urdf_path
    fixed_link_names = [link.replace('Joint', 'Link') for link in fixed_joints]
    if "Body_Link2" in fixed_link_names:
        cfg_hydra.domain_rand.randomize_base_com = False
    
    not_randomize_link_body_names_count = 0 
    for fixed_link_name in fixed_link_names:
        if fixed_link_name in cfg_hydra.domain_rand.randomize_link_body_names:
            cfg_hydra.domain_rand.randomize_link_body_names.remove(fixed_link_name)
            not_randomize_link_body_names_count += 1
        if not cfg_hydra.train.distill:
            cfg_hydra.motion.teleop_selected_keypoints_names.remove(fixed_link_name)
        else:
            cfg_hydra.train.distill_model_config.teleop_selected_keypoints_names.remove(fixed_link_name)
    
    cfg_hydra.env.num_actions -= num_fixed_joints

    if not cfg_hydra.train.distill:
        cfg_hydra.env.num_observations -= num_fixed_joints * 40
        cfg_hydra.env.num_privileged_obs = cfg_hydra.env.num_privileged_obs - num_fixed_joints * 40 - 3* num_fixed_joints - not_randomize_link_body_names_count #todo xmh self._ground_friction_values = num_bodies
    
    cfg_hydra.motion.num_dof_pos_reference -= num_fixed_joints
    cfg_hydra.motion.num_dof_vel_reference -= num_fixed_joints
    
    joint_list = list(cfg_hydra.init_state.default_joint_angles.keys())
    fix_dof_index = []
    fix_pos_index = []
    for fixed_joint in fixed_joints:
        fix_pos_index.append(joint_list.index(fixed_joint) + 1)
        fix_dof_index.append(joint_list.index(fixed_joint))
        cfg_hydra.rewards.teleop_joint_pos_selection.pop(fixed_joint)
    cfg_hydra.motion.visualize_config.marker_joint_colors = [item for idx, item in enumerate(cfg_hydra.motion.visualize_config.marker_joint_colors) if idx not in fix_pos_index]

    cfg_hydra.fix_dof_indices = fix_dof_index
    cfg_hydra.fix_pos_indices = fix_pos_index
    cfg_hydra.asset.clip_motion_goal = False  #xmh
    
    
    
    cfg_hydra.leg_joint1_indice = 18 - num_fixed_joints #11
    
    default_selected_keypoints = [9, 11, 18]
    filtered_list1 = [x for x in default_selected_keypoints if x not in fix_pos_index]
    cfg_hydra.selected_keypoints = [x - sum(1 for y in fix_pos_index if y < x) for x in filtered_list1]
    #todo 去掉domainrand和noise
    num_selected_keypoints = len(cfg_hydra.selected_keypoints)
    if cfg_hydra.train.distill_model_config:
        cfg_hydra.train.distill_model_config.num_observations = 1236 - num_fixed_joints * 40  
        cfg_hydra.train.distill_model_config.num_privileged_obs  = 1353 - num_fixed_joints * 40 - 3* num_fixed_joints - not_randomize_link_body_names_count
        cfg_hydra.env.num_observations = num_selected_keypoints * (3 + 3 + 3) + (30 - num_fixed_joints) * 3 + 3 +3 + cfg_hydra.env.short_history_length* (96 - 3*cfg_hydra.num_fixed)
        cfg_hydra.env.num_privileged_obs = cfg_hydra.env.num_observations + 3+ 2 +6 +1 + len(cfg_hydra.domain_rand.randomize_link_body_names) + 3*(30-num_fixed_joints)

# env_cfg, train_cfg = task_registry.get_cfgs(name=cfg_hydra.task)
# import ipdb; ipdb.set_trace()

env_cfg, train_cfg = cfg_hydra, cfg_hydra.train
train_cfg.seed = 123145

##### Compare two configs. 
# env_cfg_, train_cfg_prev = task_registry.get_cfgs(name=cfg_hydra.task)
# env_cfg_, train_cfg_prev = class_to_dict(env_cfg_), class_to_dict(train_cfg_prev)
# for k, v in env_cfg_.items():
#     if isinstance(v, dict):
#         for kk, vv in v.items():
#             if not vv == env_cfg[k][kk]:
#                 print(k, kk)
#                 import ipdb; ipdb.set_trace()
#                 print('...')
    
#     elif not v == env_cfg[k]:
#         import ipdb; ipdb.set_trace()
#         print('...')

# override some parameters for testing
# env_cfg.env.num_envs = min(env_cfg.env.num_envs, 100)

# if not env_cfg.train_velocity_estimation:
env_cfg.env.num_envs = 1
env_cfg.viewer.debug_viz = True
env_cfg.motion.visualize = False
# env_cfg.terrain.num_rows = 5
# env_cfg.terrain.num_cols = 5
env_cfg.terrain.curriculum = False
env_cfg.terrain.mesh_type = 'trimesh'
# env_cfg.terrain.mesh_type = 'plane'
# if env_cfg.terrain.mesh_type == 'trimesh':
#     env_cfg.terrain.terrain_types = ['flat', 'rough', 'low_obst']  # do not duplicate!
#     env_cfg.terrain.terrain_proportions = [1.0, 0.0, 0.0]
env_cfg.noise.add_noise = False
env_cfg.domain_rand.randomize_friction = False
env_cfg.domain_rand.push_robots = False
env_cfg.domain_rand.randomize_base_mass = False
env_cfg.env.episode_length_s = 20
env_cfg.domain_rand.randomize_rfi_lim = False
env_cfg.domain_rand.randomize_pd_gain = False
env_cfg.domain_rand.randomize_link_mass = False
env_cfg.domain_rand.randomize_base_com = False
env_cfg.domain_rand.randomize_ctrl_delay = False
env_cfg.domain_rand.ctrl_delay_step_range = [1, 3]



# env_cfg.asset.termination_scales.max_ref_motion_distance = 1



env_cfg.env.test = True

if env_cfg.motion.realtime_vr_keypoints:
    env_cfg.asset.terminate_by_1time_motion = False
    env_cfg.asset.terminate_by_ref_motion_distance = False
    rospy.init_node("avppose_subscriber")
    from avp_pose_subscriber import AVPPoseInfo
    avpposeinfo = AVPPoseInfo()
    rospy.Subscriber("avp_pose", Float64MultiArray, avpposeinfo.avp_callback, queue_size=1)
if cfg_hydra.joystick:
    env_cfg.commands.ranges.lin_vel_x = [0.0, 0.0]
    env_cfg.commands.ranges.lin_vel_y = [0.0, 0.0]
    env_cfg.commands.ranges.ang_vel_yaw = [0.0, 0.0]
    from pynput import keyboard
    from legged_gym.utils import key_response_fn

# prepare environment

env, _ = task_registry.make_env_hydra(name=cfg_hydra.task, hydra_cfg=cfg_hydra, env_cfg=env_cfg)

logger = Logger(env.dt)
robot_index = 0 # which robot is used for logging
joint_index = 3 # which joint is used for logging
stop_state_log = 200 # number of steps before plotting states
stop_rew_log = env.max_episode_length + 1 # number of steps before print average episode rewards


obs = env.get_observations()

if env_cfg.motion.realtime_vr_keypoints:
    init_root_pos = env._rigid_body_pos[..., 0, :].clone()
    init_avp_pos = avpposeinfo.avp_pose.copy()
    init_root_offset = init_root_pos[0, :2] - init_avp_pos[2, :2]
# import ipdb; ipdb.set_trace()
# obs[:, 9:12] = torch.Tensor([0.5, 0, 0])
# load policy
train_cfg.runner.resume = True


ppo_runner, train_cfg = task_registry.make_alg_runner(env=env, name=cfg_hydra.task, args=cfg_hydra, train_cfg=train_cfg)

policy = ppo_runner.get_inference_policy(device=env.device)
exported_policy_name = str(task_registry.loaded_policy_path.split('/')[-2]) + str(task_registry.loaded_policy_path.split('/')[-1])
print('Loaded policy from: ', task_registry.loaded_policy_path)

# export policy as a jit module (used to run it from C++)
if EXPORT_POLICY:
    path = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name, 'exported', 'policies')
    export_policy_as_jit(ppo_runner.alg.actor_critic, path, exported_policy_name)
    print('Exported policy as jit script to: ', os.path.join(path, exported_policy_name))
if EXPORT_ONNX:
    exported_onnx_name = exported_policy_name.replace('.pt', '.onnx')
    path = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name, 'exported', 'policies')
    export_policy_as_onnx(ppo_runner.alg.actor_critic, path, exported_onnx_name, onnx_num_observations=env_cfg.env.num_observations)
    print('Exported policy as onnx to: ', os.path.join(path, exported_onnx_name))
    

if cfg_hydra.joystick:
    print(colored("joystick on", "green"))
    key_response = key_response_fn(mode='vel')
    def on_press(key):
        global command_state
        try:
            # print(key.char)
            key_response(key, command_state, env)
        except AttributeError:
            pass

    listener = keyboard.Listener(on_press=on_press)
    listener.start()

i = 0

while (not NOROSPY and not rospy.is_shutdown()) or (NOROSPY):
    # for i in range(1000*int(env.max_episode_length)):

    # obs[:, -19:] = 0 # will destroy the performance
    
    actions = policy(obs.detach())
    
    # print(torch.sum(torch.square(env.projected_gravity[:, :2]), dim=1))
    obs, _, rews, dones, infos = env.step(actions.detach())

    if env_cfg.motion.realtime_vr_keypoints:
        avpposeinfo.check()
        keypoints_pos = avpposeinfo.avp_pose.copy()
        keypoints_pos[:, 0] += init_root_offset[0].item()
        keypoints_pos[:, 1] += init_root_offset[1].item()
        # import ipdb; ipdb.set_trace()
        keypoints_vel = avpposeinfo.avp_vel.copy()
        print(keypoints_pos)
        env._update_realtime_vr_keypoints(keypoints_pos, keypoints_vel)
    # print("obs = ", obs)
    # print("actions = ", actions)
    # print()
    # exit()
    if override: 
        obs[:,9] = 0.5
        obs[:,10] = 0.0
        obs[:,11] = 0.0
    
    # overwrite linear velocity - z and angular velocity - xy
    # obs[:, 40] = 0.
    # obs[: 41:43] = 0.
    

    if i < stop_state_log:
        logger.log_states(
            {
                'dof_pos_target': actions[robot_index, joint_index].item() * env.cfg.control.action_scale + env.default_dof_pos[robot_index, joint_index].item(),
                # 'dof_pos_target': env.actions[robot_index, joint_index].item() * env.cfg.control.action_scale + env.default_dof_pos[robot_index, joint_index].item(),
                'dof_pos': env.dof_pos[robot_index, joint_index].item(),
                'dof_vel': env.dof_vel[robot_index, joint_index].item(),
                'dof_torque': env.torques[robot_index, joint_index].item(),
                'base_vel_x': env.base_lin_vel[robot_index, 0].item(),
                'base_vel_y': env.base_lin_vel[robot_index, 1].item(),
                'base_vel_z': env.base_lin_vel[robot_index, 2].item(),
                'base_vel_yaw': env.base_ang_vel[robot_index, 2].item(),
                'contact_forces_z': env.contact_forces[robot_index, env.feet_indices, 2].cpu().numpy()
            }
        )
    elif i==stop_state_log:
        logger.plot_states()
    if  0 < i < stop_rew_log:
        if infos["episode"]:
            num_episodes = torch.sum(env.reset_buf).item()
            if num_episodes>0:
                logger.log_rewards(infos["episode"], num_episodes)
    elif i==stop_rew_log:
        pass
        # logger.print_rewards()
    i += 1

if name == 'main':
EXPORT_POLICY = False
RECORD_FRAMES = False
MOVE_CAMERA = False
# args = get_args()
play()我想要实现遍历pkl的所有motion 如果有出现 reference motion distance > 0.5m的情况就把这个motionid记录下来怎么做

def _resample_motion_times(self, env_ids):
if len(env_ids) == 0:
return
# self.motion_ids[env_ids] = self._motion_lib.sample_motions(len(env_ids))
# self.motion_ids[env_ids] = torch.randint(0, self._motion_lib._num_unique_motions, (len(env_ids),), device=self.device)
# print(self.motion_ids[:10])
self.motion_len[env_ids] = self._motion_lib.get_motion_length(self.motion_ids[env_ids])
# self.env_origins_init_3Doffset[env_ids, :2] = torch_rand_float(-1., 1., (len(env_ids), 2), device=self.device) # xy position within 1m of the center
if self.cfg.env.test:
self.motion_start_times[env_ids] = 0
else:
self.motion_start_times[env_ids] = self._motion_lib.sample_time(self.motion_ids[env_ids])
# self.motion_start_times[env_ids] = self._motion_lib.sample_time(self.motion_ids[env_ids])
offset=(self.env_origins + self.env_origins_init_3Doffset)
motion_times = (self.episode_length_buf ) * self.dt + self.motion_start_times # next frames so +1
# motion_res = self._get_state_from_motionlib_cache(self.motion_ids, motion_times, offset= offset)
motion_res = self._get_state_from_motionlib_cache_trimesh(self.motion_ids, motion_times, offset= offset)

text
    self.ref_base_pos_init[env_ids] = motion_res["root_pos"][env_ids]
    self.ref_base_rot_init[env_ids] = motion_res["root_rot"][env_ids]
    self.ref_base_vel_init[env_ids] = motion_res["root_vel"][env_ids]
    self.ref_base_ang_vel_init[env_ids] = motion_res["root_ang_vel"][env_ids] 解释这个东西

legged gym training的时候是一个env就是一个motionid吗

在整个训练过程中怎么看envid对应的motionid是多少

self.motion_ids[env_ids]是不是envid当前时刻的motionid

是不是每次leggedgym训练load的motionid都是固定的

def _resample_motion_times(self, env_ids):
if len(env_ids) == 0:
return
self.motion_ids[env_ids] = self._motion_lib.sample_motions(len(env_ids))
self.motion_ids[env_ids] = torch.randint(0, self._motion_lib._num_unique_motions, (len(env_ids),), device=self.device)
# print(self.motion_ids[:10])
self.motion_len[env_ids] = self._motion_lib.get_motion_length(self.motion_ids[env_ids])解释self.motion_ids[env_ids] = self._motion_lib.sample_motions(len(env_ids))
self.motion_ids[env_ids] = torch.randint(0, self._motion_lib._num_unique_motions, (len(env_ids),), device=self.device)
# print(self.motion_ids[:10])是什么

一个n, 3的tensor张量 返回值小于0.5的行的index

一个1，n, 3的tensor张量 返回3那一维度值小于0.5的行的index

一个n，2, 3的tensor张量 返回3那一维度值小于0.5的n行的index