diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..4dd9a82 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,84 @@ +# AKA-Sim2Real + +基于前视视角模拟器的数据采集与 ACT (Action Chunking Transformer) 模型训练系统。 + +## 🛠️ 技术栈 + +- **后端**: FastAPI, Socket.IO, PyTorch +- **前端**: React, TypeScript, Socket.IO Client +- **文档**: mdbook + +## 🚫 核心红线 + +- **Python 环境**: 使用 `pip`，确保 `requirements.txt` 依赖完整，引入新依赖更新`requirements.txt` +- **模型层**: `policies/models/act/` 为纯模型定义，禁止在此目录外修改模型结构 +- **API 调用**: 通过显式 runtime 接口调用推理，禁止直接依赖隐式模块状态 +- **测试**: 修改 ACT 相关代码后必须运行 `python3 backend/run_act_checks.py` + +## ⚙️ 常用命令 + +```bash +# 启动后端 +cd backend && python main.py + +# 启动前端 +cd ui && npm run dev + +# 运行测试 +coverage erase && python3 -m pytest tests/ --cov=backend --cov-report=html + +# 打开报告网页 +open htmlcov/index.html +``` + +## 🏗️ 架构与规范 + +### 目录结构 + +``` +├── backend/ # Python 后端 (FastAPI + Socket.IO) +│ ├── api/ # REST API 路由 (domains/inference, training, episode) +│ ├── sio_handlers/ # Socket.IO 事件处理 (sim/real 命名空间) +│ └── services/ +│ ├── inference/ # checkpoint.py, preprocess.py, execution.py, runtime.py +│ ├── training/ # 训练编排 +│ └── episode/ # 数据采集/导出 +├── policies/models/act/ # ACT 模型 (modeling_act, configuration_act, defaults, train_act) +├── ui/ # React 前端 +├── tests/act/ # ACT 测试 +└── docs/ # mdbook 文档 +``` + +### 代码规范 + +## 🧹 代码整洁规范 (Clean Code Principles) + +在编写任何代码时，必须严格遵守以下 Clean Code 原则： + +### 1. 命名即文档 +- **意图清晰**: 变量、函数、类的命名必须揭示其意图。避免使用魔术数字和模糊缩写（如 `d` 代表 `data`，`val` 代表 `value`）。 +- **可搜索性**: 避免使用魔法数字。如果必须使用，请定义为常量（如 `const DAYS_IN_WEEK = 7`）。 +- **类型命名**: 类名和对象名必须是名词（`Customer`, `AccountStatus`），函数名必须是动词（`getUser`, `calculateTotal`）。 + +### 2. 函数设计 +- **单一职责**: 一个函数只做一件事。如果一个函数超过 **20 行**（或屏幕一屏），请尝试拆分。 +- **无副作用**: 尽量避免函数产生副作用（如修改全局变量、隐式修改传入的对象）。如果必须修改，请在命名中体现（如 `updateUser` vs `getUser`）。 +- **命令与查询分离**: 函数要么执行操作（Command），要么返回数据（Query），不要同时做这两件事。 + +### 3. 注释与文档 +- **代码自解释**: 优先通过重构代码（如提取函数、优化变量名）来消除对注释的需求，而不是写注释来解释烂代码。 +- **禁止废话**: 不要写“获取用户 ID”这种显而易见的注释。 +- **解释“为什么”**: 只有在解释复杂的业务逻辑、权衡取舍或解决特定 Bug 时才写注释（解释“为什么这么做”，而不是“做了什么”）。 +- **TODO 规范**: 遗留代码必须标记 `// TODO: [描述] - [作者/日期]`。 + +### 4. 错误处理 +- **使用异常而非返回码**: 不要通过返回 `null` 或 `-1` 来表示错误。使用 `try/catch` 或 Result/Either 模式。 +- **不要返回/传递 null**: 除非是特定语言特性（如 Optional），否则不要返回 `null`。返回空列表 `[]` 或空对象。 +- **精确捕获**: 不要捕获通用的 `Exception`，只捕获你明确知道如何处理的特定错误类型。 + +### 5. 结构与格式 +- **垂直格式**: 相关的代码行应该放在一起（变量声明靠近使用处）。 +- **依赖倒置**: 高层模块不应依赖低层模块，二者都应依赖抽象（接口）。 +- **DRY**: 杜绝重复代码。如果你发现自己在复制粘贴代码，请立即提取为公共函数或组件。 + +--- \ No newline at end of file diff --git a/backend/api/domains/episode/routes.py b/backend/api/domains/episode/routes.py index 4e28bff..234989d 100644 --- a/backend/api/domains/episode/routes.py +++ b/backend/api/domains/episode/routes.py @@ -2,6 +2,7 @@ AKA-Sim 后端 - Episode/数据采集域 API """ +import json import logging from pathlib import Path @@ -27,10 +28,13 @@ async def list_dataset_dirs(user_id: str): datasets = [] for item in user_dataset_path.iterdir(): - if item.is_dir() and (item / "data").exists() or (item / "meta").exists(): - datasets.append(item.name) + if item.is_dir() and ((item / "data").exists() or (item / "meta").exists()): + info_path = item / "meta" / "info.json" + sort_time = info_path.stat().st_mtime if info_path.exists() else item.stat().st_mtime + datasets.append((item.name, sort_time)) - return {"datasets": sorted(datasets)} + datasets.sort(key=lambda entry: entry[1], reverse=True) + return {"datasets": [name for name, _ in datasets]} @router.get("/models") @@ -43,15 +47,52 @@ async def list_models(user_id: str, dataset_name: str = "default"): if not train_path.exists(): return {"models": []} + if dataset_name: + dataset_model_path = train_path / dataset_name + has_model = (dataset_model_path / "model.pt").exists() or (dataset_model_path / "final_model.pt").exists() + if dataset_model_path.is_dir() and has_model: + return {"models": [dataset_name]} + return {"models": []} + # 返回子文件夹名称 models = [] for item in train_path.iterdir(): - if item.is_dir(): + has_model = (item / "model.pt").exists() or (item / "final_model.pt").exists() + if item.is_dir() and has_model: models.append(item.name) return {"models": sorted(models)} +@router.get("/info") +async def get_dataset_info(user_id: str, dataset_name: str): + """读取数据集元信息。""" + if not user_id: + raise HTTPException(status_code=400, detail="user_id is required") + if not dataset_name or ".." in dataset_name or "/" in dataset_name or "\\" in dataset_name: + raise HTTPException(status_code=400, detail="invalid dataset_name") + + project_root = Path(__file__).resolve().parents[4] + info_path = project_root / "output" / "dataset" / user_id / dataset_name / "meta" / "info.json" + if not info_path.exists(): + return { + "dataset_name": dataset_name, + "total_frames": 0, + "total_episodes": 0, + "exists": False, + } + + with open(info_path, "r") as f: + info = json.load(f) + + return { + "dataset_name": dataset_name, + "total_frames": int(info.get("total_frames", 0) or 0), + "total_episodes": int(info.get("total_episodes", 0) or 0), + "exists": True, + } + + @router.post("/collect") async def collect_image(payload: CollectImagePayload): """将前端直接采集到的图像写入当前 episode。""" diff --git a/backend/api/domains/training/routes.py b/backend/api/domains/training/routes.py index 51294f3..6f845a8 100644 --- a/backend/api/domains/training/routes.py +++ b/backend/api/domains/training/routes.py @@ -22,6 +22,13 @@ def set_sio_server(sio): _sio_server = sio +async def _run_training_task(**kwargs): + try: + await training.train_model(_sio_server, **kwargs) + except Exception: + logger.exception("后台训练任务异常退出") + + @router.post("") async def start_training(request: TrainRequest, user_id: str = Query(...)): """启动训练。""" @@ -59,21 +66,25 @@ async def start_training(request: TrainRequest, user_id: str = Query(...)): if request.resume_from: resume_path = str(project_root / request.resume_from) - logger.info(f"收到开始训练请求: user={user_id}, epochs={request.epochs}, batch_size={request.batch_size}, lr={request.lr}, resume_from={resume_path}") - - asyncio.create_task( - training.train_model( - _sio_server, - user_id=user_id, - data_dir=str(data_path), - output_dir=str(output_path), - epochs=request.epochs, - batch_size=request.batch_size, - lr=request.lr, - resume_from=resume_path, - ) + logger.info( + "收到开始训练请求: user=%s, epochs=%s, batch_size=%s, lr=%s, resume_from=%s", + user_id, + request.epochs, + request.batch_size, + request.lr, + resume_path, ) + asyncio.create_task(_run_training_task( + user_id=user_id, + data_dir=str(data_path), + output_dir=str(output_path), + epochs=request.epochs, + batch_size=request.batch_size, + lr=request.lr, + resume_from=resume_path, + )) + resume_msg = f"，从模型继续: {request.resume_from}" if request.resume_from else "" return { "success": True, diff --git a/backend/config.py b/backend/config.py index cb7ecdd..4d3122c 100644 --- a/backend/config.py +++ b/backend/config.py @@ -10,7 +10,7 @@ class Config: # 服务器配置 HOST = os.getenv("HOST", "0.0.0.0") - PORT = int(os.getenv("PORT", "8000")) + PORT = int(os.getenv("PORT", "8001")) # 模型配置 MODEL_PATH = os.getenv("MODEL_PATH", None) diff --git a/backend/run_act_checks.py b/backend/run_act_checks.py index 5abcafa..a1c7c1c 100644 --- a/backend/run_act_checks.py +++ b/backend/run_act_checks.py @@ -6,7 +6,7 @@ REPO_ROOT = Path(__file__).resolve().parents[1] -PYTEST_CMD = ["python3", "-m", "pytest", "tests/act", "-q"] +PYTEST_CMD = [sys.executable, "-m", "pytest", "tests/act", "-q"] def main() -> int: diff --git a/backend/services/episode/exporter.py b/backend/services/episode/exporter.py index 76e85fd..d02be75 100644 --- a/backend/services/episode/exporter.py +++ b/backend/services/episode/exporter.py @@ -176,17 +176,23 @@ def _build_stats_entry(self, key: str, values: np.ndarray) -> Dict[str, Any]: batch_count = int(values.shape[0]) batch_min = values.min(axis=0).astype(np.float64) batch_max = values.max(axis=0).astype(np.float64) + batch_sum = values.sum(axis=0).astype(np.float64) + batch_sumsq = np.square(values).sum(axis=0).astype(np.float64) existing = self._stats_accumulators.get(key) if existing is None: accum = { "count": batch_count, + "sum": batch_sum, + "sumsq": batch_sumsq, "min": batch_min, "max": batch_max, } else: accum = { "count": existing["count"] + batch_count, + "sum": existing.get("sum", np.zeros_like(batch_sum)) + batch_sum, + "sumsq": existing.get("sumsq", np.zeros_like(batch_sumsq)) + batch_sumsq, "min": np.minimum(existing["min"], batch_min), "max": np.maximum(existing["max"], batch_max), } @@ -205,15 +211,52 @@ def _build_stats_entry(self, key: str, values: np.ndarray) -> Dict[str, Any]: all_values = self._stats_accumulators[values_key] q01 = np.percentile(all_values, 1, axis=0).astype(np.float64) q99 = np.percentile(all_values, 99, axis=0).astype(np.float64) + mean = accum["sum"] / max(accum["count"], 1) + variance = accum["sumsq"] / max(accum["count"], 1) - np.square(mean) + std = np.sqrt(np.maximum(variance, 0.0)) return { "count": int(accum["count"]), + "sum": accum["sum"].tolist(), + "sumsq": accum["sumsq"].tolist(), "min": accum["min"].tolist(), "max": accum["max"].tolist(), + "mean": mean.tolist(), + "std": std.tolist(), "q01": q01.tolist(), "q99": q99.tolist(), } + def _read_saved_numeric_columns(self) -> tuple[np.ndarray, np.ndarray] | None: + """从已写入 parquet 的全量数据重建状态和动作数组。""" + if not self.data_dir.exists(): + return None + + states = [] + actions = [] + for parquet_file in sorted(self.data_dir.glob("chunk-*/file-*.parquet")): + df = pd.read_parquet(parquet_file, columns=["observation.state", "action"]) + states.extend(np.asarray(value, dtype=np.float32) for value in df["observation.state"]) + actions.extend(np.asarray(value, dtype=np.float32) for value in df["action"]) + + if not states or not actions: + return None + + return np.stack(states), np.stack(actions) + + def _rebuild_stats_from_saved_data(self): + """确保 stats.json 反映整个数据集，而不是最近一次导出的 episode。""" + saved_data = self._read_saved_numeric_columns() + if saved_data is None: + return + + states_array, actions_array = saved_data + self._stats_accumulators = {} + self._stats = { + "observation.state": self._build_stats_entry("observation.state", states_array), + "action": self._build_stats_entry("action", actions_array), + } + def _ensure_dirs(self): """确保目录结构存在""" self.data_dir.mkdir(parents=True, exist_ok=True) @@ -415,6 +458,7 @@ def _save_all_metadata(self): """保存所有元数据文件""" # 保存 info.json self._ensure_dirs() + self._rebuild_stats_from_saved_data() info = self.info with open(self.meta_dir / "info.json", "w") as f: json.dump(info, f, indent=2) diff --git a/backend/services/training/dataset.py b/backend/services/training/dataset.py index 1128310..b60b724 100644 --- a/backend/services/training/dataset.py +++ b/backend/services/training/dataset.py @@ -72,4 +72,4 @@ def __getitem__(self, idx): "state": state, }, "action": action, - } \ No newline at end of file + } diff --git a/backend/services/training/orchestrator.py b/backend/services/training/orchestrator.py index 6bf448b..a993815 100644 --- a/backend/services/training/orchestrator.py +++ b/backend/services/training/orchestrator.py @@ -4,6 +4,7 @@ import asyncio import logging +import os from pathlib import Path from typing import Optional @@ -22,6 +23,9 @@ logger = logging.getLogger(__name__) +TRAINING_MAX_CPU_THREADS = 4 +TRAINING_RESERVED_CPU_THREADS = 2 + def _extract_checkpoint_payload(checkpoint): """兼容旧格式 state_dict 和新格式 checkpoint dict。""" @@ -30,6 +34,28 @@ def _extract_checkpoint_payload(checkpoint): return checkpoint, {} +def _configure_training_cpu_threads(): + cpu_count = os.cpu_count() or 1 + available_threads = max(1, cpu_count - TRAINING_RESERVED_CPU_THREADS) + training_threads = min(TRAINING_MAX_CPU_THREADS, available_threads) + + if torch.get_num_threads() != training_threads: + torch.set_num_threads(training_threads) + + try: + if torch.get_num_interop_threads() != training_threads: + torch.set_num_interop_threads(training_threads) + except RuntimeError as exc: + logger.debug("PyTorch interop 线程数已锁定，保持当前设置: %s", exc) + + logger.info( + "训练 CPU 线程配置: torch=%s, interop=%s, cpu=%s", + torch.get_num_threads(), + torch.get_num_interop_threads(), + cpu_count, + ) + + def _train_model_sync( sio_server, user_id: str, @@ -48,6 +74,8 @@ def _train_model_sync( training_state["total_epochs"] = epochs training_state["loss"] = 0.0 training_state["progress"] = 0.0 + training_state["error"] = None + training_state["message"] = "训练中" callbacks = TrainingCallbacks(sio_server, loop=loop, namespace="/", training_state=training_state) @@ -57,8 +85,12 @@ def _train_model_sync( if resume_from: logger.info(f"从已有模型继续训练: {resume_from}") logger.info("=" * 20) + _configure_training_cpu_threads() data = load_dataset(data_dir) + if not data: + raise ValueError(f"训练集为空或无法加载: {data_dir}") + action_dim = data["action"].shape[-1] raw_state_dim = data["observation.state"].shape[-1] state_dim = 2 @@ -172,10 +204,13 @@ def _train_model_sync( logger.info(f"模型已保存到: {final_path}") callbacks.on_train_end(str(final_path)) + training_state["message"] = f"训练完成: {final_path}" return model except Exception as exc: - logger.error(f"训练失败: {exc}") + logger.exception("训练失败") training_state["is_running"] = False + training_state["error"] = str(exc) + training_state["message"] = f"训练失败: {exc}" raise diff --git a/backend/services/training/state.py b/backend/services/training/state.py index d65e154..87f037a 100644 --- a/backend/services/training/state.py +++ b/backend/services/training/state.py @@ -1,20 +1,37 @@ """Shared training progress state.""" +from __future__ import annotations + # 按用户隔离的训练状态 # 格式: { user_id: training_state } _user_training_states: dict[str, dict] = {} +TRAINING_STATE_VERSION = "2026-06-26-error-reporting" + + +def _create_initial_state() -> dict: + return { + "is_running": False, + "epoch": 0, + "total_epochs": 0, + "loss": 0.0, + "progress": 0.0, + "error": None, + "message": "", + "version": TRAINING_STATE_VERSION, + } + def _ensure_user_state(user_id: str) -> dict: """确保用户训练状态存在""" if user_id not in _user_training_states: - _user_training_states[user_id] = { - "is_running": False, - "epoch": 0, - "total_epochs": 0, - "loss": 0.0, - "progress": 0.0, - } + _user_training_states[user_id] = _create_initial_state() + else: + _user_training_states[user_id].update({ + key: value + for key, value in _create_initial_state().items() + if key not in _user_training_states[user_id] + }) return _user_training_states[user_id] @@ -27,3 +44,4 @@ def stop_training(user_id: str): """Stop training for user.""" if user_id in _user_training_states: _user_training_states[user_id]["is_running"] = False + _user_training_states[user_id]["message"] = "训练已停止" diff --git a/ui/README.md b/ui/README.md index 1077987..508b0c8 100644 --- a/ui/README.md +++ b/ui/README.md @@ -28,7 +28,7 @@ npm run dev npm run build ``` -前端运行在 **http://localhost:5173**，会自动代理 API 请求到后端 `http://localhost:8000`。 +前端运行在 **http://localhost:5175**，会自动代理 API 请求到后端 `http://localhost:8000`。 ## 页面 diff --git a/ui/package-lock.json b/ui/package-lock.json index 8dee93a..3bb439f 100644 --- a/ui/package-lock.json +++ b/ui/package-lock.json @@ -80,6 +80,7 @@ "integrity": "sha512-CGOfOJqWjg2qW/Mb6zNsDm+u5vFQ8DxXfbM09z69p5Z6+mE1ikP2jUXw+j42Pf1XTYED2Rni5f95npYeuwMDQA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@babel/code-frame": "^7.29.0", "@babel/generator": "^7.29.0", @@ -1749,6 +1750,7 @@ "integrity": "sha512-OowOUbD1lBCOFIPOZ8xnMIhgqA4sCutMiYOmPHL1PTLt5+y1XA+g2+yC9OOyz8p+deMZqPZLxfMjYIfrKsPeFg==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "undici-types": "~7.16.0" } @@ -1759,6 +1761,7 @@ "integrity": "sha512-ilcTH/UniCkMdtexkoCN0bI7pMcJDvmQFPvuPvmEaYA/NSfFTAgdUSLAoVjaRJm7+6PvcM+q1zYOwS4wTYMF9w==", "devOptional": true, "license": "MIT", + "peer": true, "dependencies": { "csstype": "^3.2.2" } @@ -1844,6 +1847,7 @@ "integrity": "sha512-klQbnPAAiGYFyI02+znpBRLyjL4/BrBd0nyWkdC0s/6xFLkXYQ8OoRrSkqacS1ddVxf/LDyODIKbQ5TgKAf/Fg==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@typescript-eslint/scope-manager": "8.56.1", "@typescript-eslint/types": "8.56.1", @@ -2121,6 +2125,7 @@ "integrity": "sha512-UVJyE9MttOsBQIDKw1skb9nAwQuR5wuGD3+82K6JgJlm/Y+KI92oNsMNGZCYdDsVtRHSak0pcV5Dno5+4jh9sw==", "dev": true, "license": "MIT", + "peer": true, "bin": { "acorn": "bin/acorn" }, @@ -2239,6 +2244,7 @@ } ], "license": "MIT", + "peer": true, "dependencies": { "baseline-browser-mapping": "^2.9.0", "caniuse-lite": "^1.0.30001759", @@ -2518,6 +2524,7 @@ "integrity": "sha512-VmQ+sifHUbI/IcSopBCF/HO3YiHQx/AVd3UVyYL6weuwW+HvON9VYn5l6Zl1WZzPWXPNZrSQpxwkkZ/VuvJZzg==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@eslint-community/eslint-utils": "^4.8.0", "@eslint-community/regexpp": "^4.12.1", @@ -3515,6 +3522,7 @@ "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "dev": true, "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, @@ -3576,6 +3584,7 @@ "resolved": "https://registry.npmjs.org/react/-/react-19.2.4.tgz", "integrity": "sha512-9nfp2hYpCwOjAN+8TZFGhtWEwgvWHXqESH8qT89AT/lWklpLON22Lc8pEtnpsZz7VmawabSU0gCjnj8aC0euHQ==", "license": "MIT", + "peer": true, "engines": { "node": ">=0.10.0" } @@ -3585,6 +3594,7 @@ "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.4.tgz", "integrity": "sha512-AXJdLo8kgMbimY95O2aKQqsz2iWi9jMgKJhRBAxECE4IFxfcazB2LmzloIoibJI3C12IlY20+KFaLv+71bUJeQ==", "license": "MIT", + "peer": true, "dependencies": { "scheduler": "^0.27.0" }, @@ -3890,6 +3900,7 @@ "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", "dev": true, "license": "Apache-2.0", + "peer": true, "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" @@ -3976,6 +3987,7 @@ "integrity": "sha512-w+N7Hifpc3gRjZ63vYBXA56dvvRlNWRczTdmCBBa+CotUzAPf5b7YMdMR/8CQoeYE5LX3W4wj6RYTgonm1b9DA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "esbuild": "^0.27.0", "fdir": "^6.5.0", @@ -4126,6 +4138,7 @@ "integrity": "sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg==", "dev": true, "license": "MIT", + "peer": true, "funding": { "url": "https://github.com/sponsors/colinhacks" } diff --git a/ui/src/api/api.ts b/ui/src/api/api.ts index 5b020aa..061ae5c 100644 --- a/ui/src/api/api.ts +++ b/ui/src/api/api.ts @@ -18,6 +18,18 @@ export const listDatasetDirs = (userId: string) => .get(`dataset/dirs?user_id=${encodeURIComponent(userId)}`) .json(); +export interface DatasetInfoResponse { + dataset_name: string; + total_frames: number; + total_episodes: number; + exists: boolean; +} + +export const getDatasetInfo = (userId: string, datasetName: string) => + api + .get(`dataset/info?user_id=${encodeURIComponent(userId)}&dataset_name=${encodeURIComponent(datasetName)}`) + .json(); + export interface ModelsResponse { models?: string[]; } @@ -49,6 +61,21 @@ export const startTraining = (userId: string, params: TrainingParams) => .post(`train?user_id=${encodeURIComponent(userId)}`, { json: params }) .json(); +export interface TrainingStatus { + is_running: boolean; + epoch: number; + total_epochs: number; + loss: number; + progress: number; + error?: string | null; + message?: string; +} + +export const getTrainingStatus = (userId: string) => + api + .get(`train/status?user_id=${encodeURIComponent(userId)}`) + .json(); + export interface StopTrainingResponse { success?: boolean; message?: string; diff --git a/ui/src/pages/MujocoPage/ControlPanel.tsx b/ui/src/pages/MujocoPage/ControlPanel.tsx index f8da201..8dfa237 100644 --- a/ui/src/pages/MujocoPage/ControlPanel.tsx +++ b/ui/src/pages/MujocoPage/ControlPanel.tsx @@ -9,6 +9,10 @@ interface Props { onTurnSpeedChange: (v: number) => void; } +const DRIVE_SPEED_MIN = 0.5; +const DRIVE_SPEED_MAX = 5; +const DRIVE_SPEED_STEP = 0.1; + export function ControlPanel({ isLoaded, reset, @@ -27,10 +31,13 @@ export function ControlPanel({

Drive Settings

onDriveSpeedChange(parseFloat(e.target.value))} className="w-full h-1.5 bg-slate-700 rounded-lg appearance-none cursor-pointer accent-emerald-500" diff --git a/ui/src/pages/MujocoPage/MujocoRenderer.tsx b/ui/src/pages/MujocoPage/MujocoRenderer.tsx index 9017540..8976ac0 100644 --- a/ui/src/pages/MujocoPage/MujocoRenderer.tsx +++ b/ui/src/pages/MujocoPage/MujocoRenderer.tsx @@ -1,4 +1,4 @@ -import { useEffect, useRef, useCallback } from "react"; +import { useEffect, useRef, useCallback, type RefObject } from "react"; import * as THREE from "three"; import type { MainModule, MjModel, MjData } from "@mujoco/mujoco"; @@ -32,7 +32,7 @@ function createGeomMesh( switch (type) { case MJ_GEOM_PLANE: { - const geometry = new THREE.PlaneGeometry(20, 20); + const geometry = new THREE.PlaneGeometry(size[0] * 2, size[1] * 2); geometry.rotateX(-Math.PI / 2); const material = new THREE.MeshStandardMaterial({ color, @@ -188,6 +188,9 @@ interface Props { isLoaded: boolean; onStep: () => void; showJointOverlay: boolean; + onPlaceBall?: (position: { x: number; y: number }) => void; + onFirstPersonCanvasReady?: (canvas: HTMLCanvasElement | null) => void; + firstPersonContainerRef?: RefObject; } export default function MujocoRenderer({ @@ -197,14 +200,17 @@ export default function MujocoRenderer({ isLoaded, onStep, showJointOverlay, + onPlaceBall, + onFirstPersonCanvasReady, + firstPersonContainerRef, }: Props) { const containerRef = useRef(null); - const fpContainerRef = useRef(null); + const fallbackFpContainerRef = useRef(null); const axesContainerRef = useRef(null); const sceneRef = useRef(null); const orbitCameraRef = useRef(null); const fpCameraRef = useRef(null); - const axesCameraRef = useRef(null); + const axesCameraRef = useRef(null); const axesSceneRef = useRef(null); const axesGroupRef = useRef(null); const mainRendererRef = useRef(null); @@ -213,6 +219,7 @@ export default function MujocoRenderer({ const bodyGroupsRef = useRef([]); const animRef = useRef(0); const draggingRef = useRef(false); + const draggingBallRef = useRef(false); const lastMouseRef = useRef({ x: 0, y: 0 }); const orbitStateRef = useRef({ azimuth: 0.5, @@ -222,15 +229,61 @@ export default function MujocoRenderer({ }); const fpCamIdxRef = useRef(-1); const fpCamBodyIdRef = useRef(-1); + const carBodyIdRef = useRef(-1); const jointGroupRef = useRef(null); const jointOverlaysRef = useRef([]); const jointFrameCountRef = useRef(0); + const prevLoadedRef = useRef(false); const showJointOverlayRef = useRef(showJointOverlay); - showJointOverlayRef.current = showJointOverlay; + + useEffect(() => { + showJointOverlayRef.current = showJointOverlay; + }, [showJointOverlay]); + + const getRaycasterFromPointer = useCallback((e: React.PointerEvent) => { + const container = containerRef.current; + const camera = orbitCameraRef.current; + if (!container || !camera) return null; + + const rect = container.getBoundingClientRect(); + const pointer = new THREE.Vector2( + ((e.clientX - rect.left) / rect.width) * 2 - 1, + -((e.clientY - rect.top) / rect.height) * 2 + 1, + ); + const raycaster = new THREE.Raycaster(); + raycaster.setFromCamera(pointer, camera); + return raycaster; + }, []); + + const findTargetBallGroup = useCallback(() => { + return bodyGroupsRef.current.find((group) => group.name === "target_ball") ?? null; + }, []); + + const isPointerOnTargetBall = useCallback((e: React.PointerEvent) => { + const ballGroup = findTargetBallGroup(); + const raycaster = getRaycasterFromPointer(e); + if (!ballGroup || !raycaster) return false; + + return raycaster.intersectObject(ballGroup, true).length > 0; + }, [findTargetBallGroup, getRaycasterFromPointer]); + + const placeBallFromPointer = useCallback((e: React.PointerEvent) => { + const raycaster = getRaycasterFromPointer(e); + if (!raycaster || !onPlaceBall) return; + + const groundPlane = new THREE.Plane(new THREE.Vector3(0, 1, 0), 0); + const hitPoint = new THREE.Vector3(); + if (!raycaster.ray.intersectPlane(groundPlane, hitPoint)) return; + + onPlaceBall({ + x: hitPoint.x, + y: -hitPoint.z, + }); + }, [getRaycasterFromPointer, onPlaceBall]); const setupScene = useCallback(() => { const container = containerRef.current; - const fpContainer = fpContainerRef.current; + const fpContainer = firstPersonContainerRef?.current ?? fallbackFpContainerRef.current; if (!container || !fpContainer) return; const scene = new THREE.Scene(); @@ -254,17 +307,22 @@ export default function MujocoRenderer({ mainRenderer.domElement.style.left = "0"; container.appendChild(mainRenderer.domElement); - const fpRenderer = new THREE.WebGLRenderer({ antialias: true, alpha: false }); + const fpRenderer = new THREE.WebGLRenderer({ + antialias: true, + alpha: false, + preserveDrawingBuffer: true, + }); fpRenderer.setClearColor(0x2a2a4e, 1); fpRenderer.setSize(fpW, fpH); fpRenderer.setPixelRatio(Math.min(window.devicePixelRatio, 2)); fpRenderer.shadowMap.enabled = true; fpContainer.appendChild(fpRenderer.domElement); + onFirstPersonCanvasReady?.(fpRenderer.domElement); // Axes gizmo (bottom-left) const axesContainer = axesContainerRef.current; let axesScene: THREE.Scene | null = null; - let axesCamera: THREE.PerspectiveCamera | null = null; + let axesCamera: THREE.OrthographicCamera | null = null; let axesRenderer: THREE.WebGLRenderer | null = null; if (axesContainer) { axesScene = new THREE.Scene(); @@ -277,8 +335,9 @@ export default function MujocoRenderer({ axesGroupRef.current = axesGroup; const size = 120; const halfSize = 1.3; - axesCamera = new THREE.OrthographicCamera(-halfSize, halfSize, halfSize, -halfSize, 0.1, 10); - axesCamera.position.set(0, 0, 3); + const orthographicCamera = new THREE.OrthographicCamera(-halfSize, halfSize, halfSize, -halfSize, 0.1, 10); + orthographicCamera.position.set(0, 0, 3); + axesCamera = orthographicCamera; axesRenderer = new THREE.WebGLRenderer({ antialias: true, alpha: false }); axesRenderer.setClearColor(0x1a1a2e, 1); axesRenderer.setSize(size, size); @@ -301,7 +360,7 @@ export default function MujocoRenderer({ const hemi = new THREE.HemisphereLight(0x87ceeb, 0x3a3a3a, 0.7); scene.add(hemi); - const grid = new THREE.GridHelper(10, 20, 0x444466, 0x222244); + const grid = new THREE.GridHelper(20, 40, 0x444466, 0x222244); scene.add(grid); const axes = new THREE.AxesHelper(2); @@ -312,7 +371,7 @@ export default function MujocoRenderer({ fpCameraRef.current = fpCamera; mainRendererRef.current = mainRenderer; fpRendererRef.current = fpRenderer; - }, []); + }, [firstPersonContainerRef, onFirstPersonCanvasReady]); const updateOrbitCamera = useCallback(() => { const camera = orbitCameraRef.current; @@ -336,7 +395,7 @@ export default function MujocoRenderer({ setupScene(); const handleResize = () => { const container = containerRef.current; - const fpContainer = fpContainerRef.current; + const fpContainer = firstPersonContainerRef?.current ?? fallbackFpContainerRef.current; const mainRenderer = mainRendererRef.current; const fpRenderer = fpRendererRef.current; const orbitCamera = orbitCameraRef.current; @@ -367,10 +426,11 @@ export default function MujocoRenderer({ fpRendererRef.current?.dispose(); axesRendererRef.current?.domElement?.remove(); axesRendererRef.current?.dispose(); + onFirstPersonCanvasReady?.(null); + prevLoadedRef.current = false; }; - }, [setupScene]); + }, [firstPersonContainerRef, setupScene, onFirstPersonCanvasReady]); - const prevLoadedRef = useRef(false); useEffect(() => { if (!isLoaded || !model.current || prevLoadedRef.current) return; prevLoadedRef.current = true; @@ -418,6 +478,7 @@ export default function MujocoRenderer({ const { fpIdx, tdIdx } = resolveCameraIndices(m); fpCamIdxRef.current = fpIdx; fpCamBodyIdRef.current = fpIdx >= 0 ? (m.cam_bodyid?.[fpIdx] ?? -1) : -1; + carBodyIdRef.current = bodyGroups.findIndex((group) => group.name === "car"); const orbitCam = orbitCameraRef.current; const fpCam = fpCameraRef.current; @@ -459,7 +520,6 @@ export default function MujocoRenderer({ const canvas = document.createElement("canvas"); canvas.width = 256; canvas.height = 64; - const ctx = canvas.getContext("2d")!; const texture = new THREE.CanvasTexture(canvas); texture.minFilter = THREE.LinearFilter; const spriteMat = new THREE.SpriteMaterial({ map: texture, depthTest: false }); @@ -473,7 +533,7 @@ export default function MujocoRenderer({ jointOverlaysRef.current = overlays; updateOrbitCamera(); - }, [isLoaded, model, updateOrbitCamera]); + }, [isLoaded, model, showJointOverlay, updateOrbitCamera]); useEffect(() => { let running = true; @@ -504,11 +564,13 @@ export default function MujocoRenderer({ } // Track car body position for orbit camera target - const carBodyId = 1; + const carBodyId = carBodyIdRef.current; const carPos3 = new THREE.Vector3(); - mjPosToThree(d.xpos, carBodyId, carPos3); - orbitStateRef.current.target.set(carPos3.x, carPos3.y + 0.3, carPos3.z); - updateOrbitCamera(); + if (carBodyId >= 0) { + mjPosToThree(d.xpos, carBodyId, carPos3); + orbitStateRef.current.target.set(carPos3.x, carPos3.y + 0.3, carPos3.z); + updateOrbitCamera(); + } // First-person camera: reuse body group transform (same proven quaternion path) const fpIdx = fpCamIdxRef.current; @@ -552,16 +614,16 @@ export default function MujocoRenderer({ const qposVal = d.qpos[ov.qposAdr]?.toFixed(3) ?? "?"; const canvas = (ov.sprite.material as THREE.SpriteMaterial).map?.image as HTMLCanvasElement; if (canvas) { - const ctx = canvas.getContext("2d"); - if (ctx) { - ctx.clearRect(0, 0, canvas.width, canvas.height); - ctx.fillStyle = "rgba(0,0,0,0.75)"; - ctx.fillRect(0, 0, canvas.width, canvas.height); - ctx.fillStyle = "#ffffff"; - ctx.font = "20px monospace"; - ctx.textAlign = "center"; - ctx.textBaseline = "middle"; - ctx.fillText(`${ov.name}: ${qposVal}`, 128, 32); + const labelCtx = canvas.getContext("2d"); + if (labelCtx) { + labelCtx.clearRect(0, 0, canvas.width, canvas.height); + labelCtx.fillStyle = "rgba(0,0,0,0.75)"; + labelCtx.fillRect(0, 0, canvas.width, canvas.height); + labelCtx.fillStyle = "#ffffff"; + labelCtx.font = "20px monospace"; + labelCtx.textAlign = "center"; + labelCtx.textBaseline = "middle"; + labelCtx.fillText(`${ov.name}: ${qposVal}`, 128, 32); (ov.sprite.material as THREE.SpriteMaterial).map!.needsUpdate = true; } } @@ -599,16 +661,28 @@ export default function MujocoRenderer({ running = false; cancelAnimationFrame(animRef.current); }; - }, [isLoaded, mujoco, data, onStep]); + }, [isLoaded, mujoco, data, model, onStep, updateOrbitCamera]); const handlePointerDown = useCallback((e: React.PointerEvent) => { + if (isPointerOnTargetBall(e)) { + draggingBallRef.current = true; + (e.target as HTMLElement).setPointerCapture(e.pointerId); + placeBallFromPointer(e); + return; + } + draggingRef.current = true; lastMouseRef.current = { x: e.clientX, y: e.clientY }; (e.target as HTMLElement).setPointerCapture(e.pointerId); - }, []); + }, [isPointerOnTargetBall, placeBallFromPointer]); const handlePointerMove = useCallback( (e: React.PointerEvent) => { + if (draggingBallRef.current) { + placeBallFromPointer(e); + return; + } + if (!draggingRef.current) return; const dx = e.clientX - lastMouseRef.current.x; const dy = e.clientY - lastMouseRef.current.y; @@ -622,10 +696,15 @@ export default function MujocoRenderer({ ); updateOrbitCamera(); }, - [updateOrbitCamera], + [placeBallFromPointer, updateOrbitCamera], ); const handlePointerUp = useCallback(() => { + if (draggingBallRef.current) { + draggingBallRef.current = false; + return; + } + draggingRef.current = false; }, []); @@ -660,11 +739,13 @@ export default function MujocoRenderer({ ref={axesContainerRef} className="absolute bottom-3 left-3 w-[120px] h-[120px] z-10 rounded-md overflow-hidden border border-slate-600/50" /> -
+ {!firstPersonContainerRef && ( +
+ )}
); } diff --git a/ui/src/pages/MujocoPage/carArmXml.ts b/ui/src/pages/MujocoPage/carArmXml.ts index 3351a20..7191dff 100644 --- a/ui/src/pages/MujocoPage/carArmXml.ts +++ b/ui/src/pages/MujocoPage/carArmXml.ts @@ -1,5 +1,3 @@ -import { MAZE_XML } from './mazeXml'; - export const CAR_ARM_XML = ` `; diff --git a/ui/src/pages/MujocoPage/index.tsx b/ui/src/pages/MujocoPage/index.tsx index d01f32c..b170a25 100644 --- a/ui/src/pages/MujocoPage/index.tsx +++ b/ui/src/pages/MujocoPage/index.tsx @@ -1,7 +1,11 @@ import { useState, useEffect, useRef, useCallback } from "react"; -import { useMujoco } from "./useMujoco"; +import { startTraining, stopTraining, getTrainingStatus, loadTrainedModel, listDatasetDirs, listModels, getDatasetInfo, type TrainingStatus } from "../../api/api"; +import { runInferenceWithSocket, simSocket } from "../../api/socket"; +import { getDatasetPath, getTrainPath } from "../../lib/constants"; +import { useSimCarStore } from "../../stores/simCarStore"; +import { useMujoco, type CarToBallState, type MjPosition } from "./useMujoco"; +import { useDataCollection } from "./useDataCollection"; import MujocoRenderer from "./MujocoRenderer"; -import { ControlPanel } from "./ControlPanel"; const DRIVE_KEYS = [ "KeyW", "KeyA", "KeyS", "KeyD", @@ -9,21 +13,401 @@ const DRIVE_KEYS = [ "Space", ]; +type TaskMode = "follow" | "avoid"; +type ControlMode = "manual" | "auto" | "model"; +type ModelTaskRunState = { + startedAt: number; + followCompleteSince: number | null; + avoidDangerSeen: boolean; + avoidSafeSince: number | null; +}; + +const TASK_DATASET_NAMES: Record = { + follow: "mujoco-follow", + avoid: "mujoco-avoid", +}; +const MANAGED_TASK_DATASET_NAMES = new Set(["mujoco-default", ...Object.values(TASK_DATASET_NAMES)]); + +const clamp = (value: number, min: number, max: number) => Math.max(min, Math.min(max, value)); +const GROUND_GRID_HALF_SIZE = 10; +const BALL_GRID_MARGIN = 0.5; +const BALL_GRID_LIMIT = GROUND_GRID_HALF_SIZE - BALL_GRID_MARGIN; +const AVOID_CAR_GRID_LIMIT = GROUND_GRID_HALF_SIZE - 1.1; +const DRIVE_SPEED_MIN = 0.5; +const DRIVE_SPEED_MAX = 5; +const DRIVE_SPEED_STEP = 0.1; +const DEFAULT_DRIVE_SPEED = 0.7; +const TURN_SPEED_MIN = 0.5; +const TURN_SPEED_MAX = 8; +const TURN_SPEED_STEP = 0.5; +const DEFAULT_TURN_SPEED = 4; +const MOTOR_CONTROL_LIMIT = 8; +const TRAINING_BALL_EDGE_BUFFER = 0.2; +const BALL_PLACEMENT_ATTEMPTS = 40; +const AUTO_BALL_RESET_COOLDOWN_MS = 700; +const BALL_APPROACH_DISTANCE = 1.4; +const MODEL_INFERENCE_INTERVAL_MS = 100; +const MODEL_TASK_TIMEOUT_MS = 15000; +const FOLLOW_TASK_COMPLETE_DISTANCE = BALL_APPROACH_DISTANCE; +const FOLLOW_TASK_COMPLETE_ANGLE = 0.65; +const FOLLOW_TASK_COMPLETE_STABLE_MS = 400; +const AVOID_TARGET_DISTANCE_MIN = 5.5; +const AVOID_TARGET_DISTANCE_MAX = 8; +const AVOID_TRIGGER_DISTANCE = 3; +const AVOID_TASK_MIN_DURATION_MS = 700; +const AVOID_TASK_SAFE_DISTANCE = AVOID_TRIGGER_DISTANCE + 0.8; +const AVOID_TASK_SAFE_STABLE_MS = 700; +const AVOID_VISIBLE_FRONTNESS = 0.25; +const AVOID_VISIBLE_ANGLE = 1.2; +const AVOID_RELOCATE_MIN_SETTLE_MS = 1800; +const AVOID_RELOCATE_MAX_SETTLE_MS = 4200; +const AVOID_HEADING_SAMPLE_MS = 180; +const AVOID_HEADING_STABLE_MS = 550; +const AVOID_HEADING_STABLE_RATE = 0.25; +const AVOID_MISPLACED_RETRY_MS = 1600; +const AVOID_EDGE_MARGIN = 1.0; +const AVOID_EDGE_RELEASE_MARGIN = 1.4; +const AVOID_EDGE_LOOKAHEAD_DISTANCE = 2.2; +const AVOID_EDGE_CENTER_ANGLE = 0.45; +const AVOID_EDGE_TURN_GAIN = 1.7; +const AVOID_EDGE_TURN_LIMIT = 1.35; +const AVOID_EDGE_FORWARD_SPEED_SCALE = 0.35; +const BALL_RELOCATION_MIN_DISTANCE = 3.2; +const FOLLOW_VISIBLE_ANGLES = [-0.72, -0.45, -0.22, 0, 0.22, 0.45, 0.72]; +const AVOID_VISIBLE_ANGLES = [-0.58, -0.3, 0, 0.3, 0.58]; +const AVOID_AHEAD_ANGLES = [0, -0.18, 0.18, -0.35, 0.35, -0.52, 0.52]; + +function pickNumber(values: number[]) { + return values[Math.floor(Math.random() * values.length)]; +} + +function pickNumberInRange(min: number, max: number) { + return min + Math.random() * (max - min); +} + +function roundCoordinate(value: number) { + return Number(value.toFixed(2)); +} + +function clampCoordinateToGrid(value: number) { + return clamp(value, -BALL_GRID_LIMIT, BALL_GRID_LIMIT); +} + +function clampBallPositionToGrid(position: MjPosition): MjPosition { + return { + x: roundCoordinate(clampCoordinateToGrid(position.x)), + y: roundCoordinate(clampCoordinateToGrid(position.y)), + z: position.z, + }; +} + +function getPlanarDistance(from: MjPosition, to: MjPosition) { + return Math.hypot(from.x - to.x, from.y - to.y); +} + +function getGridEdgeMargin(position: MjPosition, gridLimit: number) { + return Math.min(gridLimit - Math.abs(position.x), gridLimit - Math.abs(position.y)); +} + +function getSmallestAngleDelta(from: number, to: number) { + return Math.atan2(Math.sin(to - from), Math.cos(to - from)); +} + +function isAvoidBallInView(state: CarToBallState) { + return state.frontness > AVOID_VISIBLE_FRONTNESS && Math.abs(state.angleError) < AVOID_VISIBLE_ANGLE; +} + +function createRandomBallPosition(label = "随机位置") { + const position = clampBallPositionToGrid({ + x: roundCoordinate(Math.random() * 8 - 4), + y: roundCoordinate(Math.random() * 9 - 2), + z: 0.25, + }); + return { label, ...position }; +} + +function getTrainingBallDistanceRange(taskMode: TaskMode) { + return taskMode === "follow" + ? { min: 2.3, max: 5.5 } + : { min: AVOID_TARGET_DISTANCE_MIN, max: AVOID_TARGET_DISTANCE_MAX }; +} + +function getRayDistanceToGridBounds(origin: MjPosition, angle: number, gridLimit: number) { + const worldMin = -gridLimit; + const worldMax = gridLimit; + const directionX = Math.cos(angle); + const directionY = Math.sin(angle); + const distances: number[] = []; + + if (directionX > 0) distances.push((worldMax - origin.x) / directionX); + if (directionX < 0) distances.push((worldMin - origin.x) / directionX); + if (directionY > 0) distances.push((worldMax - origin.y) / directionY); + if (directionY < 0) distances.push((worldMin - origin.y) / directionY); + + const positiveDistances = distances.filter((distance) => Number.isFinite(distance) && distance > 0); + return positiveDistances.length > 0 ? Math.min(...positiveDistances) : 0; +} + +function getRayDistanceToTrainingBounds(origin: MjPosition, angle: number) { + return getRayDistanceToGridBounds(origin, angle, BALL_GRID_LIMIT); +} + +function createBallPositionFromRay(origin: MjPosition, angle: number, distance: number): MjPosition { + return clampBallPositionToGrid({ + x: roundCoordinate(origin.x + Math.cos(angle) * distance), + y: roundCoordinate(origin.y + Math.sin(angle) * distance), + z: 0.25, + }); +} + +function createFallbackVisibleBallPosition() { + return clampBallPositionToGrid({ + x: roundCoordinate(Math.random() * 5 - 2.5), + y: roundCoordinate(Math.random() * 3 - 1), + z: 0.25, + }); +} + +function createVisibleBallPosition(state: CarToBallState | null, taskMode: TaskMode): MjPosition { + if (!state) { + return createFallbackVisibleBallPosition(); + } + + const angles = taskMode === "follow" ? FOLLOW_VISIBLE_ANGLES : AVOID_VISIBLE_ANGLES; + const distanceRange = getTrainingBallDistanceRange(taskMode); + let farthestPosition: MjPosition | null = null; + let farthestDistance = 0; + + for (let attempt = 0; attempt < BALL_PLACEMENT_ATTEMPTS; attempt++) { + const relativeAngle = pickNumber(angles) + (Math.random() * 0.12 - 0.06); + const worldAngle = state.headingAngle + relativeAngle; + const maxDistance = getRayDistanceToTrainingBounds(state.car, worldAngle) - TRAINING_BALL_EDGE_BUFFER; + if (taskMode === "follow" && maxDistance < distanceRange.min) continue; + if (taskMode === "avoid" && maxDistance < 1.5) continue; + + const distanceLimit = Math.min(distanceRange.max, maxDistance); + const distance = taskMode === "avoid" + ? distanceLimit + : distanceRange.min + Math.random() * (distanceLimit - distanceRange.min); + const position = createBallPositionFromRay(state.car, worldAngle, distance); + const relocationDistance = getPlanarDistance(position, state.ball); + if (relocationDistance > farthestDistance) { + farthestPosition = position; + farthestDistance = relocationDistance; + } + if (relocationDistance >= BALL_RELOCATION_MIN_DISTANCE) { + return position; + } + } + + if (farthestPosition) return farthestPosition; + + const fallbackAngle = Math.atan2(-state.car.y, -state.car.x); + const fallbackDistance = Math.min(distanceRange.min, Math.max(1.5, getRayDistanceToTrainingBounds(state.car, fallbackAngle) - TRAINING_BALL_EDGE_BUFFER)); + return createBallPositionFromRay(state.car, fallbackAngle, fallbackDistance); +} + +function createAvoidBallPositionAhead(state: CarToBallState | null): MjPosition { + if (!state) { + return createFallbackVisibleBallPosition(); + } + + const desiredDistance = pickNumberInRange(AVOID_TARGET_DISTANCE_MIN, AVOID_TARGET_DISTANCE_MAX); + const minimumUsefulDistance = Math.max(AVOID_TRIGGER_DISTANCE + 0.8, desiredDistance * 0.7); + let bestInGridPosition: MjPosition | null = null; + let bestInGridDistance = 0; + + for (const relativeAngle of AVOID_AHEAD_ANGLES) { + const worldAngle = state.headingAngle + relativeAngle; + const maxDistance = getRayDistanceToTrainingBounds(state.car, worldAngle) - TRAINING_BALL_EDGE_BUFFER; + if (maxDistance < 1.5) continue; + + const distance = Math.min(desiredDistance, maxDistance); + const position = createBallPositionFromRay(state.car, worldAngle, distance); + if (distance > bestInGridDistance) { + bestInGridPosition = position; + bestInGridDistance = distance; + } + if (distance >= minimumUsefulDistance) return position; + } + + if (bestInGridPosition) return bestInGridPosition; + + const fallbackAngle = Math.atan2(-state.car.y, -state.car.x); + const fallbackDistance = Math.min(desiredDistance, Math.max(1.5, getRayDistanceToTrainingBounds(state.car, fallbackAngle) - TRAINING_BALL_EDGE_BUFFER)); + return createBallPositionFromRay(state.car, fallbackAngle, fallbackDistance); +} + +function getAvoidBoundaryReturnDrive(state: CarToBallState, maxSpeed: number, turnSpeed: number) { + const edgeMargin = getGridEdgeMargin(state.car, AVOID_CAR_GRID_LIMIT); + const distanceToEdgeAhead = getRayDistanceToGridBounds(state.car, state.headingAngle, AVOID_CAR_GRID_LIMIT); + const shouldReturnToGrid = edgeMargin < AVOID_EDGE_MARGIN || distanceToEdgeAhead < AVOID_EDGE_LOOKAHEAD_DISTANCE; + if (!shouldReturnToGrid) return null; + + const centerAngle = Math.atan2(-state.car.y, -state.car.x); + const centerAngleError = getSmallestAngleDelta(state.headingAngle, centerAngle); + const turnCommand = clamp( + centerAngleError * turnSpeed * AVOID_EDGE_TURN_GAIN, + -turnSpeed * AVOID_EDGE_TURN_LIMIT, + turnSpeed * AVOID_EDGE_TURN_LIMIT, + ); + const forwardSpeed = Math.abs(centerAngleError) < AVOID_EDGE_CENTER_ANGLE + ? maxSpeed * AVOID_EDGE_FORWARD_SPEED_SCALE + : 0; + + return { + leftVel: clampMotorCommand(forwardSpeed - turnCommand), + rightVel: clampMotorCommand(forwardSpeed + turnCommand), + }; +} + +function clampMotorCommand(value: number) { + return clamp(value, -MOTOR_CONTROL_LIMIT, MOTOR_CONTROL_LIMIT); +} + +const defaultTrainingStatus: TrainingStatus = { + is_running: false, + epoch: 0, + total_epochs: 50, + loss: 0, + progress: 0, + error: null, + message: "", +}; + +function timestamp() { + return new Date().toLocaleTimeString("zh-CN", { hour12: false }); +} + export default function MujocoPage() { - const { isLoaded, mujoco, model, data, step, setControl, reset } = - useMujoco(); + const userId = useSimCarStore((state) => state.userId); + const { + isLoaded, + mujoco, + model, + data, + step, + setControl, + reset, + getBodyPosition, + getCarToBallState, + getWheelVelocityState, + setTargetBallPosition, + } = useMujoco(); const keysRef = useRef>(new Set()); + const firstPersonCanvasRef = useRef(null); + const firstPersonPreviewRef = useRef(null); + const lastDriveCommandRef = useRef({ left: Number.NaN, right: Number.NaN }); + const modelDriveRef = useRef({ left: 0, right: 0 }); + const inferenceTimerRef = useRef(null); + const modelTaskRunRef = useRef(null); + const modelTaskInferenceInFlightRef = useRef(false); + const lastAutoBallResetRef = useRef(0); + const avoidBallWasVisibleRef = useRef(false); + const avoidBallLostViewAtRef = useRef(null); + const avoidAllowMisplacedRetryRef = useRef(false); + const avoidHeadingSettleRef = useRef<{ + heading: number; + sampledAt: number; + stableSince: number | null; + } | null>(null); const [showJointOverlay, setShowJointOverlay] = useState(false); - const [driveSpeed, setDriveSpeed] = useState(5); - const [turnSpeed, setTurnSpeed] = useState(3); + const [driveSpeed, setDriveSpeed] = useState(DEFAULT_DRIVE_SPEED); + const [turnSpeed, setTurnSpeed] = useState(DEFAULT_TURN_SPEED); const [fps, setFps] = useState(0); + const [collectionFps, setCollectionFps] = useState(30); + const [datasetName, setDatasetName] = useState(TASK_DATASET_NAMES.follow); + const [savedDatasetFrameCount, setSavedDatasetFrameCount] = useState(0); + const [datasets, setDatasets] = useState([]); + const [trainingDatasetName, setTrainingDatasetName] = useState(""); + const [episodeId, setEpisodeId] = useState(1); + const [taskMode, setTaskMode] = useState("follow"); + const [controlMode, setControlMode] = useState("manual"); + const [trainingEpochs, setTrainingEpochs] = useState(50); + const [trainingStatus, setTrainingStatus] = useState(defaultTrainingStatus); + const [models, setModels] = useState([]); + const [selectedModel, setSelectedModel] = useState(""); + const [isModelLoaded, setIsModelLoaded] = useState(false); + const [autoInference, setAutoInference] = useState(false); + const [isTaskInferenceRunning, setIsTaskInferenceRunning] = useState(false); + const [inferenceResult, setInferenceResult] = useState("模型未加载"); + const [wheelState, setWheelState] = useState({ left: 0, right: 0 }); + const [carPosition, setCarPosition] = useState(null); + const [logs, setLogs] = useState([ + `[${timestamp()}] MuJoCo 页面就绪，等待仿真初始化...`, + ]); const fpsFramesRef = useRef(0); const fpsLastTimeRef = useRef(performance.now()); - const applyDrive = useCallback(() => { + const addLog = useCallback((message: string) => { + setLogs((items) => [`[${timestamp()}] ${message}`, ...items].slice(0, 120)); + }, []); + + const refreshDatasets = useCallback(async () => { + try { + const result = await listDatasetDirs(userId); + const nextDatasets = result.datasets || []; + setDatasets(nextDatasets); + setTrainingDatasetName((current) => { + if (current && nextDatasets.includes(current)) return current; + return nextDatasets[0] || ""; + }); + } catch (error) { + setDatasets([]); + addLog(`数据集列表刷新失败: ${error instanceof Error ? error.message : String(error)}`); + } + }, [addLog, userId]); + + const refreshDatasetFrameCount = useCallback(async (targetDatasetName = datasetName) => { + const normalizedDatasetName = targetDatasetName.trim(); + if (!normalizedDatasetName) { + setSavedDatasetFrameCount(0); + return; + } + + try { + const result = await getDatasetInfo(userId, normalizedDatasetName); + setSavedDatasetFrameCount(result.total_frames || 0); + } catch (error) { + setSavedDatasetFrameCount(0); + addLog(`数据集帧数刷新失败: ${error instanceof Error ? error.message : String(error)}`); + } + }, [addLog, datasetName, userId]); + + const collection = useDataCollection({ + userId, + datasetName, + episodeId, + fps: collectionFps, + taskName: taskMode === "follow" ? "mujoco_follow_ball" : "mujoco_avoid_ball", + modelRef: model, + dataRef: data, + fpCanvasRef: firstPersonCanvasRef, + onLog: addLog, + onEpisodeEnded: (payload) => { + if (!payload.error) { + setSavedDatasetFrameCount((value) => value + Math.max(0, payload.frame_count ?? 0)); + setEpisodeId((value) => value + 1); + refreshDatasets(); + refreshDatasetFrameCount(datasetName); + } + }, + }); + const { + isRecording, + frameCount, + socketStatus, + disconnect, + startRecording, + stopRecording, + setAction, + captureFrame, + } = collection; + const totalCollectedFrameCount = savedDatasetFrameCount + (isRecording ? frameCount : 0); + + const computeDrive = useCallback(() => { let leftVel = 0; let rightVel = 0; - const k = keysRef.current; if (k.has("KeyW") || k.has("ArrowUp")) { @@ -47,11 +431,82 @@ export default function MujocoPage() { rightVel = 0; } + return { + leftVel: clampMotorCommand(leftVel), + rightVel: clampMotorCommand(rightVel), + }; + }, [driveSpeed, turnSpeed]); + + const computeAutoDrive = useCallback(() => { + const state = getCarToBallState(); + if (!state) return { leftVel: 0, rightVel: 0 }; + + const maxSpeed = clamp(driveSpeed, DRIVE_SPEED_MIN, DRIVE_SPEED_MAX); + const turnCommand = clamp(state.angleError * turnSpeed * 1.6, -turnSpeed * 1.4, turnSpeed * 1.4); + + if (taskMode === "follow") { + const forwardSpeed = state.distance > 1.1 ? maxSpeed * 0.65 : 0; + return { + leftVel: clampMotorCommand(forwardSpeed - turnCommand), + rightVel: clampMotorCommand(forwardSpeed + turnCommand), + }; + } + + const boundaryReturnDrive = getAvoidBoundaryReturnDrive(state, maxSpeed, turnSpeed); + if (boundaryReturnDrive) return boundaryReturnDrive; + + const ballIsAhead = isAvoidBallInView(state); + const ballIsClose = state.distance < AVOID_TRIGGER_DISTANCE; + if (ballIsAhead && ballIsClose) { + const awayTurn = (state.angleError >= 0 ? -1 : 1) * turnSpeed * 1.25; + const forwardSpeed = maxSpeed * 0.35; + return { + leftVel: clampMotorCommand(forwardSpeed - awayTurn), + rightVel: clampMotorCommand(forwardSpeed + awayTurn), + }; + } + + return { + leftVel: clampMotorCommand(maxSpeed * 0.45), + rightVel: clampMotorCommand(maxSpeed * 0.45), + }; + }, [driveSpeed, getCarToBallState, taskMode, turnSpeed]); + + const computeActiveDrive = useCallback(() => { + if (controlMode === "model") { + return { + leftVel: modelDriveRef.current.left, + rightVel: modelDriveRef.current.right, + }; + } + if (controlMode === "auto") return computeAutoDrive(); + return computeDrive(); + }, [computeAutoDrive, computeDrive, controlMode]); + + const refreshTelemetry = useCallback(() => { + const { velLeft, velRight } = getWheelVelocityState(); + setWheelState({ left: velLeft, right: velRight }); + setCarPosition(getBodyPosition("car")); + }, [getBodyPosition, getWheelVelocityState]); + + const applyDrive = useCallback(() => { + const { leftVel, rightVel } = computeActiveDrive(); + setAction(leftVel, rightVel); + setControl("motor_wheel_fl", leftVel); setControl("motor_wheel_rl", leftVel); setControl("motor_wheel_fr", rightVel); setControl("motor_wheel_rr", rightVel); - }, [setControl, driveSpeed, turnSpeed]); + + const last = lastDriveCommandRef.current; + if (Math.abs(last.left - leftVel) > 0.05 || Math.abs(last.right - rightVel) > 0.05) { + lastDriveCommandRef.current = { left: leftVel, right: rightVel }; + } + }, [setAction, setControl, computeActiveDrive]); + + useEffect(() => { + refreshTelemetry(); + }, [refreshTelemetry]); useEffect(() => { const handleKeyDown = (e: KeyboardEvent) => { @@ -64,7 +519,6 @@ export default function MujocoPage() { const handleKeyUp = (e: KeyboardEvent) => { keysRef.current.delete(e.code); }; - // capture:true ensures we intercept before the browser scrolls window.addEventListener("keydown", handleKeyDown, { capture: true }); window.addEventListener("keyup", handleKeyUp); return () => { @@ -73,11 +527,102 @@ export default function MujocoPage() { }; }, []); - // Apply drive controls every frame before mj_step, with real FPS + const placeVisibleTrainingBall = useCallback((reason: string) => { + const state = getCarToBallState(); + const position = taskMode === "avoid" + ? createAvoidBallPositionAhead(state) + : createVisibleBallPosition(state, taskMode); + if (setTargetBallPosition(position)) { + lastAutoBallResetRef.current = performance.now(); + avoidBallWasVisibleRef.current = taskMode !== "avoid"; + avoidBallLostViewAtRef.current = null; + avoidAllowMisplacedRetryRef.current = taskMode === "avoid"; + avoidHeadingSettleRef.current = null; + addLog(`${reason}: (${position.x.toFixed(1)}, ${position.y.toFixed(1)})`); + } + }, [addLog, getCarToBallState, setTargetBallPosition, taskMode]); + + const avoidHeadingIsSettled = useCallback((state: CarToBallState, now: number) => { + const sample = avoidHeadingSettleRef.current; + if (!sample) { + avoidHeadingSettleRef.current = { + heading: state.headingAngle, + sampledAt: now, + stableSince: null, + }; + return false; + } + + const elapsedSeconds = (now - sample.sampledAt) / 1000; + if (elapsedSeconds * 1000 < AVOID_HEADING_SAMPLE_MS) return false; + + const headingDelta = Math.abs(getSmallestAngleDelta(sample.heading, state.headingAngle)); + const headingRate = headingDelta / Math.max(elapsedSeconds, 0.001); + const stableSince = headingRate <= AVOID_HEADING_STABLE_RATE + ? sample.stableSince ?? now + : null; + + avoidHeadingSettleRef.current = { + heading: state.headingAngle, + sampledAt: now, + stableSince, + }; + + return stableSince !== null && now - stableSince >= AVOID_HEADING_STABLE_MS; + }, []); + + const maybeRandomizeBallForAutoCollection = useCallback(() => { + if (!isRecording || controlMode !== "auto") return; + + const state = getCarToBallState(); + if (!state) return; + + const now = performance.now(); + const timeSinceLastReset = now - lastAutoBallResetRef.current; + if (timeSinceLastReset < AUTO_BALL_RESET_COOLDOWN_MS) return; + + if (taskMode === "follow") { + if (state.distance > BALL_APPROACH_DISTANCE) return; + placeVisibleTrainingBall("自动采集换目标"); + return; + } + + const ballIsInAvoidView = isAvoidBallInView(state); + + if (ballIsInAvoidView) { + avoidBallWasVisibleRef.current = true; + avoidBallLostViewAtRef.current = null; + avoidHeadingSettleRef.current = null; + return; + } + + if (!avoidBallWasVisibleRef.current) { + if (!avoidAllowMisplacedRetryRef.current) return; + if (timeSinceLastReset < AVOID_MISPLACED_RETRY_MS) return; + placeVisibleTrainingBall("自动采集重置前方目标"); + return; + } + + if (avoidBallLostViewAtRef.current === null) { + avoidBallLostViewAtRef.current = now; + avoidHeadingSettleRef.current = null; + return; + } + + const lostViewDuration = now - avoidBallLostViewAtRef.current; + const reachedMinimumWait = lostViewDuration >= AVOID_RELOCATE_MIN_SETTLE_MS; + const reachedMaximumWait = lostViewDuration >= AVOID_RELOCATE_MAX_SETTLE_MS; + if (!reachedMinimumWait) return; + if (getGridEdgeMargin(state.car, AVOID_CAR_GRID_LIMIT) < AVOID_EDGE_RELEASE_MARGIN) return; + if (!reachedMaximumWait && !avoidHeadingIsSettled(state, now)) return; + + placeVisibleTrainingBall("自动采集换目标"); + }, [avoidHeadingIsSettled, controlMode, getCarToBallState, isRecording, placeVisibleTrainingBall, taskMode]); + const stepWithDrive = useCallback(() => { + maybeRandomizeBallForAutoCollection(); applyDrive(); - // Real FPS counter fpsFramesRef.current++; const now = performance.now(); if (now - fpsLastTimeRef.current >= 1000) { @@ -87,39 +632,807 @@ export default function MujocoPage() { } step(); - }, [applyDrive, step]); + refreshTelemetry(); + captureFrame(); + }, [applyDrive, captureFrame, maybeRandomizeBallForAutoCollection, refreshTelemetry, step]); + + const handleReset = useCallback(() => { + reset(); + keysRef.current.clear(); + lastDriveCommandRef.current = { left: Number.NaN, right: Number.NaN }; + avoidBallWasVisibleRef.current = false; + avoidBallLostViewAtRef.current = null; + avoidAllowMisplacedRetryRef.current = false; + avoidHeadingSettleRef.current = null; + setWheelState({ left: 0, right: 0 }); + setCarPosition(getBodyPosition("car")); + addLog("场景已复位"); + }, [addLog, getBodyPosition, reset]); + + const handleSetBallPosition = useCallback((position: { label: string; x: number; y: number; z: number }) => { + const boundedPosition = clampBallPositionToGrid(position); + if (!setTargetBallPosition(boundedPosition)) { + addLog("目标球位置设置失败：MuJoCo 尚未加载完成"); + return; + } + addLog(`目标球已移动到${position.label}: (${boundedPosition.x.toFixed(1)}, ${boundedPosition.y.toFixed(1)})`); + }, [addLog, setTargetBallPosition]); + + const handleRandomizeBall = useCallback(() => { + handleSetBallPosition(createRandomBallPosition()); + }, [handleSetBallPosition]); + + const handlePlaceBallOnGround = useCallback((position: { x: number; y: number }) => { + const clampedX = clamp(position.x, -8, 8); + const clampedY = clamp(position.y, -8, 8); + handleSetBallPosition({ + label: "鼠标位置", + x: clampedX, + y: clampedY, + z: 0.25, + }); + }, [handleSetBallPosition]); + + const handleFirstPersonCanvasReady = useCallback((canvas: HTMLCanvasElement | null) => { + firstPersonCanvasRef.current = canvas; + }, []); + + const handleToggleRecording = useCallback(() => { + if (isRecording) { + stopRecording(); + } else { + if (!isLoaded) { + addLog("仿真尚未加载完成，不能开始采集"); + return; + } + startRecording(); + } + }, [addLog, isLoaded, isRecording, startRecording, stopRecording]); + + const handleStartAutoCollection = useCallback(() => { + if (!isLoaded) { + addLog("仿真尚未加载完成，不能开始自动采集"); + return; + } + + setControlMode("auto"); + setAutoInference(false); + if (taskMode === "follow") { + placeVisibleTrainingBall("自动采集起始目标"); + } else { + const state = getCarToBallState(); + avoidBallWasVisibleRef.current = state ? isAvoidBallInView(state) : false; + avoidBallLostViewAtRef.current = null; + avoidAllowMisplacedRetryRef.current = false; + avoidHeadingSettleRef.current = null; + lastAutoBallResetRef.current = performance.now(); + addLog("避开球自动采集起始：保留当前目标球位置"); + } + if (!isRecording) { + startRecording(); + } + addLog(`自动采集已启用: ${taskMode === "follow" ? "跟踪球" : "避开球"}`); + }, [addLog, getCarToBallState, isLoaded, isRecording, placeVisibleTrainingBall, startRecording, taskMode]); + + const refreshModels = useCallback(async () => { + try { + const modelDatasetName = trainingDatasetName || datasetName; + const result = await listModels(userId, modelDatasetName); + const nextModels = result.models || []; + setModels(nextModels); + setSelectedModel((current) => { + if (current && nextModels.includes(current)) return current; + return nextModels.length === 1 ? nextModels[0] : ""; + }); + } catch (error) { + setModels([]); + setSelectedModel(""); + addLog(`模型列表刷新失败: ${error instanceof Error ? error.message : String(error)}`); + } + }, [addLog, datasetName, trainingDatasetName, userId]); + + useEffect(() => { + setIsModelLoaded(false); + setSelectedModel(""); + setAutoInference(false); + setIsTaskInferenceRunning(false); + modelDriveRef.current = { left: 0, right: 0 }; + modelTaskRunRef.current = null; + modelTaskInferenceInFlightRef.current = false; + refreshModels(); + }, [refreshModels]); + + useEffect(() => { + refreshDatasets(); + }, [refreshDatasets]); + + useEffect(() => { + refreshDatasetFrameCount(); + }, [refreshDatasetFrameCount]); + + const stopModelInference = useCallback(() => { + setAutoInference(false); + setIsTaskInferenceRunning(false); + setControlMode("manual"); + if (inferenceTimerRef.current) { + window.clearInterval(inferenceTimerRef.current); + inferenceTimerRef.current = null; + } + modelTaskRunRef.current = null; + modelTaskInferenceInFlightRef.current = false; + modelDriveRef.current = { left: 0, right: 0 }; + setControl("motor_wheel_fl", 0); + setControl("motor_wheel_rl", 0); + setControl("motor_wheel_fr", 0); + setControl("motor_wheel_rr", 0); + }, [setControl]); + + const handleSetTaskMode = useCallback((nextTaskMode: TaskMode) => { + setTaskMode(nextTaskMode); + setIsModelLoaded(false); + setSelectedModel(""); + stopModelInference(); + + const nextDatasetName = TASK_DATASET_NAMES[nextTaskMode]; + if (!isRecording && MANAGED_TASK_DATASET_NAMES.has(datasetName)) { + setDatasetName(nextDatasetName); + } + addLog(`任务切换为: ${nextTaskMode === "follow" ? "跟踪球" : "避开球"}`); + }, [addLog, datasetName, isRecording, stopModelInference]); + + const handleLoadModel = useCallback(async () => { + if (!selectedModel) { + addLog("请先选择一个训练模型"); + return; + } + + if (!trainingDatasetName) { + addLog("请先选择训练数据集，再加载模型"); + return; + } + + const modelDatasetName = trainingDatasetName; + const dataDir = getDatasetPath(userId, modelDatasetName); + const modelPath = `output/train/${userId}/${selectedModel}/model.pt`; + addLog(`加载推理模型: ${modelPath}`); + + try { + const result = await loadTrainedModel(userId, dataDir, modelPath); + if (!result.success) { + addLog(`模型加载失败: ${result.message || result.detail || "unknown error"}`); + return; + } + setIsModelLoaded(true); + setInferenceResult("模型已加载"); + addLog("模型加载成功，可以单次推理或自动推理控制小车"); + } catch (error) { + addLog(`模型加载请求失败: ${error instanceof Error ? error.message : String(error)}`); + } + }, [addLog, selectedModel, trainingDatasetName, userId]); + + const runModelInference = useCallback(async () => { + if (!isModelLoaded) { + addLog("请先加载模型"); + return; + } + + const canvas = firstPersonCanvasRef.current; + if (!canvas) { + addLog("第一视角画面还没准备好，不能推理"); + return; + } + + const { velLeft, velRight } = getWheelVelocityState(); + const image = canvas.toDataURL("image/jpeg", 0.75); + const result = await runInferenceWithSocket(simSocket, [velLeft, velRight], image, userId); + + if (!result.success || !Array.isArray(result.action)) { + throw new Error(result.error || "推理失败"); + } + + const [leftRaw, rightRaw] = result.action; + if (typeof leftRaw !== "number" || typeof rightRaw !== "number") { + throw new Error(`模型输出格式不正确: ${JSON.stringify(result.action)}`); + } + + const left = clampMotorCommand(leftRaw); + const right = clampMotorCommand(rightRaw); + modelDriveRef.current = { left, right }; + setControlMode("model"); + setInferenceResult(`L:${left.toFixed(2)}, R:${right.toFixed(2)}`); + addLog(`推理动作: L=${left.toFixed(2)}, R=${right.toFixed(2)}`); + }, [addLog, getWheelVelocityState, isModelLoaded, userId]); + + const getModelTaskCompletionMessage = useCallback(() => { + const taskRun = modelTaskRunRef.current; + if (!taskRun) return null; + + const state = getCarToBallState(); + if (!state) return null; + + const now = performance.now(); + if (now - taskRun.startedAt >= MODEL_TASK_TIMEOUT_MS) { + return "单次推理超时，已自动停车"; + } + + if (taskMode === "follow") { + const reachedTarget = state.distance <= FOLLOW_TASK_COMPLETE_DISTANCE + && Math.abs(state.angleError) <= FOLLOW_TASK_COMPLETE_ANGLE; + if (!reachedTarget) { + taskRun.followCompleteSince = null; + return null; + } + + taskRun.followCompleteSince = taskRun.followCompleteSince ?? now; + if (now - taskRun.followCompleteSince >= FOLLOW_TASK_COMPLETE_STABLE_MS) { + return "追踪任务完成，已到达目标球附近"; + } + return null; + } + + const ballIsInDangerZone = isAvoidBallInView(state) && state.distance < AVOID_TRIGGER_DISTANCE; + if (ballIsInDangerZone) { + taskRun.avoidDangerSeen = true; + taskRun.avoidSafeSince = null; + return null; + } + + if (!taskRun.avoidDangerSeen) return null; + + const ballIsClear = !isAvoidBallInView(state) || state.distance >= AVOID_TASK_SAFE_DISTANCE; + if (!ballIsClear) { + taskRun.avoidSafeSince = null; + return null; + } + if (now - taskRun.startedAt < AVOID_TASK_MIN_DURATION_MS) return null; + + taskRun.avoidSafeSince = taskRun.avoidSafeSince ?? now; + if (now - taskRun.avoidSafeSince >= AVOID_TASK_SAFE_STABLE_MS) { + return "避让任务完成，目标球已离开前方危险区"; + } + return null; + }, [getCarToBallState, taskMode]); + + const finishModelTask = useCallback((message: string) => { + stopModelInference(); + setInferenceResult(message); + addLog(message); + }, [addLog, stopModelInference]); + + const runModelTaskTick = useCallback(async () => { + if (modelTaskInferenceInFlightRef.current) return; + + const beforeInferenceMessage = getModelTaskCompletionMessage(); + if (beforeInferenceMessage) { + finishModelTask(beforeInferenceMessage); + return; + } + + modelTaskInferenceInFlightRef.current = true; + try { + await runModelInference(); + } catch (error) { + addLog(`单次推理失败: ${error instanceof Error ? error.message : String(error)}`); + stopModelInference(); + return; + } finally { + modelTaskInferenceInFlightRef.current = false; + } + + const afterInferenceMessage = getModelTaskCompletionMessage(); + if (afterInferenceMessage) { + finishModelTask(afterInferenceMessage); + } + }, [addLog, finishModelTask, getModelTaskCompletionMessage, runModelInference, stopModelInference]); + + const handleSingleInference = useCallback(async () => { + if (isTaskInferenceRunning) { + stopModelInference(); + setInferenceResult("单次推理已停止"); + addLog("已停止单次推理"); + return; + } + + if (!isModelLoaded) { + addLog("请先加载模型"); + return; + } + if (autoInference) return; + if (!firstPersonCanvasRef.current) { + addLog("第一视角画面还没准备好，不能执行任务"); + return; + } + + const state = getCarToBallState(); + if (!state) { + addLog("无法读取车辆与目标球状态"); + return; + } + + modelTaskRunRef.current = { + startedAt: performance.now(), + followCompleteSince: null, + avoidDangerSeen: taskMode === "avoid" && isAvoidBallInView(state) && state.distance < AVOID_TRIGGER_DISTANCE, + avoidSafeSince: null, + }; + modelTaskInferenceInFlightRef.current = false; + setControlMode("model"); + setAutoInference(false); + setIsTaskInferenceRunning(true); + addLog(`单次推理已启动: ${taskMode === "follow" ? "追踪球" : "避开球"}`); + + await runModelTaskTick(); + if (!modelTaskRunRef.current) return; + + if (inferenceTimerRef.current) { + window.clearInterval(inferenceTimerRef.current); + } + inferenceTimerRef.current = window.setInterval(() => { + runModelTaskTick(); + }, MODEL_INFERENCE_INTERVAL_MS); + }, [ + addLog, + autoInference, + getCarToBallState, + isModelLoaded, + isTaskInferenceRunning, + runModelTaskTick, + stopModelInference, + taskMode, + ]); + + const handleToggleAutoInference = useCallback(async () => { + if (!isModelLoaded) { + addLog("请先加载模型"); + return; + } + if (isTaskInferenceRunning) { + addLog("单次推理执行中，请先停止"); + return; + } + + if (autoInference) { + stopModelInference(); + setInferenceResult("自动推理已停止"); + addLog("已停止模型自动推理"); + return; + } + + setControlMode("model"); + setAutoInference(true); + addLog("模型自动推理已启动"); + try { + await runModelInference(); + } catch (error) { + addLog(`推理失败: ${error instanceof Error ? error.message : String(error)}`); + } + inferenceTimerRef.current = window.setInterval(() => { + runModelInference().catch((error) => { + addLog(`自动推理失败: ${error instanceof Error ? error.message : String(error)}`); + stopModelInference(); + }); + }, MODEL_INFERENCE_INTERVAL_MS); + }, [addLog, autoInference, isModelLoaded, isTaskInferenceRunning, runModelInference, stopModelInference]); + + const handleStartTraining = useCallback(async () => { + if (isRecording) { + addLog("请先结束采集，再开始训练"); + return; + } + + if (!trainingDatasetName) { + addLog("请先选择训练数据集"); + return; + } + + const dataDir = getDatasetPath(userId, trainingDatasetName); + const outputDir = getTrainPath(userId, trainingDatasetName); + addLog(`启动训练: ${dataDir}`); + + try { + const result = await startTraining(userId, { + data_dir: dataDir, + output_dir: outputDir, + epochs: trainingEpochs, + batch_size: 8, + lr: 1e-4, + }); + + if (!result.success) { + addLog(`训练启动失败: ${result.message || "unknown error"}`); + return; + } + + setTrainingStatus({ + is_running: true, + epoch: 0, + total_epochs: trainingEpochs, + loss: 0, + progress: 0, + error: null, + message: "训练中", + }); + addLog("训练任务已提交"); + } catch (error) { + addLog(`训练请求失败: ${error instanceof Error ? error.message : String(error)}`); + } + }, [addLog, isRecording, trainingDatasetName, trainingEpochs, userId]); + + const handleStopTraining = useCallback(async () => { + await stopTraining(userId); + setTrainingStatus((status) => ({ ...status, is_running: false, message: "训练已停止" })); + addLog("已发送停止训练请求"); + }, [addLog, userId]); + + useEffect(() => { + if (!trainingStatus.is_running) return; + const timer = window.setInterval(async () => { + try { + const status = await getTrainingStatus(userId); + setTrainingStatus(status); + if (!status.is_running && status.error) { + addLog(`训练失败: ${status.error}`); + return; + } + if (!status.is_running && status.progress >= 1) { + addLog(`训练完成: output/train/${userId}/${trainingDatasetName}/model.pt`); + setSelectedModel(trainingDatasetName); + refreshModels(); + } + } catch (error) { + addLog(`训练状态获取失败: ${error instanceof Error ? error.message : String(error)}`); + } + }, 1000); + return () => window.clearInterval(timer); + }, [addLog, refreshModels, trainingDatasetName, trainingStatus.is_running, userId]); + + useEffect(() => { + return () => { + if (inferenceTimerRef.current) { + window.clearInterval(inferenceTimerRef.current); + } + disconnect(); + }; + }, [disconnect]); return ( -
-
+
+
-
+
MJC
-

- MuJoCo WASM + Three.js -

-

Browser-native simulation

+

AKA MuJoCo 自动采集模拟器

+

Browser-native MuJoCo WASM + Three.js

- - FPS: {fps || "--"} - -
- - - {isLoaded ? "Running" : "Loading..."} - +
+ + {isLoaded ? "Simulation Active" : "Loading"}
+ {fps || "--"} FPS
-
+
+ +
+ + +
+
+
WASD / 方向键移动
+
拖动红球移动目标；拖动空地旋转视角
+
-
+
+ {[ + ["指令: 前进", "KeyW"], + ["指令: 左转", "KeyA"], + ["指令: 右转", "KeyD"], + ["指令: 后退", "KeyS"], + ].map(([label, key]) => ( + + ))} + + 本轮帧数: {frameCount} +
+ -
- -

- WASD / Arrows: drive · Space: brake -

-
-
+ +
); } diff --git a/ui/src/pages/MujocoPage/useDataCollection.ts b/ui/src/pages/MujocoPage/useDataCollection.ts index 5887673..dab1f23 100644 --- a/ui/src/pages/MujocoPage/useDataCollection.ts +++ b/ui/src/pages/MujocoPage/useDataCollection.ts @@ -1,53 +1,138 @@ import { useRef, useCallback, useState } from "react"; import { io, Socket } from "socket.io-client"; -import type { MjData } from "@mujoco/mujoco"; +import type { MjModel, MjData } from "@mujoco/mujoco"; + +interface UseDataCollectionOptions { + userId: string; + datasetName: string; + episodeId: number; + fps: number; + taskName: string; + modelRef: React.RefObject; + dataRef: React.RefObject; + fpCanvasRef: React.RefObject; + onLog?: (message: string) => void; + onEpisodeEnded?: (payload: { output_path?: string; frame_count?: number; error?: string }) => void; +} + +const LEFT_WHEEL_JOINTS = ["wheel_fl_joint", "wheel_rl_joint"]; +const RIGHT_WHEEL_JOINTS = ["wheel_fr_joint", "wheel_rr_joint"]; + +function toUint8Array(buf: unknown): Uint8Array | null { + if (typeof buf === "string") { + return new TextEncoder().encode(buf); + } + if (buf instanceof ArrayBuffer) { + return new Uint8Array(buf); + } + if (ArrayBuffer.isView(buf)) { + return new Uint8Array(buf.buffer, buf.byteOffset, buf.byteLength); + } + return null; +} -const CAPTURE_INTERVAL_MS = 100; // 10 FPS -const USER_ID = "sim_user"; +function resolveName(names: unknown, addr: number): string { + const buf = toUint8Array(names); + if (!buf || addr < 0) return ""; + let end = addr; + while (end < buf.length && buf[end] !== 0) end++; + return new TextDecoder().decode(buf.slice(addr, end)); +} -interface Frame { - timestamp: number; - image: string; // base64 JPEG - action: [number, number, number]; // [left, right, gripper] - state: { vel_left: number; vel_right: number }; +function getJointQvel(model: MjModel, data: MjData, jointName: string) { + for (let jointId = 0; jointId < model.njnt; jointId++) { + if (resolveName(model.names, model.name_jntadr[jointId]) === jointName) { + return Number(data.qvel[model.jnt_dofadr[jointId]] ?? 0); + } + } + return 0; } -export function useDataCollection( - dataRef: React.RefObject, - fpCanvasRef: React.RefObject, -) { +function averageJointVelocity(model: MjModel, data: MjData, jointNames: string[]) { + const values = jointNames.map((jointName) => getJointQvel(model, data, jointName)); + return values.reduce((sum, value) => sum + value, 0) / values.length; +} + +function readWheelState(model: MjModel | null, data: MjData | null) { + if (!data) return { velLeft: 0, velRight: 0 }; + if (model) { + return { + velLeft: averageJointVelocity(model, data, LEFT_WHEEL_JOINTS), + velRight: averageJointVelocity(model, data, RIGHT_WHEEL_JOINTS), + }; + } + + const qvel = data.qvel; + const get = (index: number) => (qvel.length > index ? Number(qvel[index]) : 0); + const velLeft = (get(0) + get(2)) / 2; + const velRight = (get(1) + get(3)) / 2; + return { velLeft, velRight }; +} + +export function useDataCollection({ + userId, + datasetName, + episodeId, + fps, + taskName, + modelRef, + dataRef, + fpCanvasRef, + onLog, + onEpisodeEnded, +}: UseDataCollectionOptions) { const [isRecording, setIsRecording] = useState(false); const [frameCount, setFrameCount] = useState(0); - const [episodeId, setEpisodeId] = useState(1); - const [socketStatus, setSocketStatus] = useState("disconnected"); + const [socketStatus, setSocketStatus] = useState<"disconnected" | "connecting" | "connected" | "error">("disconnected"); const socketRef = useRef(null); const recordingRef = useRef(false); const lastCaptureRef = useRef(0); - const episodeIdRef = useRef(1); - const leftVelRef = useRef(0); - const rightVelRef = useRef(0); + const leftTargetRef = useRef(0); + const rightTargetRef = useRef(0); + const localFrameCountRef = useRef(0); - // Connect to backend const connect = useCallback(() => { - if (socketRef.current?.connected) return; - const s = io("http://localhost:8000", { - transports: ["websocket"], - namespace: "/sim", + if (socketRef.current?.connected) return socketRef.current; + + setSocketStatus("connecting"); + const socket = io("/sim", { + path: "/socket.io", + transports: ["websocket", "polling"], + auth: { clientId: `mujoco_${userId}` }, + }); + + socket.on("connect", () => { + setSocketStatus("connected"); + onLog?.("Socket connected to /sim"); }); - s.on("connect", () => setSocketStatus("connected")); - s.on("disconnect", () => setSocketStatus("disconnected")); - s.on("connect_error", () => setSocketStatus("error")); - s.on("episode_started", (data: { episode_id: number }) => { - setEpisodeId(data.episode_id); - episodeIdRef.current = data.episode_id; + socket.on("disconnect", () => { + setSocketStatus("disconnected"); + onLog?.("Socket disconnected"); }); - s.on("collection_count", (data: { count: number }) => { - setFrameCount(data.count); + socket.on("connect_error", () => { + setSocketStatus("error"); + onLog?.("Socket connection failed"); }); - socketRef.current = s; - return s; - }, []); + socket.on("collection_count", (payload: { count: number }) => { + setFrameCount(payload.count); + localFrameCountRef.current = payload.count; + }); + socket.on("episode_started", (payload: { episode_id: number }) => { + onLog?.(`Episode ${payload.episode_id} started`); + }); + socket.on("episode_ended", (payload: { output_path?: string; frame_count?: number; error?: string }) => { + if (payload.error) { + onLog?.(`Export failed: ${payload.error}`); + } else { + onLog?.(`Dataset exported: ${payload.output_path || "unknown path"}`); + } + onEpisodeEnded?.(payload); + }); + + socketRef.current = socket; + return socket; + }, [onEpisodeEnded, onLog, userId]); const disconnect = useCallback(() => { socketRef.current?.disconnect(); @@ -55,92 +140,78 @@ export function useDataCollection( setSocketStatus("disconnected"); }, []); - // Start recording const startRecording = useCallback(() => { - const s = connect(); + const socket = connect(); recordingRef.current = true; setIsRecording(true); setFrameCount(0); + localFrameCountRef.current = 0; lastCaptureRef.current = 0; - s.emit("start_episode", { - user_id: USER_ID, - episode_id: episodeIdRef.current, - task_name: "maze_driving", + socket.emit("start_episode", { + user_id: userId, + episode_id: episodeId, + task_name: taskName, }); - }, [connect]); + onLog?.(`Start collection: ${datasetName}, episode ${episodeId}`); + }, [connect, datasetName, episodeId, onLog, taskName, userId]); - // Stop recording const stopRecording = useCallback(() => { recordingRef.current = false; setIsRecording(false); - const s = socketRef.current; - if (s) { - s.emit("end_episode", { - user_id: USER_ID, - episode_id: episodeIdRef.current, - }); - episodeIdRef.current += 1; - } - }, []); - // Set current action values (called from drive loop) + socketRef.current?.emit("end_episode", { + user_id: userId, + episode_id: episodeId, + }); + onLog?.(`Stop collection: ${localFrameCountRef.current} frames`); + }, [episodeId, onLog, userId]); + const setAction = useCallback((left: number, right: number) => { - leftVelRef.current = left; - rightVelRef.current = right; + leftTargetRef.current = left; + rightTargetRef.current = right; }, []); - // Capture frame + send data (called each physics step) const captureFrame = useCallback(() => { if (!recordingRef.current) return; const now = performance.now(); - if (now - lastCaptureRef.current < CAPTURE_INTERVAL_MS) return; + const intervalMs = 1000 / Math.max(1, fps); + if (now - lastCaptureRef.current < intervalMs) return; lastCaptureRef.current = now; - // Capture FPV canvas - const fpCanvas = fpCanvasRef.current; - if (!fpCanvas) return; - const image = fpCanvas.toDataURL("image/jpeg", 0.7); - - // Read wheel velocities from MuJoCo data - const d = dataRef.current; - let velLeft = 0; - let velRight = 0; - if (d) { - // Wheel joints are at indices 0-3 (fl, fr, rl, rr) - // Joint velocity = d.qvel[joint_dof_adr] - const getJointVel = (name: string) => { - // Simplified: use the first 4 hinge joints - // wheel_fl=0, wheel_fr=1, wheel_rl=2, wheel_rr=3 - const idx = ["wheel_fl", "wheel_fr", "wheel_rl", "wheel_rr"].indexOf(name); - if (idx >= 0 && d.qvel.length > idx) return d.qvel[idx]; - return 0; - }; - velLeft = (getJointVel("wheel_fl") + getJointVel("wheel_rl")) / 2; - velRight = (getJointVel("wheel_fr") + getJointVel("wheel_rr")) / 2; + const canvas = fpCanvasRef.current; + if (!canvas) { + onLog?.("Skip frame: first-person canvas is not ready"); + return; } - const action: [number, number, number] = [ - leftVelRef.current, - rightVelRef.current, - 0, // gripper (not used for car) - ]; + const image = canvas.toDataURL("image/jpeg", 0.75); + const { velLeft, velRight } = readWheelState(modelRef.current, dataRef.current); socketRef.current?.emit("collect_data", { - user_id: USER_ID, - image: image.replace(/^data:image\/jpeg;base64,/, ""), - dataset_name: "maze", - timestamp: now, - state: { vel_left: velLeft, vel_right: velRight }, - action, + user_id: userId, + dataset_name: datasetName, + image, + timestamp: Date.now(), + state: { + vel_left: velLeft, + vel_right: velRight, + }, + action: [ + leftTargetRef.current, + rightTargetRef.current, + 0, + ], }); - }, [dataRef, fpCanvasRef]); + + localFrameCountRef.current += 1; + setFrameCount(localFrameCountRef.current); + }, [dataRef, datasetName, fpCanvasRef, fps, modelRef, onLog, userId]); return { isRecording, frameCount, - episodeId, socketStatus, connect, disconnect, diff --git a/ui/src/pages/MujocoPage/useMujoco.ts b/ui/src/pages/MujocoPage/useMujoco.ts index a9b9eee..c78a977 100644 --- a/ui/src/pages/MujocoPage/useMujoco.ts +++ b/ui/src/pages/MujocoPage/useMujoco.ts @@ -3,6 +3,103 @@ import type { MainModule, MjModel, MjData } from "@mujoco/mujoco"; import wasmUrl from "@mujoco/mujoco/mujoco.wasm?url"; import { CAR_ARM_XML } from "./carArmXml"; +export interface MjPosition { + x: number; + y: number; + z: number; +} + +export interface CarToBallState { + car: MjPosition; + ball: MjPosition; + distance: number; + angleError: number; + frontness: number; + headingAngle: number; + forward: { + x: number; + y: number; + }; +} + +const TARGET_BALL_JOINT = "target_ball_free"; +const LEFT_WHEEL_JOINTS = ["wheel_fl_joint", "wheel_rl_joint"]; +const RIGHT_WHEEL_JOINTS = ["wheel_fr_joint", "wheel_rr_joint"]; + +function toUint8Array(buf: unknown): Uint8Array | null { + if (typeof buf === "string") { + return new TextEncoder().encode(buf); + } + if (buf instanceof ArrayBuffer) { + return new Uint8Array(buf); + } + if (ArrayBuffer.isView(buf)) { + return new Uint8Array(buf.buffer, buf.byteOffset, buf.byteLength); + } + return null; +} + +function resolveName(names: unknown, addr: number): string { + const buf = toUint8Array(names); + if (!buf || addr < 0) return ""; + + let end = addr; + while (end < buf.length && buf[end] !== 0) end++; + return new TextDecoder().decode(buf.slice(addr, end)); +} + +function findNamedIndex( + names: unknown, + nameAddresses: ArrayLike | undefined, + count: number, + targetName: string, +) { + if (!nameAddresses) return -1; + + for (let index = 0; index < count; index++) { + if (resolveName(names, nameAddresses[index]) === targetName) return index; + } + return -1; +} + +function readBodyPosition(model: MjModel, data: MjData, bodyName: string): MjPosition | null { + const bodyId = findNamedIndex(model.names, model.name_bodyadr, model.nbody, bodyName); + if (bodyId < 0) return null; + + return { + x: Number(data.xpos[bodyId * 3]), + y: Number(data.xpos[bodyId * 3 + 1]), + z: Number(data.xpos[bodyId * 3 + 2]), + }; +} + +function readBodyForward2D(model: MjModel, data: MjData, bodyName: string) { + const bodyId = findNamedIndex(model.names, model.name_bodyadr, model.nbody, bodyName); + if (bodyId < 0) return null; + + const quatIndex = bodyId * 4; + const w = Number(data.xquat[quatIndex]); + const x = Number(data.xquat[quatIndex + 1]); + const y = Number(data.xquat[quatIndex + 2]); + const z = Number(data.xquat[quatIndex + 3]); + + return { + x: 1 - 2 * (y * y + z * z), + y: 2 * (x * y + w * z), + }; +} + +function readJointVelocity(model: MjModel, data: MjData, jointName: string) { + const jointId = findNamedIndex(model.names, model.name_jntadr, model.njnt, jointName); + if (jointId < 0) return 0; + return Number(data.qvel[model.jnt_dofadr[jointId]] ?? 0); +} + +function readAverageJointVelocity(model: MjModel, data: MjData, jointNames: string[]) { + const total = jointNames.reduce((sum, jointName) => sum + readJointVelocity(model, data, jointName), 0); + return total / jointNames.length; +} + export function useMujoco() { const [isLoaded, setIsLoaded] = useState(false); const mujocoRef = useRef(null); @@ -81,6 +178,72 @@ export function useMujoco() { m.mj_resetData(model, data); }, []); + const getBodyPosition = useCallback((bodyName: string) => { + const model = modelRef.current; + const data = dataRef.current; + if (!model || !data) return null; + return readBodyPosition(model, data, bodyName); + }, []); + + const getCarToBallState = useCallback((): CarToBallState | null => { + const model = modelRef.current; + const data = dataRef.current; + if (!model || !data) return null; + + const car = readBodyPosition(model, data, "car"); + const ball = readBodyPosition(model, data, "target_ball"); + const forward = readBodyForward2D(model, data, "car"); + if (!car || !ball || !forward) return null; + + const dx = ball.x - car.x; + const dy = ball.y - car.y; + const distance = Math.hypot(dx, dy); + const targetAngle = Math.atan2(dy, dx); + const headingAngle = Math.atan2(forward.y, forward.x); + const angleError = Math.atan2(Math.sin(targetAngle - headingAngle), Math.cos(targetAngle - headingAngle)); + const frontness = distance > 0 ? (forward.x * dx + forward.y * dy) / distance : 0; + + return { car, ball, distance, angleError, frontness, headingAngle, forward }; + }, []); + + const getWheelVelocityState = useCallback(() => { + const model = modelRef.current; + const data = dataRef.current; + if (!model || !data) return { velLeft: 0, velRight: 0 }; + + return { + velLeft: readAverageJointVelocity(model, data, LEFT_WHEEL_JOINTS), + velRight: readAverageJointVelocity(model, data, RIGHT_WHEEL_JOINTS), + }; + }, []); + + const setTargetBallPosition = useCallback((position: MjPosition) => { + const m = mujocoRef.current; + const model = modelRef.current; + const data = dataRef.current; + if (!m || !model || !data) return false; + + const jointId = findNamedIndex(model.names, model.name_jntadr, model.njnt, TARGET_BALL_JOINT); + if (jointId < 0) return false; + + const qposAddress = model.jnt_qposadr[jointId]; + const qvelAddress = model.jnt_dofadr[jointId]; + data.qpos[qposAddress] = position.x; + data.qpos[qposAddress + 1] = position.y; + data.qpos[qposAddress + 2] = position.z; + data.qpos[qposAddress + 3] = 1; + data.qpos[qposAddress + 4] = 0; + data.qpos[qposAddress + 5] = 0; + data.qpos[qposAddress + 6] = 0; + + for (let offset = 0; offset < 6; offset++) { + data.qvel[qvelAddress + offset] = 0; + } + + m.mj_forward(model, data); + return true; + }, []); + return { isLoaded, mujoco: mujocoRef, @@ -89,5 +252,9 @@ export function useMujoco() { step, setControl, reset, + getBodyPosition, + getCarToBallState, + getWheelVelocityState, + setTargetBallPosition, }; } diff --git a/ui/vite.config.ts b/ui/vite.config.ts index 39a3352..4688a8f 100644 --- a/ui/vite.config.ts +++ b/ui/vite.config.ts @@ -1,26 +1,26 @@ -import { defineConfig } from 'vite' -import react from '@vitejs/plugin-react' +import { defineConfig } from "vite"; +import react from "@vitejs/plugin-react"; // https://vite.dev/config/ export default defineConfig({ plugins: [ react({ babel: { - plugins: [['babel-plugin-react-compiler']], + plugins: [["babel-plugin-react-compiler"]], }, }), ], server: { port: 5175, proxy: { - '/api': { - target: 'http://localhost:8000', + "/api": { + target: "http://localhost:8001", changeOrigin: true, }, - '/socket.io': { - target: 'http://localhost:8000', + "/socket.io": { + target: "http://localhost:8001", ws: true, }, }, }, -}) +});