diff --git a/src/box/simulator.py b/src/box/simulator.py index 866969922e..e803fe7319 100644 --- a/src/box/simulator.py +++ b/src/box/simulator.py @@ -1,4 +1,42 @@ import os + my-featurn-branch +import logging +from typing import Dict, Any, Optional, Tuple, List + +import numpy as np +from gymnasium import spaces + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + +# 全局依赖检测 +try: + import mujoco + HAS_MUJOCO = True +except ImportError: + mujoco = None + HAS_MUJOCO = False + +try: + import mujoco_viewer + HAS_MUJOCO_VIEWER = True +except ImportError: + mujoco_viewer = None + HAS_MUJOCO_VIEWER = False + + +class Simulator: # 第27行：类定义后必须有缩进的代码块 + """Mechanical hand simulator with optional MuJoCo backend. + Behavior: + - If `mujoco` is available, generate a simple hand MJCF, load model/data, + and expose `reset`/`step` using MuJoCo stepping routines. + - If `mujoco` is not available, fall back to a lightweight placeholder + implementation so the module remains importable and testable. + """ + # --------------- 核心：类内所有方法必须缩进（通常4个空格）--------------- + def __init__(self, render_mode: str = "human", simulator_folder: str = "./mujoco_models"): + """初始化模拟器（类的构造方法，必须缩进）""" + self.render_mode = render_mode import numpy as np import pygame import mujoco @@ -9,12 +47,339 @@ class Simulator: def __init__(self, simulator_folder: str, render_mode: Optional[str] = None): + main self.simulator_folder = simulator_folder self.render_mode = render_mode self.step_count = 0 self.terminated = False self.truncated = False self.last_reward = 0.0 + + # ---------------- MuJoCo后端实现（必须缩进）---------------- + def _init_mujoco(self): + """初始化MuJoCo仿真环境""" + self.config: Dict[str, Any] = { + "simulation": { + "max_steps": 1000, + "model_path": "hand_model.mjcf", + "target_joint_pos": [] + } + } + + # 5根手指的基座位置和连杆长度（单位：米） + bases = [(-0.20, 0.06), (-0.08, 0.06), (0.04, 0.06), (0.16, 0.06), (0.28, 0.06)] + lengths = [[0.06, 0.05, 0.04] for _ in bases] + self.finger_bases = bases + self.link_lengths = lengths + + # 生成MJCF模型文件路径 + model_path = os.path.join(self.simulator_folder, self.config["simulation"]["model_path"]) + os.makedirs(os.path.dirname(model_path), exist_ok=True) + + # 构建MJCF模型内容 + mjcf_lines: List[str] = [ + '', + ' ') + + # 保存MJCF文件 + with open(model_path, 'w', encoding='utf-8') as f: + f.write('\n'.join(mjcf_lines)) + + # 加载MuJoCo模型和数据 + try: + self.model = mujoco.MjModel.from_xml_path(model_path) + self.data = mujoco.MjData(self.model) + except Exception as e: + logger.exception(f'Failed to load MuJoCo model — falling back to placeholder: {e}') + self._init_placeholder() + return + + # 定义动作/观测空间 + self.action_space = spaces.Box( + low=-1.0, high=1.0, + shape=(int(self.model.nu),), dtype=np.float32 + ) + obs_dim = int(self.model.nq) + int(self.model.nv) + self.observation_space = spaces.Box( + low=-np.inf, high=np.inf, + shape=(obs_dim,), dtype=np.float32 + ) + + # 初始化MuJoCo Viewer + self.viewer = None + if self.render_mode == 'human' and HAS_MUJOCO_VIEWER: + try: + self.viewer = mujoco_viewer.MujocoViewer(self.model, self.data) + self.viewer.cam.azimuth = 45 + self.viewer.cam.elevation = -20 + self.viewer.cam.distance = 0.5 + except Exception as e: + logger.exception(f'Failed to initialize mujoco_viewer: {e}') + + # ---------------- 占位实现（必须缩进）---------------- + def _init_placeholder(self): + """轻量级占位实现""" + self.model = type('M', (), {'nq': 15, 'nu': 15, 'nv': 15})() + self.data = type('D', (), { + 'qpos': np.zeros(self.model.nq), + 'qvel': np.zeros(self.model.nv) + })() + + # 定义兼容的动作/观测空间 + self.action_space = spaces.Box( + low=-1.0, high=1.0, + shape=(int(self.model.nu),), dtype=np.float32 + ) + obs_dim = int(self.model.nq) + int(self.model.nv) + self.observation_space = spaces.Box( + low=-np.inf, high=np.inf, + shape=(obs_dim,), dtype=np.float32 + ) + + # 占位的手指参数 + self.finger_bases = [(-0.20, 0.06), (-0.08, 0.06), (0.04, 0.06), (0.16, 0.06), (0.28, 0.06)] + self.link_lengths = [[0.06, 0.05, 0.04] for _ in self.finger_bases] + + # ---------------- 通用API（必须缩进）---------------- + def reset(self, seed: Optional[int] = None) -> Tuple[np.ndarray, dict]: + """重置仿真环境""" + if seed is not None: + np.random.seed(int(seed)) + if HAS_MUJOCO and hasattr(mujoco, 'set_rng_seed'): + try: + mujoco.set_rng_seed(int(seed)) + except Exception as e: + logger.warning(f'Failed to set MuJoCo seed: {e}') + + self.step_count = 0 + self.terminated = False + self.truncated = False + self.last_reward = 0.0 + + if HAS_MUJOCO and isinstance(self.data, mujoco.MjData): + mujoco.mj_resetData(self.model, self.data) + else: + self.data.qpos = np.zeros(int(self.model.nq)) + self.data.qvel = np.zeros(int(self.model.nv)) + + obs = self._get_obs() + return obs, {'step': self.step_count} + + def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, dict]: + """执行一步仿真""" + action = np.asarray(action, dtype=np.float32) + action = np.clip(action, self.action_space.low, self.action_space.high) + + # MuJoCo步进 + if HAS_MUJOCO and isinstance(self.data, mujoco.MjData): + self.data.ctrl[:] = action * 5.0 + try: + mujoco.mj_step(self.model, self.data) + except Exception: + try: + mujoco.mj_step1(self.model, self.data) + mujoco.mj_step2(self.model, self.data) + except Exception as e: + logger.exception(f'MuJoCo stepping failed: {e}') + else: + # 占位动力学 + gain = 0.01 + self.data.qvel = self.data.qvel + action[:int(self.model.nv)] * gain + self.data.qpos = self.data.qpos + self.data.qvel + + self.step_count += 1 + reward = self._compute_reward() + self.last_reward = float(reward) + + self.terminated = self._check_terminated() + max_steps = int(self.config['simulation'].get('max_steps', 1000)) if hasattr(self, 'config') else 1000 + self.truncated = self.step_count >= max_steps + + obs = self._get_obs() + info = {'step': self.step_count, 'reward': float(reward)} + return obs, float(reward), bool(self.terminated), bool(self.truncated), info + + def _get_obs(self) -> np.ndarray: + """获取当前观测""" + if HAS_MUJOCO and isinstance(self.data, mujoco.MjData): + qpos = self.data.qpos.copy() + qvel = self.data.qvel.copy() + else: + qpos = np.asarray(self.data.qpos).copy() + qvel = np.asarray(self.data.qvel).copy() + return np.concatenate([qpos, qvel]).astype(np.float32) + + def _compute_reward(self) -> float: + """计算奖励""" + if hasattr(self, 'config') and 'simulation' in self.config: + target = np.array(self.config['simulation'].get('target_joint_pos', [0.0] * int(self.model.nq))) + else: + target = np.zeros(int(self.model.nq)) + + current = np.asarray(self.data.qpos).copy() + if len(target) != len(current): + target = np.zeros_like(current) + + return float(-np.sum((current - target) ** 2)) + + def _check_terminated(self) -> bool: + """检查是否终止""" + return False + + def render(self): + """渲染仿真画面""" + if self.render_mode != 'human': + return + + # MuJoCo Viewer渲染 + if HAS_MUJOCO and self.viewer is not None: + try: + self.viewer.render() + return + except Exception as e: + logger.warning(f'MuJoCo Viewer render failed: {e}') + + # 占位渲染（Pygame） + try: + import pygame + except ImportError: + logger.warning("Pygame not installed — skip placeholder rendering") + return + + # 初始化Pygame + if not (pygame.get_init() and pygame.display.get_init()): + try: + pygame.init() + self.screen = pygame.display.set_mode((800, 600)) + pygame.display.set_caption('Mechanical Hand (Placeholder)') + except Exception as e: + logger.warning(f'Pygame init failed: {e}') + return + + # 绘制背景 + screen = self.screen + screen.fill((240, 240, 255)) + + # 绘制手掌 + center_x, center_y = 400, 420 + palm_w, palm_h = 360, 80 + pygame.draw.rect(screen, (160, 120, 90), + (center_x - palm_w // 2, center_y - palm_h // 2, palm_w, palm_h)) + + # 绘制手指连杆 + q = np.asarray(self.data.qpos).copy() + ptr = 0 + for i, base in enumerate(self.finger_bases): + bx, by = base + sx = center_x + int(bx * 1000) + sy = center_y - int(by * 1000) + angle = 0.0 + x0, y0 = sx, sy + + for j, L in enumerate(self.link_lengths[i]): + if ptr < len(q): + angle += float(q[ptr]) + ex = x0 + int(np.cos(angle) * (L * 1000)) + ey = y0 - int(np.sin(angle) * (L * 1000)) + + pygame.draw.line(screen, (10, 10, 10), (x0, y0), (ex, ey), max(1, 6 - j)) + pygame.draw.circle(screen, (200, 80, 80), (x0, y0), 6) + + x0, y0 = ex, ey + ptr += 1 + + # 绘制状态文本 + font = pygame.font.Font(None, 28) + txt = font.render(f"Step: {self.step_count} Reward: {self.last_reward:.3f}", True, (0, 0, 0)) + screen.blit(txt, (10, 10)) + + # 更新画面 + pygame.display.flip() + + # 处理Pygame事件 + for event in pygame.event.get(): + if event.type == pygame.QUIT: + self.close() + + def close(self): + """关闭模拟器""" + # 关闭MuJoCo Viewer + if HAS_MUJOCO and self.viewer is not None: + try: + self.viewer.close() + except Exception: + pass + + # 关闭Pygame + if self.render_mode == 'human': + try: + import pygame + pygame.quit() + except Exception: + pass + + +# 测试代码（类外代码无需缩进） +if __name__ == "__main__": + # 初始化模拟器 + sim = Simulator(render_mode="human", simulator_folder="./mujoco_hand") + + # 重置环境 + obs, info = sim.reset(seed=42) + print(f"Initial observation shape: {obs.shape}") + + # 运行100步仿真 + for _ in range(100): + # 随机动作 + action = sim.action_space.sample() + obs, reward, terminated, truncated, info = sim.step(action) + + # 渲染画面 + sim.render() + + # 打印状态 + if _ % 10 == 0: + print(f"Step: {info['step']}, Reward: {reward:.3f}") + + if terminated or truncated: + print(f"Simulation ended at step {info['step']}") + break + + # 关闭模拟器 + sim.close() +======= self.model = None self.data = None self.viewer = None @@ -275,4 +640,4 @@ def render(self): def close(self): """关闭环境""" if self.render_mode == "human": - pygame.quit() \ No newline at end of file + pygame.quit() diff --git a/src/box/test.simulator.py b/src/box/test.simulator.py index 233864623b..8850b343da 100644 --- a/src/box/test.simulator.py +++ b/src/box/test.simulator.py @@ -1,5 +1,85 @@ import sys import os +import argparse +import importlib.util +import traceback +import numpy as np + + +def load_simulator_module(path: str): + if not os.path.exists(path): + raise FileNotFoundError(f"simulator.py not found: {path}") + spec = importlib.util.spec_from_file_location("simulator", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def main(): + parser = argparse.ArgumentParser(description="Run a short smoke test of the Simulator") + parser.add_argument("--steps", type=int, default=500, help="Number of simulation steps") + parser.add_argument("--seed", type=int, default=42, help="Random seed for repeatability") + parser.add_argument("--render", action="store_true", help="Enable human rendering (pygame)") + parser.add_argument("--sim-folder", type=str, default=None, help="Simulator folder (overrides default)") + parser.add_argument("--log-interval", type=int, default=50, help="How often to print status") + args = parser.parse_args() + + script_dir = os.path.dirname(os.path.abspath(__file__)) + sim_path = os.path.join(script_dir, "simulator.py") + module = load_simulator_module(sim_path) + Simulator = getattr(module, "Simulator") + + sim_folder = args.sim_folder or os.path.join(script_dir, "simulators", "arm_simulation") + os.makedirs(sim_folder, exist_ok=True) + + print("=" * 50) + print("Mechanical arm simulator smoke test") + print(f"simulator folder: {sim_folder}") + + # seed for reproducible samples + np.random.seed(args.seed) + + try: + env = Simulator.get(simulator_folder=sim_folder, render_mode=("human" if args.render else None)) + print("Environment created") + print(f"nq={env.model.nq} nu={env.model.nu} nv={env.model.nv}") + + obs, info = env.reset(seed=args.seed) + print(f"initial obs shape: {obs.shape}") + + for step in range(args.steps): + action = env.action_space.sample() + obs, reward, terminated, truncated, info = env.step(action) + + if step % args.log_interval == 0: + print(f"Step {step:4d} | reward={reward:.3f} terminated={terminated} truncated={truncated}") + + if terminated or truncated: + print(f"Terminated at step {step}, resetting environment") + obs, info = env.reset() + + print("Simulation finished") + + except KeyboardInterrupt: + print("Interrupted by user") + except Exception as e: + print(f"Error: {type(e).__name__}: {e}") + traceback.print_exc() + finally: + try: + if 'env' in locals(): + env.close() + except Exception: + pass + + +if __name__ == "__main__": + main() +import sys +import os + +main import importlib.util import numpy as np @@ -83,7 +163,6 @@ traceback.print_exc() finally: - # 关闭环境 print("\n关闭仿真环境...") try: if 'env' in locals():