()
| 93 | return single_obs, single_obs_dim |
| 94 | |
| 95 | def main(): |
| 96 | # Load configuration |
| 97 | config = load_config("g1.yaml") |
| 98 | |
| 99 | # Load robot model |
| 100 | m = mujoco.MjModel.from_xml_path(config['xml_path']) |
| 101 | d = mujoco.MjData(m) |
| 102 | m.opt.timestep = config['simulation_dt'] |
| 103 | |
| 104 | # Check number of joints |
| 105 | n_joints = d.qpos.shape[0] - 7 |
| 106 | # print(f"Robot has {n_joints} joints in MuJoCo model") |
| 107 | |
| 108 | # Initialize variables |
| 109 | action = np.zeros(config['num_actions'], dtype=np.float32) |
| 110 | target_dof_pos = config['default_angles'].copy() |
| 111 | cmd = config['cmd_init'].copy() |
| 112 | height_cmd = config['height_cmd'] |
| 113 | |
| 114 | # Initialize observation history |
| 115 | single_obs, single_obs_dim = compute_observation(d, config, action, cmd, height_cmd, n_joints) |
| 116 | obs_history = collections.deque(maxlen=config['obs_history_len']) |
| 117 | for _ in range(config['obs_history_len']): |
| 118 | obs_history.append(np.zeros(single_obs_dim, dtype=np.float32)) |
| 119 | |
| 120 | # Prepare full observation vector |
| 121 | obs = np.zeros(config['num_obs'], dtype=np.float32) |
| 122 | |
| 123 | # Load policy |
| 124 | policy = torch.jit.load(config['policy_path']) |
| 125 | |
| 126 | counter = 0 |
| 127 | |
| 128 | with mujoco.viewer.launch_passive(m, d) as viewer: |
| 129 | start = time.time() |
| 130 | while viewer.is_running() and time.time() - start < config['simulation_duration']: |
| 131 | step_start = time.time() |
| 132 | |
| 133 | # Control leg joints with policy |
| 134 | leg_tau = pd_control( |
| 135 | target_dof_pos, |
| 136 | d.qpos[7:7+config['num_actions']], |
| 137 | config['kps'], |
| 138 | np.zeros_like(config['kps']), |
| 139 | d.qvel[6:6+config['num_actions']], |
| 140 | config['kds'] |
| 141 | ) |
| 142 | |
| 143 | d.ctrl[:config['num_actions']] = leg_tau |
| 144 | |
| 145 | # Keep other joints at zero positions if they exist |
| 146 | if n_joints > config['num_actions']: |
| 147 | arm_kp = 100.0 |
| 148 | arm_kd = 0.5 |
| 149 | arm_target_positions = np.zeros(n_joints - config['num_actions'], dtype=np.float32) |
| 150 | |
| 151 | arm_tau = pd_control( |
| 152 | arm_target_positions, |
no test coverage detected