[software] 添加16DOF早期训练仿真与Sim2Real闭环

2026-07-21 16:15:14 +08:00
parent 9bd22225f9
commit e9e2c946b3
681 changed files with 137221 additions and 8 deletions
@@ -0,0 +1,108 @@
+import time
+import torch
+import numpy as np
+import mujoco.viewer
+from pathlib import Path
+
+# 导入我们的模块化组件
+from input_dev.keyboard import KeyboardCommandController
+from policy.policy_runner import PolicyRunner
+from interface.mujoco_io import MuJoCoIO
+from tools.logger import SimpleLogger
+
+def main():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # 路径设置 (兼容 rc_mjlab/sim2sim 目录结构)
+    project_root = Path(__file__).parent.parent.absolute()
+    terrain_dir = Path(__file__).parent / "terrain"
+    terrain_xml = terrain_dir / "scene_terrain.xml"
+    robot_xml = project_root / "mjcf" / "wheelleg.xml"
+    policy_path = {
+        "rough": project_root / "model_rough.pt",
+        "crawl": project_root / "model_crawl.pt"
+    }
+
+    # 1. 初始化 MuJoCo IO 接口
+    print("\n[Main] Initializing MuJoCo Environment...")
+    io = MuJoCoIO(terrain_xml, robot_xml, terrain_dir)
+    
+    # 2. 初始化 Policy 推理层
+    print("\n[Main] Initializing Policy Runner...")
+    runner = PolicyRunner(policy_path, device)
+    
+    # 3. 初始化键盘控制器 (带平滑加减速)
+    print("\n[Main] Initializing Input Controller...")
+    kb = KeyboardCommandController(max_x_vel=1.0, max_yaw_vel=1.0)
+    kb.start()
+    
+    # 4. 初始化数据日志记录器
+    print("\n[Main] Initializing Data Logger...")
+    logger = SimpleLogger(log_dir=str(project_root / "sim2sim"))
+
+    # 机器人复位
+    io.reset_robot(runner.default_dof_pos)
+    runner.reset()
+
+    # 时序控制计算
+    control_dt = io.control_dt  # 通常是 0.02s (50Hz)
+    sim_steps_per_control = int(round(control_dt / io.m.opt.timestep)) # 通常是 10
+    
+    next_exec_time = time.perf_counter()
+    viewer_counter = 0
+
+    print(f"\n[Main] Starting Control Loop (Control DT: {control_dt:.3f}s)")
+    
+    try:
+        with mujoco.viewer.launch_passive(io.m, io.d) as viewer:
+            # 初始相机视角
+            viewer.cam.distance = 5.0
+            viewer.cam.elevation = -20.0
+            viewer.cam.azimuth = 45.0
+            
+            while viewer.is_running():
+                # [1] 获取用户指令
+                command = kb.get_command()
+                
+                # [2] 读取环境观测值
+                obs = io.get_obs_53d(command, runner.default_dof_pos, runner.last_actions)
+                
+                # [3] 神经网络推理 (包含历史堆叠处理)
+                scaled_actions, raw_actions = runner.step(obs)
+                
+                # [4] 下发动作到仿真器 (包含低通滤波)
+                io.send_actions(scaled_actions, runner.default_dof_pos)
+                
+                # [5] 记录日志
+                logger.update(io.d.time, io.m, io.d, command)
+                
+                # [6] 推进物理仿真
+                for _ in range(sim_steps_per_control):
+                    mujoco.mj_step(io.m, io.d)
+                
+                # [7] 降低渲染频率以节省性能 (25Hz 渲染)
+                viewer_counter += 1
+                if viewer_counter >= 2:
+                    base_id = mujoco.mj_name2id(io.m, mujoco.mjtObj.mjOBJ_BODY, "base_link")
+                    if base_id != -1:
+                        viewer.cam.lookat[:] = io.d.xpos[base_id]
+                    viewer.sync()
+                    viewer_counter = 0
+
+                # [8] 高精度时序锁帧 (完全对标真机 RTOS 逻辑)
+                next_exec_time += control_dt
+                now = time.perf_counter()
+                sleep_time = next_exec_time - now
+                
+                if sleep_time > 0:
+                    time.sleep(sleep_time)
+                elif sleep_time < -control_dt:
+                    # 如果发生严重掉帧，重置时钟，避免疯狂快进
+                    next_exec_time = now
+    finally:
+        print("\n[Main] Shutting down...")
+        kb.stop()
+        logger.save()  # 保存数据到 txt
+
+if __name__ == "__main__":
+    main()