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

This commit is contained in:
2026-07-21 16:15:14 +08:00
parent 9bd22225f9
commit e9e2c946b3
681 changed files with 137221 additions and 8 deletions
@@ -0,0 +1,3 @@
# Cloud Training
See the [Cloud Training](https://mjlab.readthedocs.io/en/latest/source/training/cloud.html) documentation for setup and usage.
@@ -0,0 +1,17 @@
# SkyPilot job that runs a single W&B sweep agent on one GPU.
#
# Submit to an existing cluster provisioned by sweep-cluster.yaml:
# sky exec mjlab-sweep scripts/cloud/sweep-agent.yaml \
# --gpus A100:1 --env SWEEP_ID=<entity/project/sweep_id> -d
resources:
accelerators: A100:1
envs:
SWEEP_ID: ""
MUJOCO_GL: egl
run: |
source "$HOME/.local/bin/env" 2>/dev/null || true
cd ~/sky_workdir
uv run wandb agent "$SWEEP_ID"
@@ -0,0 +1,34 @@
# SkyPilot cluster definition for W&B sweep agents.
#
# Provisions a multi-GPU instance and installs dependencies. Does not
# start any jobs. Use sweep-agent.yaml with sky exec for that.
#
# Usage:
# sky launch scripts/cloud/sweep-cluster.yaml -c mjlab-sweep --gpus A100:8
name: mjlab-sweep-cluster
resources:
cloud: lambda
accelerators: A100:8
autostop:
idle_minutes: 5
down: true
workdir: .
file_mounts:
~/.netrc: ~/.netrc
envs:
MUJOCO_GL: egl
setup: |
# Install EGL for MuJoCo headless rendering.
sudo apt-get update && sudo apt-get install -y libegl-dev
# Install uv if not present.
command -v uv || curl -LsSf https://astral.sh/uv/install.sh | sh
source "$HOME/.local/bin/env" 2>/dev/null || true
uv sync --locked --no-dev
@@ -0,0 +1,61 @@
#!/usr/bin/env bash
# Launch a W&B sweep via SkyPilot.
#
# Provisions a single multi-GPU cluster and runs one sweep agent per GPU
# using SkyPilot's job queue. Each agent pulls hyperparameters from the
# W&B sweep controller independently.
#
# Usage:
# ./scripts/cloud/sweep-launch.sh [GPUS [CLOUD]]
#
# Examples:
# ./scripts/cloud/sweep-launch.sh A100:4 # 4 agents, default cloud
# ./scripts/cloud/sweep-launch.sh A100:8 gcp # 8 agents on GCP
# ./scripts/cloud/sweep-launch.sh A100:8 lambda # 8 agents on Lambda
# ./scripts/cloud/sweep-launch.sh # defaults to A100:4
set -euo pipefail
GPUS="${1:-A100:4}"
CLOUD="${2:-}"
GPU_TYPE="${GPUS%%:*}"
NUM_AGENTS="${GPUS##*:}"
CLUSTER_NAME="mjlab-sweep"
echo "Creating W&B sweep..."
SWEEP_ID=$(uv run wandb sweep scripts/cloud/sweep.yaml 2>&1 | grep "wandb agent" | awk '{print $NF}')
if [ -z "$SWEEP_ID" ]; then
echo "Failed to create sweep."
exit 1
fi
echo "Sweep created: $SWEEP_ID"
echo "Provisioning $GPUS cluster..."
# Provision the cluster and run setup (no run section in this YAML).
CLOUD_FLAG=${CLOUD:+--cloud "$CLOUD"}
sky launch scripts/cloud/sweep-cluster.yaml \
-c "$CLUSTER_NAME" \
--gpus "$GPUS" \
${CLOUD_FLAG} \
-y --retry-until-up
echo "Submitting $NUM_AGENTS agents to job queue..."
for i in $(seq 1 "$NUM_AGENTS"); do
echo " Agent $i/$NUM_AGENTS"
sky exec "$CLUSTER_NAME" \
--gpus "${GPU_TYPE}:1" \
--env "SWEEP_ID=$SWEEP_ID" \
-d \
scripts/cloud/sweep-agent.yaml
done
echo ""
echo "All agents launched. Monitor at:"
echo " sky queue $CLUSTER_NAME"
echo " sky logs $CLUSTER_NAME <JOB_ID>"
echo " W&B dashboard: https://wandb.ai/$SWEEP_ID"
echo ""
echo "When done: sky down $CLUSTER_NAME"
@@ -0,0 +1,37 @@
# Example W&B sweep configuration. Customize the task, parameters, and
# search space for your own experiment.
#
# Usage:
# ./scripts/cloud/sweep-launch.sh A100:4
name: velocity-lr-entropy-sweep
project: mjlab
program: train
method: random
metric:
name: Train/mean_reward
goal: maximize
parameters:
agent.algorithm.learning-rate:
distribution: log_uniform_values
min: 1e-4
max: 1e-2
agent.algorithm.entropy-coef:
distribution: log_uniform_values
min: 0.001
max: 0.1
command:
- ${env}
- uv
- run
- ${program}
- Mjlab-Velocity-Flat-Unitree-G1
- --env.scene.num-envs
- "4096"
- --agent.max-iterations
- "6000"
- ${args}
run_cap: 8
@@ -0,0 +1,47 @@
# SkyPilot task for launching mjlab training on Lambda Cloud.
#
# Uses the pre-built Docker image from GHCR.
#
# Usage:
# sky launch scripts/cloud/train-docker.yaml \
# --env TASK=Mjlab-Velocity-Flat-Unitree-G1
name: mjlab-train
resources:
cloud: lambda
accelerators: A100:1
autostop:
idle_minutes: 5
down: true # Terminates the instance when idle (stops billing).
workdir: .
file_mounts:
~/.netrc: ~/.netrc
envs:
TASK: Mjlab-Velocity-Flat-Unitree-G1
NUM_ENVS: "4096"
MAX_ITERATIONS: "6000"
MUJOCO_GL: egl
setup: |
# Configure NVIDIA runtime for Docker if not already set up.
if ! sudo docker info 2>/dev/null | grep -q "nvidia"; then
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker
sleep 3 # Wait for the daemon to be ready before pulling.
fi
sudo docker pull ghcr.io/mujocolab/mjlab:latest
run: |
sudo docker run --rm --runtime=nvidia --gpus all \
-v "$HOME/.netrc:/root/.netrc:ro" \
-e MUJOCO_GL=egl \
ghcr.io/mujocolab/mjlab:latest \
uv run --no-dev train "$TASK" \
--env.scene.num-envs "$NUM_ENVS" \
--agent.max-iterations "$MAX_ITERATIONS" \
--gpu-ids all
@@ -0,0 +1,44 @@
# SkyPilot task for launching mjlab training on Lambda Cloud.
#
# Installs mjlab directly with uv (no Docker).
#
# Usage:
# sky launch scripts/cloud/train.yaml \
# --env TASK=Mjlab-Velocity-Flat-Unitree-G1
name: mjlab-train
resources:
cloud: lambda
accelerators: A100:1
autostop:
idle_minutes: 5
down: true
workdir: .
file_mounts:
~/.netrc: ~/.netrc
envs:
TASK: Mjlab-Velocity-Flat-Unitree-G1
NUM_ENVS: "4096"
MAX_ITERATIONS: "6000"
MUJOCO_GL: egl
setup: |
# Install EGL for MuJoCo headless rendering.
sudo apt-get update && sudo apt-get install -y libegl-dev
# Install uv if not present.
command -v uv || curl -LsSf https://astral.sh/uv/install.sh | sh
source "$HOME/.local/bin/env" 2>/dev/null || true
uv sync --locked --no-dev
run: |
source "$HOME/.local/bin/env" 2>/dev/null || true
uv run train "$TASK" \
--env.scene.num-envs "$NUM_ENVS" \
--agent.max-iterations "$MAX_ITERATIONS" \
--gpu-ids all