diff --git a/.gitignore b/.gitignore index d88ac7c..2c6888e 100644 --- a/.gitignore +++ b/.gitignore @@ -236,6 +236,5 @@ docs/superpowers MJLAB-FEASIBILITY.md # Agent configs -AGENTS.md CLAUDE.md .claude/ \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..8bba993 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,101 @@ +# AGENTS.md — so101-nexus + +## Quick commands + +```bash +make format # ruff format + ruff check --fix +make lint # ruff check +make typecheck # ty check (NOT mypy/pyright) +make test # pytest with 84% coverage threshold (COV_FAIL_UNDER=84) +``` + +Run `make format && make lint && make typecheck && make test` before committing. Format first; `lint` and `typecheck` can run in either order, but `test` should run last because it is slowest and depends on valid code. CI runs lint/test/docs separately and does not run `typecheck`, so the local `make typecheck` step is the safety net. + +## Focused testing + +```bash +uv run pytest tests/core/test_rewards.py -q # single file +uv run pytest tests/mujoco/test_envs.py::test_gymnasium_contract -q # single test +uv run pytest tests/core -q # single package/folder +uv run pytest -m "not slow and not visual" -q # skip expensive suites +``` + +## Test environment quirks + +- Tests require headless MuJoCo rendering. `tests/conftest.py` sets `MUJOCO_GL=egl`. +- CI runs tests under `xvfb-run`. Locally, `MUJOCO_GL=egl` should suffice; if you get GL errors, prefix with `xvfb-run -a`. +- `pytest` config lives in `pyproject.toml`: `--import-mode=importlib`, `pythonpath = ["."]`, `testpaths = ["tests"]`. +- Three test markers: `slow` (integration), `visual` (LLM-verified rendering), `warp` (GPU/Warp backend). +- Visual tests need a vision model. Default is `anthropic/claude-sonnet-4-20250514`, gated by `ANTHROPIC_API_KEY`. Override with `VISUAL_TEST_MODEL` and `VISUAL_TEST_API_BASE`; supported providers also include `openai` and `openrouter`. Tests skip cleanly if no model/key is configured. +- Warp tests auto-skip if `mujoco_warp`/`torch` are absent and run on the Warp CPU device. To run: `make test-warp` or `uv run --extra warp pytest tests/warp tests/core/test_rewards_tensor.py -q`. +- Docs-consistency tests: `make docs-check` (requires `pnpm` for the Next.js docs build). + +## Package layout + +- **Source**: `src/so101_nexus/` +- **MuJoCo envs**: `src/so101_nexus/mujoco/` — register via `import so101_nexus.mujoco` +- **Warp envs**: `src/so101_nexus/warp/` — register via `import so101_nexus.warp` +- **Public API**: all top-level exports live in `src/so101_nexus/__init__.py`, NOT in backend submodules. Docs and user code must import from `so101_nexus`, not `so101_nexus.mujoco.*`. +- **Assets**: vendored SO101 model + MuJoCo Menagerie under `src/so101_nexus/assets/` (excluded from wheel; force-included selectively). +- **CLI**: `so101-nexus` entry point → `so101_nexus.cli:main`. + +## Code style & source guardrails + +- Ruff, line-length 100, target `py312`. +- NumPy-style docstrings (`[tool.ruff.lint.pydocstyle] convention = "numpy"`). Google-style `Args:` blocks are forbidden and caught by `tests/core/test_source_guardrails.py`. +- Every package `__init__.py` must define `__all__` (an empty list is acceptable). +- No em dashes, en dashes, or emoji in project-owned source, tests, or user-facing docs. `tests/core/test_source_guardrails.py` and `tests/test_docs_consistency.py` enforce this. +- Ruff enforces: `E`, `F`, `I`, `W`, `D`, `ERA`, `PLR`, `PLW`, `PLC`, `PGH004`, `C4`, `C90`, `SIM`, `RUF`, `UP`, `B`, `PIE`, `PT`, `RET`, `TCH`. +- Key ignores: `F403`/`F405` (wildcard imports in `__init__`), `PLR0913`, `PLR2004`, `PLC0415` (deferred imports in optional deps), `RUF012`, `RUF059`. +- Tests skip `D` (docstring) and `PT006` rules. + +## Type checking + +- Use `ty` (from the `ty` package), not pyright or mypy. +- Many third-party libs (`mujoco`, `mujoco_warp`, `cv2`, `gradio`, `plotly`, `warp`, `tyro`, `transforms3d`) are replaced with `Any` via `[tool.ty.analysis] replace-imports-with-any`. +- `pyrightconfig.json` is present for IDE support only; CI does not run it. + +## Optional extras and dependency groups + +| Extra / group | Purpose | Install | +|---------------|---------|---------| +| `teleop` | Leader-arm teleoperation | `uv sync --extra teleop` | +| `train` | PPO training (torch, wandb, tensorboard) | `uv sync --extra train` | +| `warp` | GPU-parallel MuJoCo Warp | `uv sync --extra warp` | +| `viz` | Pillow visualisation | `uv sync --extra viz` | +| `molmoact` | MolmoAct policy | `uv sync --extra molmoact` | +| `dev` | ruff, ty, test deps | `uv sync --group dev` | +| `visual` | litellm + imageio for visual tests | `uv sync --group visual` | + +CI installs with `uv sync --locked --group dev`. + +## Docs + +- Next.js/Fumadocs site in `docs/`. Build with `pnpm install --frozen-lockfile && pnpm build` (CI uses pnpm 9 and Node 22). +- `make docs-check` runs the Python docs-consistency tests and then builds the Next.js site. +- `test_docs_consistency.py` guards against docs-code drift: no em dashes/emoji in user docs, env nav completeness, public API import paths, and `max_episode_steps` documented as a `gym.make` kwarg rather than a config field. + +## ROCm support (AMD GPUs) + +Two venvs coexist: `.venv` (CUDA, default) and `.venv-rocm` (ROCm 7.2). + +```bash +make rocm-sync # create/update .venv-rocm +make rocm-test # run tests in ROCm venv +make rocm-test-warp # run warp tests in ROCm venv +make rocm-format # format in ROCm venv +make rocm-lint # lint in ROCm venv +make rocm-typecheck # typecheck in ROCm venv +``` + +Or manually: `scripts/setup-rocm.sh` (creates `.venv-rocm` from scratch). Activate: `source .venv-rocm/bin/activate`. + +The setup script installs `torch` last with `--torch-backend rocm7.2 --force-reinstall` so ROCm torch is not overridden by a CUDA torch pulled in transitively. + +## Gotchas + +- `AGENTS.md` and `CLAUDE.md` are gitignored. Keep agent instructions local. +- Coverage threshold is 84% (`COV_FAIL_UNDER`). Run `make test` or `COV_FAIL_UNDER=0 make test` to bypass locally. +- The `warp` extra pins `mujoco-warp>=3.9.0.1,<3.10` and requires `torch`. +- Public configs/classes (`EnvironmentConfig`, `Pose`, `TouchConfig`, etc.) must be imported from `so101_nexus`, not from `so101_nexus.mujoco.*`. +- If teleop crashes during `Creating LeRobot dataset ...` with `module 'triton' has no attribute 'language'`, the `triton` wheel is incomplete. Fix with `uv pip install --force-reinstall triton`. diff --git a/Makefile b/Makefile index ec0ea1b..6f28108 100644 --- a/Makefile +++ b/Makefile @@ -1,8 +1,13 @@ .PHONY: format lint typecheck test test-warp test-visual test-visual-qwen coverage docs-check +.PHONY: rocm-sync rocm-test rocm-test-warp rocm-format rocm-lint rocm-typecheck COV_FAIL_UNDER ?= 84 export COV_FAIL_UNDER +# --------------------------------------------------------------------------- +# CUDA (default) — uses .venv +# --------------------------------------------------------------------------- + format: uv run ruff format uv run ruff check --fix @@ -33,3 +38,35 @@ test-visual-qwen: docs-check: uv run pytest tests/test_docs_consistency.py cd docs && pnpm install --frozen-lockfile && pnpm build + +# --------------------------------------------------------------------------- +# ROCm — uses .venv-rocm (created by scripts/setup-rocm.sh) +# +# make rocm-sync # create/update .venv-rocm with ROCm torch +# make rocm-test # run tests in the ROCm venv +# make rocm-test-warp # run warp tests in the ROCm venv +# make rocm-format # format in the ROCm venv +# make rocm-lint # lint in the ROCm venv +# make rocm-typecheck # typecheck in the ROCm venv +# --------------------------------------------------------------------------- + +ROCM_PY := .venv-rocm/bin/python + +rocm-sync: + scripts/setup-rocm.sh + +rocm-test: + $(ROCM_PY) -m pytest --cov=so101_nexus --cov-report=term-missing --cov-fail-under=$(COV_FAIL_UNDER) + +rocm-test-warp: + $(ROCM_PY) -m pytest tests/warp tests/core/test_rewards_tensor.py -q + +rocm-format: + .venv-rocm/bin/ruff format + .venv-rocm/bin/ruff check --fix + +rocm-lint: + .venv-rocm/bin/ruff check + +rocm-typecheck: + .venv-rocm/bin/ty check diff --git a/scripts/setup-rocm.sh b/scripts/setup-rocm.sh new file mode 100755 index 0000000..8c17a98 --- /dev/null +++ b/scripts/setup-rocm.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# setup-rocm.sh — create .venv-rocm with ROCm 7.2 PyTorch +# +# Usage: +# ./scripts/setup-rocm.sh # create venv and install all deps +# ./scripts/setup-rocm.sh --sync # re-sync only (venv must exist) +set -euo pipefail + +ROCM_VENV=".venv-rocm" +PYTHON="${PYTHON:-python3.12}" + +info() { printf '\033[1;34m%s\033[0m\n' "$*"; } +die() { printf '\033[1;31merror: %s\033[0m\n' "$*" >&2; exit 1; } + +command -v uv >/dev/null || die "uv not found — https://docs.astral.sh/uv/getting-started/installation/" + +SYNC_ONLY=false +[[ "${1:-}" == "--sync" ]] && SYNC_ONLY=true + +# ── create venv if needed ─────────────────────────────────────────────── +if [[ ! -d "$ROCM_VENV" ]]; then + info "Creating $ROCM_VENV ($PYTHON) ..." + uv venv "$ROCM_VENV" --python "$PYTHON" +fi + +PY="$ROCM_VENV/bin/python" + +# ── install project + extras (no-deps first to avoid pulling CUDA torch) ── +info "Installing so101-nexus[teleop,train,warp] (editable, no-deps) ..." +uv pip install --python "$PY" --no-deps -e ".[teleop,train,warp]" + +# ── install remaining deps (excluding torch) ──────────────────────────── +info "Installing remaining dependencies ..." +uv pip install --python "$PY" \ + numpy huggingface_hub trimesh scipy "mujoco>=3.1.3" "gymnasium>=1.0.0" "tyro>=0.9.0" \ + "tensorboard>=2.0.0" "wandb[media]>=0.16.0" \ + "mujoco-warp>=3.9.0.1,<3.10" \ + "lerobot[feetech]>=0.5.0,<0.6" "gradio>=5.0.0" "plotly>=6.0.0" "opencv-python>=4.8.0" \ + "transformers>=5.8" "accelerate" + +# ── install torch (ROCm 7.2) LAST, force-reinstall to override any CUDA ─ +info "Installing torch via --torch-backend rocm7.2 ..." +uv pip install --python "$PY" --torch-backend rocm7.2 --force-reinstall torch torchvision + +# ── install dev/test deps ────────────────────────────────────────────── +info "Installing dev/test deps ..." +uv pip install --python "$PY" \ + pytest pytest-cov hypothesis "imageio[ffmpeg]" Pillow \ + ruff ty litellm + +# ── verify ────────────────────────────────────────────────────────────── +info "Verifying torch ..." +"$PY" -c " +import torch +v = torch.__version__ +hip = getattr(torch.version, 'hip', None) +print(f'torch {v}' + (f' ROCm/HIP {hip}' if hip else ' WARNING: HIP not detected')) +" + +info "Done. Activate: source $ROCM_VENV/bin/activate"