diff --git a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/1-prepare-environment.md b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/1-prepare-environment.md index 0bed838189..6eaba6a649 100644 --- a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/1-prepare-environment.md +++ b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/1-prepare-environment.md @@ -9,7 +9,7 @@ layout: learningpathall ## Inspect the Arm cloud instance -Confirm that the instance uses the Arm64 architecture: +Confirm that the instance reports the `aarch64` architecture: ```bash uname -m @@ -36,9 +36,10 @@ lscpu Save the CPU count and reserve one CPU for operating-system and runtime activity: ```bash -export CORE_COUNT=$(nproc) +CORE_COUNT="$(nproc)" +export CORE_COUNT export RESERVED_CPUS=1 -export WORKLOAD_CPUS=$((CORE_COUNT - RESERVED_CPUS)) +export WORKLOAD_CPUS=$((CORE_COUNT > RESERVED_CPUS ? CORE_COUNT - RESERVED_CPUS : 1)) ``` Verify the values: @@ -56,14 +57,14 @@ The agent count and vectorized environment count control different parts of the For example, three agents and 191 environments means that VMAS simulates 191 navigation worlds concurrently, with three agents in each world. -For CPU sampling in this Learning Path, use one VMAS environment for each workload CPU: +For the reference run, start with one VMAS environment for each workload CPU: ```bash -export N_ENVS=$WORKLOAD_CPUS +export N_ENVS="$WORKLOAD_CPUS" ``` {{% notice Note %}} -This is a workload-sizing policy. VMAS uses vectorized PyTorch operations, so an environment is not permanently mapped to one operating-system thread. +This is a starting point, not a claim that one environment maps to one operating-system thread. VMAS applies vectorized PyTorch operations across the environment batch. Reduce `N_ENVS` if memory pressure causes swapping, and compare throughput before increasing it. {{% /notice %}} ## Install system packages @@ -89,13 +90,13 @@ sudo apt-get install -y git build-essential python3-pip python3-dev pkg-config c Create a Python environment: ```bash -python3.12 -m venv $HOME/venvs/mappo +python3.12 -m venv "$HOME/venvs/mappo" ``` Activate it: ```bash -source $HOME/venvs/mappo/bin/activate +source "$HOME/venvs/mappo/bin/activate" ``` Confirm the active Python interpreter: @@ -110,10 +111,16 @@ Upgrade the Python packaging tools: python -m pip install --upgrade pip setuptools wheel packaging ``` -Install PyTorch: +Install pinned PyTorch, TorchRL, TensorDict, and VMAS versions: ```bash -python -m pip install torch torchvision torchaudio +python -m pip install \ + "torch==2.8.0" \ + "torchvision==0.23.0" \ + "torchaudio==2.8.0" \ + "torchrl==0.10.1" \ + "tensordict==0.10.0" \ + "vmas==1.5.2" ``` Verify the installation: @@ -124,25 +131,35 @@ python -c 'import platform, torch; print("Architecture:", platform.machine()); p ## Install BenchMARL and VMAS -Clone BenchMARL: +Clone the pinned BenchMARL revision: ```bash -cd $HOME -git clone https://github.com/facebookresearch/BenchMARL.git -cd $HOME/BenchMARL +export BENCHMARL_ROOT="$HOME/BenchMARL" +git clone --filter=blob:none --no-checkout --depth 1 \ + https://github.com/facebookresearch/BenchMARL.git \ + "$BENCHMARL_ROOT" +git -C "$BENCHMARL_ROOT" fetch --depth 1 origin \ + 65d649d80e0bdcbdbe2c5d6a3f02dbfed8f0bec1 +git -C "$BENCHMARL_ROOT" checkout --detach \ + 65d649d80e0bdcbdbe2c5d6a3f02dbfed8f0bec1 +cd "$BENCHMARL_ROOT" ``` -Install BenchMARL and VMAS: +Install BenchMARL and its remaining dependencies: ```bash -python -m pip install -e . -python -m pip install vmas +python -m pip install -e ".[vmas]" ``` Verify the software stack: ```bash -python -c 'import torch, torchrl, benchmarl, vmas; print("PyTorch:", torch.__version__); print("TorchRL: OK"); print("BenchMARL: OK"); print("VMAS: OK")' +python - <<'PY' +from importlib.metadata import version + +for package in ("torch", "torchrl", "tensordict", "vmas", "benchmarl"): + print(f"{package}: {version(package)}") +PY ``` Record the BenchMARL revision used for the experiment: @@ -151,4 +168,8 @@ Record the BenchMARL revision used for the experiment: git rev-parse HEAD ``` -Keep this revision with your experiment notes so you can reproduce the software environment later. +The revision must be `65d649d80e0bdcbdbe2c5d6a3f02dbfed8f0bec1`. The version pins and revision keep the checkpoint layout and exporter assumptions reproducible. + +## What you've accomplished + +You have validated the Arm cloud instance, selected an explicit workload size, and installed a pinned training stack. Next, you will configure and run the MAPPO experiment. diff --git a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/2-configure-training.md b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/2-configure-training.md index 2398548a4d..8cf718f4f4 100644 --- a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/2-configure-training.md +++ b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/2-configure-training.md @@ -12,8 +12,9 @@ layout: learningpathall Move to the BenchMARL repository: ```bash -source $HOME/venvs/mappo/bin/activate -cd $HOME/BenchMARL +source "$HOME/venvs/mappo/bin/activate" +export BENCHMARL_ROOT="$HOME/BenchMARL" +cd "$BENCHMARL_ROOT" ``` Use three agents for the reference experiment: @@ -29,17 +30,18 @@ export SAMPLING_DEVICE=cpu export TRAIN_DEVICE=cpu ``` -BenchMARL configures sampling and training devices independently. You can also test `sampling=cpu, train=cuda` or `sampling=cuda, train=cuda` on a compatible system. +BenchMARL configures sampling and training devices independently. This Learning Path keeps both workloads on the Arm CPU so the reference configuration is reproducible. ## Size the CPU sampling workload Recreate the CPU sizing variables so the configuration works after a new SSH login: ```bash -export CORE_COUNT=$(nproc) +CORE_COUNT="$(nproc)" +export CORE_COUNT export RESERVED_CPUS=1 -export WORKLOAD_CPUS=$((CORE_COUNT - RESERVED_CPUS)) -export N_ENVS=$WORKLOAD_CPUS +export WORKLOAD_CPUS=$((CORE_COUNT > RESERVED_CPUS ? CORE_COUNT - RESERVED_CPUS : 1)) +export N_ENVS="$WORKLOAD_CPUS" ``` Collect 100 frames from each environment before every MAPPO update: @@ -94,16 +96,16 @@ mkdir -p "$RUN_DIR" ## Limit CPU thread parallelism -Cap the main CPU thread pools at the number of workload CPUs: +Give PyTorch access to the workload CPUs and keep secondary numerical libraries single-threaded: ```bash -export OMP_NUM_THREADS=$WORKLOAD_CPUS -export MKL_NUM_THREADS=$WORKLOAD_CPUS -export OPENBLAS_NUM_THREADS=$WORKLOAD_CPUS -export NUMEXPR_MAX_THREADS=$WORKLOAD_CPUS +export OMP_NUM_THREADS="$WORKLOAD_CPUS" +export MKL_NUM_THREADS="$WORKLOAD_CPUS" +export OPENBLAS_NUM_THREADS=1 +export NUMEXPR_MAX_THREADS=1 ``` -This avoids library thread pools using more CPU threads than the workload allocation. +This avoids nested OpenBLAS or NumExpr pools competing with PyTorch for every CPU. Monitor memory use and frames per second during the first batches. If the instance swaps or throughput drops, stop the run and retry with a smaller `N_ENVS`. ## Validate the configuration @@ -126,7 +128,28 @@ Do not start training if a required field is blank. Start MAPPO training: ```bash -python benchmarl/run.py algorithm=mappo task=vmas/navigation task.n_agents="$AGENTS" experiment.sampling_device="$SAMPLING_DEVICE" experiment.train_device="$TRAIN_DEVICE" experiment.on_policy_n_envs_per_worker="$N_ENVS" experiment.on_policy_collected_frames_per_batch="$FRAMES_PER_BATCH" 'experiment.loggers=[csv]' experiment.render=false experiment.evaluation=true experiment.max_n_frames="$MAX_FRAMES" experiment.checkpoint_at_end=true experiment.prefer_continuous_actions=true experiment.evaluation_interval="$EVAL_INTERVAL" experiment.evaluation_episodes="$EVAL_EPISODES" experiment.save_folder="$RUN_DIR" +python benchmarl/run.py \ + algorithm=mappo \ + task=vmas/navigation \ + task.n_agents="$AGENTS" \ + experiment.sampling_device="$SAMPLING_DEVICE" \ + experiment.train_device="$TRAIN_DEVICE" \ + experiment.on_policy_n_envs_per_worker="$N_ENVS" \ + experiment.on_policy_collected_frames_per_batch="$FRAMES_PER_BATCH" \ + 'experiment.loggers=[csv]' \ + experiment.create_json=true \ + experiment.render=false \ + experiment.evaluation=true \ + experiment.max_n_frames="$MAX_FRAMES" \ + experiment.checkpoint_at_end=true \ + experiment.prefer_continuous_actions=true \ + experiment.evaluation_interval="$EVAL_INTERVAL" \ + experiment.evaluation_episodes="$EVAL_EPISODES" \ + experiment.save_folder="$RUN_DIR" ``` BenchMARL performs training and periodic evaluation and saves a checkpoint when the run completes. + +## What you've accomplished + +You have configured a CPU-only MAPPO run with bounded thread pools, periodic evaluation, and machine-readable evaluation output. Next, you will inspect the returns and validate the saved checkpoint. diff --git a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/3-validate-checkpoint.md b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/3-validate-checkpoint.md index 5c754a065a..917118398f 100644 --- a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/3-validate-checkpoint.md +++ b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/3-validate-checkpoint.md @@ -1,59 +1,109 @@ --- -title: Locate and validate the training checkpoint -description: Find the trained MAPPO checkpoint and keep the BenchMARL configuration required to reload it. +title: Evaluate and validate the training checkpoint +description: Measure the MAPPO evaluation returns and validate the checkpoint and configuration needed for actor export. weight: 4 ### FIXED, DO NOT MODIFY layout: learningpathall --- -## Find the checkpoint +## Measure the evaluation returns -BenchMARL creates an experiment directory below `RUN_DIR` with a structure similar to: - -```text -$RUN_DIR/ -└── mappo_navigation_mlp__/ - ├── config.pkl - ├── checkpoints/ - │ └── checkpoint_.pt - └── ... -``` - -List the checkpoints created by the run: +BenchMARL writes evaluation episodes to a JSON file because the training command sets `experiment.create_json=true`. Locate the file created in `RUN_DIR`: ```bash -find "$RUN_DIR" -type f -path '*/checkpoints/checkpoint_*.pt' -print +EVAL_JSON="$( +python - <<'PY' +import os +from pathlib import Path + +files = list(Path(os.environ["RUN_DIR"]).rglob("*.json")) +if len(files) != 1: + raise SystemExit(f"Expected one evaluation JSON file, found {len(files)}: {files}") +print(files[0]) +PY +)" +export EVAL_JSON +echo "Evaluation data: $EVAL_JSON" ``` -Select the most recently written checkpoint from this run: +Summarize the first, final, and best mean return across evaluation episodes: ```bash -export CHECKPOINT=$(find "$RUN_DIR" -type f -path '*/checkpoints/checkpoint_*.pt' -printf '%T@ %p\n' | sort -nr | head -1 | cut -d' ' -f2-) +python - <<'PY' +import json +import math +import os +import statistics + +with open(os.environ["EVAL_JSON"], encoding="utf-8") as file: + report = json.load(file) + +steps = [] + +def collect_steps(value): + if not isinstance(value, dict): + return + for key, nested in value.items(): + if ( + key.startswith("step_") + and isinstance(nested, dict) + and "step_count" in nested + and "return" in nested + ): + returns = [float(item) for item in nested["return"]] + if not returns or not all(math.isfinite(item) for item in returns): + raise SystemExit(f"Non-finite or empty returns in {key}") + steps.append((int(nested["step_count"]), statistics.fmean(returns))) + collect_steps(nested) + +collect_steps(report) +if not steps: + raise SystemExit("No evaluation returns found") + +steps.sort() +first_frames, first_return = steps[0] +final_frames, final_return = steps[-1] +best_frames, best_return = max(steps, key=lambda item: item[1]) + +print(f"First mean return: {first_return:.4f} at {first_frames} frames") +print(f"Final mean return: {final_return:.4f} at {final_frames} frames") +print(f"Best mean return: {best_return:.4f} at {best_frames} frames") +print(f"Best improvement: {best_return - first_return:+.4f}") +PY ``` -Display and verify the path: - -```bash -echo "$CHECKPOINT" -test -f "$CHECKPOINT" && echo "Checkpoint found: $CHECKPOINT" -``` +The output reports measured values from your run. A best return greater than the first return is evidence that training improved the evaluated policy. A single run is not a performance benchmark; repeat several seeds before drawing broader conclusions. If the best return doesn't improve, inspect the CSV logs and retry with more frames or different hyperparameters before deploying the policy. -## Keep `config.pkl` with the checkpoint +## Locate the checkpoint -Determine the BenchMARL experiment directory: +The training command requests one checkpoint at the end of the run. Resolve that checkpoint without relying on filename sorting or whitespace-sensitive shell pipelines: ```bash -export SOURCE_EXPERIMENT_DIR=$(dirname "$(dirname "$CHECKPOINT")") +CHECKPOINT="$( +python - <<'PY' +import os +from pathlib import Path + +files = list(Path(os.environ["RUN_DIR"]).glob("*/checkpoints/checkpoint_*.pt")) +if len(files) != 1: + raise SystemExit(f"Expected one final checkpoint, found {len(files)}: {files}") +print(files[0]) +PY +)" +export CHECKPOINT +test -f "$CHECKPOINT" && echo "Checkpoint found: $CHECKPOINT" ``` -Verify both artifacts: +Determine the BenchMARL experiment directory and verify its two required artifacts: ```bash +SOURCE_EXPERIMENT_DIR="$(dirname "$(dirname "$CHECKPOINT")")" +export SOURCE_EXPERIMENT_DIR ls -lh "$CHECKPOINT" "$SOURCE_EXPERIMENT_DIR/config.pkl" ``` -BenchMARL uses the following relative layout when it reconstructs an experiment: +Keep this layout intact when you archive the complete training experiment: ```text / @@ -62,30 +112,54 @@ BenchMARL uses the following relative layout when it reconstructs an experiment: └── checkpoint_.pt ``` -Do not copy only `checkpoint_.pt` when you need to reload the complete BenchMARL experiment. - -## Read deployment metadata from the experiment +## Validate the stored task configuration -Read the actual number of agents stored in the task configuration: - -```bash -export CHECKPOINT_AGENTS=$(python -c "import pickle; f=open('$SOURCE_EXPERIMENT_DIR/config.pkl','rb'); pickle.load(f); cfg=pickle.load(f); print(cfg['n_agents'])") -``` +{{% notice Security %}} +Only load `config.pkl` from a training run you trust. Python pickle files can execute code when loaded. +{{% /notice %}} -Verify it: +Read the task settings through the exported environment variable instead of interpolating a path into Python source: ```bash +CHECKPOINT_AGENTS="$( +python - <<'PY' +import os +import pickle +from pathlib import Path + +config_file = Path(os.environ["SOURCE_EXPERIMENT_DIR"]) / "config.pkl" +with config_file.open("rb") as file: + _task = pickle.load(file) + config = pickle.load(file) + +print(config["n_agents"]) +PY +)" +export CHECKPOINT_AGENTS echo "Checkpoint agents=$CHECKPOINT_AGENTS" ``` -Inspect the full VMAS task configuration: +Inspect the deployment-critical navigation settings: ```bash -python -c "import pickle; f=open('$SOURCE_EXPERIMENT_DIR/config.pkl','rb'); task=pickle.load(f); cfg=pickle.load(f); print('Task:', task); print('Task configuration:', cfg)" +python - <<'PY' +import os +import pickle +from pathlib import Path + +config_file = Path(os.environ["SOURCE_EXPERIMENT_DIR"]) / "config.pkl" +with config_file.open("rb") as file: + task = pickle.load(file) + config = pickle.load(file) + +print("Task:", task) +for key in ("n_agents", "collisions", "observe_all_goals", "lidar_range", "agent_radius", "max_steps"): + print(f"{key}: {config[key]}") +PY ``` -This value comes from the trained experiment and prevents a checkpoint from being registered with a stale shell value for the agent count. +The values must match the navigation configuration used for training. The exporter stops if a deployment-critical value is absent or incompatible. -{{% notice Security %}} -Only load `config.pkl` from a training run you trust. Python pickle files can execute code when loaded. -{{% /notice %}} +## What you've accomplished + +You have measured the policy's evaluation returns and validated the final checkpoint with its trusted task configuration. Next, you will export the actor into a smaller inference artifact. diff --git a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/4-deploy-gui.md b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/4-deploy-gui.md deleted file mode 100644 index 7639402ac1..0000000000 --- a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/4-deploy-gui.md +++ /dev/null @@ -1,239 +0,0 @@ ---- -title: Deploy the checkpoint to the cloud GUI -description: Package the full BenchMARL experiment for the MARL GUI and validate the registered model. -weight: 5 - -### FIXED, DO NOT MODIFY -layout: learningpathall ---- - -## Understand the GUI dependency - -The MARL GUI used in this stage is a separate application from BenchMARL and VMAS. It is not created automatically by the training workflow. - -If you already have the companion GUI code, use that checkout and continue with the steps below. - -If you do not have the GUI code, you can either skip this stage and continue to the actor-export section, or create an equivalent visualization application. A compatible GUI should be able to: - -- accept a BenchMARL `checkpoint_.pt` file and its associated `config.pkl`; -- reload the trained experiment with BenchMARL; -- run a single VMAS environment for interactive evaluation; -- render agent positions, goals, and episode state; -- register and select trained checkpoints; -- preserve the BenchMARL directory layout required by `Experiment.reload_from_file()`. - -The required portable model structure is: - -```text -/ -├── config.pkl -└── checkpoints/ - └── checkpoint_.pt -``` - -The companion GUI used in this Learning Path also provides these utilities: - -```text -tools/deploy_checkpoint.py -tools/inspect_checkpoint.py -``` - -`deploy_checkpoint.py` copies the trained checkpoint and its configuration into the GUI model store and registers the model. `inspect_checkpoint.py` verifies that BenchMARL can reload the checkpoint successfully. - -The remaining commands assume that the companion GUI is available locally. If you create your own GUI, replace the paths and helper-script names below with the corresponding paths in your implementation. - -## Prepare the GUI environment - -The GUI reloads the full BenchMARL experiment, so use the same Python environment that you used for training: - -```bash -source $HOME/venvs/mappo/bin/activate -``` - -Set the GUI root to your local checkout: - -```bash -export GUI_ROOT=$HOME/rl_marl_vmassmapponav_demo/app/marl_live_demo_gui -``` - -Verify that the deployment and inspection tools are available: - -```bash -test -f "$GUI_ROOT/tools/deploy_checkpoint.py" && echo "Deployment tool found" -``` - -```bash -test -f "$GUI_ROOT/tools/inspect_checkpoint.py" && echo "Inspection tool found" -``` - -Install the GUI requirements into the active MAPPO environment: - -```bash -python -m pip install -r "$GUI_ROOT/requirements.txt" -``` - -## Move training artifacts when the GUI is on another system - -The GUI import requires both of these training artifacts: - -```text -$SOURCE_EXPERIMENT_DIR/config.pkl -$SOURCE_EXPERIMENT_DIR/checkpoints/checkpoint_.pt -``` - -If training and the GUI run on different systems, copy the complete experiment directory to the GUI system and preserve this layout. After the copy, update `CHECKPOINT` and `SOURCE_EXPERIMENT_DIR` to the paths on the GUI system. - -If the GUI runs on the same system as training, keep the existing paths. - -## Validate the source checkpoint - -Move to the GUI directory: - -```bash -cd "$GUI_ROOT" -``` - -Inspect the source checkpoint: - -```bash -python tools/inspect_checkpoint.py "$CHECKPOINT" --device cpu -``` - -Find this result in the output: - -```output -"reload": "success" -``` - -If the tool reports `config.pkl file not found in experiment folder`, restore the BenchMARL directory layout before continuing. - -## Register the checkpoint - -The GUI deployment tool is part of the GUI checkout. Run it from its existing path: - -```bash -python "$GUI_ROOT/tools/deploy_checkpoint.py" --checkpoint "$CHECKPOINT" --gui-root "$GUI_ROOT" --instance-label "$(hostname)" --sampling-device "$SAMPLING_DEVICE" --train-device "$TRAIN_DEVICE" -``` - -The deployment tool performs these tasks: - -1. Locates `config.pkl` relative to the source checkpoint. -2. Reads `n_agents` from the trained task configuration. -3. Copies the checkpoint and configuration to the GUI model store. -4. Preserves the `checkpoints/checkpoint_.pt` layout required by BenchMARL. -5. Creates model metadata. -6. Adds the training instance, agent count, run configuration, and checkpoint to `configs/models.yaml`. - -The deployed model has this structure: - -```text -$GUI_ROOT/model_assets/models// -├── config.pkl -├── checkpoints/ -│ └── checkpoint_.pt -├── model_metadata.json -└── source_path.txt -``` - -The `instance` metadata is required because the GUI first filters models by training instance and then builds the available agent-count selections. - -## Verify the GUI model - -Extract the checkpoint step: - -```bash -export CHECKPOINT_STEP=$(basename "$CHECKPOINT" | sed 's/checkpoint_//;s/.pt//') -``` - -Construct the model ID generated by the deployment tool: - -```bash -export MODEL_ID="$(hostname)_agents$(printf '%03d' "$CHECKPOINT_AGENTS")_sampling_${SAMPLING_DEVICE}_train_${TRAIN_DEVICE}_ckpt${CHECKPOINT_STEP}" -``` - -Set the deployed model paths: - -```bash -export GUI_MODEL_DIR="$GUI_ROOT/model_assets/models/$MODEL_ID" -``` - -```bash -export GUI_CHECKPOINT="$GUI_MODEL_DIR/checkpoints/$(basename "$CHECKPOINT")" -``` - -Verify the artifacts: - -```bash -ls -lh "$GUI_MODEL_DIR/config.pkl" "$GUI_CHECKPOINT" "$GUI_MODEL_DIR/model_metadata.json" -``` - -Verify the registry metadata: - -```bash -python -c "import yaml; d=yaml.safe_load(open('$GUI_ROOT/configs/models.yaml')); m=next(x for x in d['models'] if x['id']=='$MODEL_ID'); print('kind=',m.get('kind')); print('instance=',m['metadata'].get('instance')); print('agents=',m['metadata'].get('agent_count')); print('run=',m['metadata'].get('run_name')); print('asset=',m['checkpoint'].get('asset_path'))" -``` - -A valid entry is similar to: - -```output -kind= benchmarl_checkpoint -instance= -agents= 3 -run= sampling_cpu_train_cpu -asset= model_assets/models//checkpoints/checkpoint_.pt -``` - -Validate the deployed checkpoint: - -```bash -python tools/inspect_checkpoint.py "$GUI_CHECKPOINT" --device cpu -``` - -Confirm that the output again contains: - -```output -"reload": "success" -``` - -## Start the GUI - -Start the application: - -```bash -cd "$GUI_ROOT" -``` - -```bash -./run_demo.sh --host 0.0.0.0 --port 8045 -``` - -Refresh the browser after registration. Select the model in this order: - -```text -Training instance - ↓ -Agent count - ↓ -Run configuration - ↓ -Checkpoint -``` - -Select **Start** to run interactive VMAS playback with the trained MAPPO policy. - -{{% notice Note %}} -The GUI uses a single VMAS environment for interactive playback. This is independent of the number of vectorized environments used during training. -{{% /notice %}} - -## Know which artifacts the GUI uses - -| Artifact | Location | Purpose | -| --- | --- | --- | -| Training checkpoint | `$CHECKPOINT` | Complete BenchMARL training state | -| Training configuration | `$SOURCE_EXPERIMENT_DIR/config.pkl` | Reconstructs the experiment | -| GUI checkpoint | `$GUI_MODEL_DIR/checkpoints/checkpoint_.pt` | Portable checkpoint used by the GUI | -| GUI configuration | `$GUI_MODEL_DIR/config.pkl` | Required for BenchMARL reload | -| GUI metadata | `$GUI_MODEL_DIR/model_metadata.json` | Describes the imported model | -| GUI registry | `$GUI_ROOT/configs/models.yaml` | Drives model selection | - -The full BenchMARL checkpoint is the correct artifact for the GUI. For a downstream inference runtime, you can export only the actor parameters, as shown in the next section. diff --git a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/5-export-actor.md b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/5-export-actor.md index cf6a83f466..d90a404a62 100644 --- a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/5-export-actor.md +++ b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/5-export-actor.md @@ -1,594 +1,165 @@ --- -title: Export the MAPPO actor for inference -description: Extract the shared MAPPO actor from a BenchMARL checkpoint and validate a lightweight inference artifact. -weight: 6 +title: Export and validate the MAPPO actor +description: Extract the shared MAPPO actor from the BenchMARL checkpoint and validate a lightweight inference artifact. +weight: 5 ### FIXED, DO NOT MODIFY layout: learningpathall --- -## Understand the actor-only artifact +## Understand the actor artifact -The full BenchMARL checkpoint contains the actor together with state used for training and experiment reload, such as the critic, replay buffer, collector state, and training counters. +The full BenchMARL checkpoint contains the actor, critic, optimizer, collector state, and training counters. A deployment runtime needs only the actor policy and the metadata that defines its input and output contract. -A downstream inference runtime normally needs only the actor policy. - -For the VMAS navigation configuration used in this Learning Path, the shared actor has this structure: +The shared actor trained in this Learning Path has the following structure: ```text 18 observation values ↓ -Linear 18 → 256 - ↓ - Tanh - ↓ -Linear 256 → 256 +Linear 18 → 256 and Tanh ↓ - Tanh +Linear 256 → 256 and Tanh ↓ Linear 256 → 4 + ↓ +2-D deterministic action ``` -For continuous MAPPO, the four raw outputs are split into two location values and two scale values for a two-dimensional action distribution. The validated configuration uses `TanhNormal` and VMAS navigation uses the default action range `[-1, 1]`. For deterministic inference, TorchRL's `TanhNormal.deterministic_sample` therefore corresponds to applying `tanh()` to the two location values. - -The 18-value VMAS navigation observation is: - -```text -2 agent-position values -2 agent-velocity values -2 agent-minus-goal position values -12 LiDAR proximity values -``` - -VMAS constructs each LiDAR proximity value as: - -```text -lidar_range - measured_range -``` +The observation contains two position values, two velocity values, two agent-minus-goal values, and 12 LiDAR proximity values. The four raw outputs define the location and scale of a two-dimensional `TanhNormal` action distribution. Deterministic inference applies `tanh()` to the two location values. -{{% notice Note %}} -The exporter in this section is intentionally specific to the validated shared MAPPO actor. It verifies the task configuration, parameter sharing, MLP layer sizes, activation type, absence of normalization layers, LiDAR-ray count, and actor tensor shapes before it creates an artifact. If any of these assumptions change, the exporter stops rather than silently creating an incompatible model. -{{% /notice %}} +The exporter is deliberately specific to this interface. It stops if the checkpoint, configuration, actor shapes, or action semantics don't match. -## Create the actor exporter +## Download the exporter -Move to the BenchMARL directory: +The reviewed [`export_mappo_actor.py` script](export_mappo_actor.py) is stored with this Learning Path. Download the same version from the Learning Paths repository: ```bash -cd $HOME/BenchMARL -``` - -Create `export_mappo_actor.py` with the following code: - -```python -#!/usr/bin/env python3 - -import argparse -import hashlib -import json -import pickle -from pathlib import Path - -import numpy as np -import torch -import torch.nn.functional as F -from tensordict.nn.distributions import NormalParamExtractor -from torchrl.modules import TanhNormal - - -def sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def cfg_get(cfg, key, default=None): - """Read a setting from a dict-like or attribute-based config object.""" - if isinstance(cfg, dict): - return cfg.get(key, default) - getter = getattr(cfg, "get", None) - if callable(getter): - try: - return getter(key, default) - except (TypeError, AttributeError): - pass - return getattr(cfg, key, default) - - -def class_name(value) -> str: - if isinstance(value, type): - return f"{value.__module__}.{value.__qualname__}" - return str(value) - - -def find_actor_loss_group(state): - matches = [] - for key, value in state.items(): - if not isinstance(key, str) or not key.startswith("loss_"): - continue - if not hasattr(value, "items"): - continue - if any( - isinstance(item_key, str) and "actor_network_params" in item_key - for item_key in value.keys() - ): - matches.append((key, value)) - if len(matches) != 1: - raise RuntimeError( - "Expected exactly one loss group containing actor_network_params; " - f"found {[name for name, _ in matches]}" - ) - return matches[0] - - -def find_actor_tensor(actor_state, suffix): - matches = [ - value - for key, value in actor_state.items() - if isinstance(key, str) - and "actor_network_params" in key - and key.endswith(suffix) - and torch.is_tensor(value) - ] - if len(matches) != 1: - raise RuntimeError( - f"Expected one actor tensor ending in {suffix}; found {len(matches)}" - ) - return matches[0].detach().cpu().to(torch.float32).contiguous() - - -def fail_if(condition, message): - if condition: - raise SystemExit(message) - - -def validate_numpy_export(output, torch_weights, scale_mapping): - """Validate tensor serialization and deterministic TanhNormal semantics.""" - generator = torch.Generator(device="cpu").manual_seed(1234) - observations = torch.randn((16, 18), generator=generator, dtype=torch.float32) - - W1_t, b1_t, W2_t, b2_t, W3_t, b3_t = torch_weights - - with torch.no_grad(): - hidden1 = torch.tanh(F.linear(observations, W1_t, b1_t)) - hidden2 = torch.tanh(F.linear(hidden1, W2_t, b2_t)) - raw_torch = F.linear(hidden2, W3_t, b3_t) - - with np.load(output, allow_pickle=False) as exported: - obs_np = observations.numpy() - hidden1_np = np.tanh(obs_np @ exported["W1"].T + exported["b1"]) - hidden2_np = np.tanh(hidden1_np @ exported["W2"].T + exported["b2"]) - raw_numpy = hidden2_np @ exported["W3"].T + exported["b3"] - - max_raw_error = float(np.max(np.abs(raw_numpy - raw_torch.numpy()))) - if not np.allclose(raw_numpy, raw_torch.numpy(), rtol=1e-5, atol=1e-6): - raise RuntimeError( - "Exported NumPy actor does not match the checkpoint actor MLP. " - f"Maximum raw-output error: {max_raw_error}" - ) - - extractor = NormalParamExtractor(scale_mapping=scale_mapping) - with torch.no_grad(): - loc, scale = extractor(raw_torch) - distribution = TanhNormal(loc, scale, low=-1.0, high=1.0) - torchrl_action = distribution.deterministic_sample - lightweight_action = torch.tanh(raw_torch[..., :2]) - - max_action_error = float( - torch.max(torch.abs(torchrl_action - lightweight_action)).item() - ) - if not torch.allclose( - torchrl_action, lightweight_action, rtol=1e-5, atol=1e-6 - ): - raise RuntimeError( - "tanh(loc) does not match TorchRL TanhNormal deterministic_sample. " - f"Maximum action error: {max_action_error}" - ) - - return max_raw_error, max_action_error - - -def main(): - parser = argparse.ArgumentParser( - description=( - "Export the validated shared BenchMARL MAPPO VMAS-navigation actor " - "to a lightweight NumPy .npz artifact." - ) - ) - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--output", help="Exact .npz output path") - parser.add_argument( - "--output-dir", - default=str(Path.home() / "mappo_actor_exports"), - help="Output directory used when --output is omitted", - ) - args = parser.parse_args() - - checkpoint = Path(args.checkpoint).expanduser().resolve() - requested_output = ( - Path(args.output).expanduser().resolve() if args.output else None - ) - output_dir = Path(args.output_dir).expanduser().resolve() - - fail_if(not checkpoint.is_file(), f"Checkpoint not found: {checkpoint}") - - experiment_dir = checkpoint.parent.parent - config_file = experiment_dir / "config.pkl" - fail_if(not config_file.is_file(), f"config.pkl not found: {config_file}") - - # config.pkl is a trusted BenchMARL artifact from the matching training run. - with config_file.open("rb") as file: - _task = pickle.load(file) - task_config = pickle.load(file) - algorithm_config = pickle.load(file) - model_config = pickle.load(file) - _seed = pickle.load(file) - experiment_config = pickle.load(file) - - n_agents_value = cfg_get(task_config, "n_agents") - fail_if(n_agents_value is None, "Task configuration does not contain n_agents") - n_agents = int(n_agents_value) - - collisions = bool(cfg_get(task_config, "collisions", True)) - observe_all_goals = bool(cfg_get(task_config, "observe_all_goals", False)) - n_lidar_rays = int(cfg_get(task_config, "n_lidar_rays", 12)) - - fail_if(not collisions, "This exporter expects navigation collisions=true") - fail_if( - observe_all_goals, - "This exporter expects observe_all_goals=false; the observation layout would differ", - ) - fail_if( - n_lidar_rays != 12, - f"This exporter expects 12 LiDAR rays; found {n_lidar_rays}", - ) - - fail_if( - not bool(getattr(experiment_config, "share_policy_params", False)), - "This exporter expects share_policy_params=true", - ) - fail_if( - not bool(getattr(algorithm_config, "use_tanh_normal", False)), - "This exporter expects MAPPO use_tanh_normal=true", - ) - - scale_mapping = str( - getattr(algorithm_config, "scale_mapping", "biased_softplus_1.0") - ) - - num_cells = list(getattr(model_config, "num_cells", [])) - activation_class = getattr(model_config, "activation_class", None) - layer_class = getattr(model_config, "layer_class", None) - norm_class = getattr(model_config, "norm_class", None) - - fail_if( - num_cells != [256, 256], - f"This exporter expects hidden layers [256, 256]; found {num_cells}", - ) - fail_if( - class_name(activation_class) != "torch.nn.modules.activation.Tanh", - "This exporter expects torch.nn.Tanh hidden activations; " - f"found {class_name(activation_class)}", - ) - fail_if( - class_name(layer_class) != "torch.nn.modules.linear.Linear", - "This exporter expects torch.nn.Linear layers; " - f"found {class_name(layer_class)}", - ) - fail_if( - norm_class is not None, - "This exporter expects no normalization layer in the actor MLP", - ) - - state = torch.load(checkpoint, map_location="cpu", weights_only=True) - loss_group_name, actor_state = find_actor_loss_group(state) - - W1_t = find_actor_tensor(actor_state, "mlp.params.0.weight") - b1_t = find_actor_tensor(actor_state, "mlp.params.0.bias") - W2_t = find_actor_tensor(actor_state, "mlp.params.2.weight") - b2_t = find_actor_tensor(actor_state, "mlp.params.2.bias") - W3_t = find_actor_tensor(actor_state, "mlp.params.4.weight") - b3_t = find_actor_tensor(actor_state, "mlp.params.4.bias") - - expected_shapes = { - "W1": (256, 18), - "b1": (256,), - "W2": (256, 256), - "b2": (256,), - "W3": (4, 256), - "b3": (4,), - } - torch_tensors = { - "W1": W1_t, - "b1": b1_t, - "W2": W2_t, - "b2": b2_t, - "W3": W3_t, - "b3": b3_t, - } - actual_shapes = {name: tuple(value.shape) for name, value in torch_tensors.items()} - fail_if( - actual_shapes != expected_shapes, - "Unexpected actor architecture. " - f"Expected {expected_shapes}, found {actual_shapes}", - ) - - training_frames = int(state.get("state", {}).get("total_frames", 0)) - agent_group = loss_group_name.removeprefix("loss_") - - output = ( - requested_output - if requested_output is not None - else output_dir / f"mappo_actor_{n_agents}agent_{training_frames}.npz" - ) - fail_if(output.suffix != ".npz", "Actor output file must use the .npz suffix") - - lidar_range = cfg_get(task_config, "lidar_range", 0.35) - agent_radius = cfg_get(task_config, "agent_radius", 0.1) - max_steps = cfg_get(task_config, "max_steps") - - metadata = { - "format": "mappo_shared_actor_numpy_v2", - "source_checkpoint_name": checkpoint.name, - "source_checkpoint_path": str(checkpoint), - "source_checkpoint_sha256": sha256_file(checkpoint), - "source_config_path": str(config_file), - "training_frames": training_frames, - "training_n_agents": n_agents, - "agent_group": agent_group, - "share_policy_params": True, - "actor_input_dim": 18, - "actor_hidden_dims": [256, 256], - "actor_raw_output_dim": 4, - "deterministic_action_dim": 2, - "activation": "tanh", - "policy_distribution": "TanhNormal", - "normal_scale_mapping": scale_mapping, - "action_bounds": [-1.0, 1.0], - "deterministic_action": "TanhNormal.deterministic_sample == tanh(loc)", - "observation_layout": [ - "x_vmas", - "y_vmas", - "vx_vmas", - "vy_vmas", - "agent_x_minus_goal_x_vmas", - "agent_y_minus_goal_y_vmas", - "lidar_proximity_0", - "lidar_proximity_1", - "lidar_proximity_2", - "lidar_proximity_3", - "lidar_proximity_4", - "lidar_proximity_5", - "lidar_proximity_6", - "lidar_proximity_7", - "lidar_proximity_8", - "lidar_proximity_9", - "lidar_proximity_10", - "lidar_proximity_11", - ], - "lidar_encoding": "lidar_range_vmas - measured_range_vmas", - "training_n_lidar_rays": n_lidar_rays, - "training_lidar_range_vmas": lidar_range, - "training_agent_radius_vmas": agent_radius, - "training_max_steps": max_steps, - "training_collisions": collisions, - "training_observe_all_goals": observe_all_goals, - } - - output.parent.mkdir(parents=True, exist_ok=True) - np.savez_compressed( - output, - W1=W1_t.numpy(), - b1=b1_t.numpy(), - W2=W2_t.numpy(), - b2=b2_t.numpy(), - W3=W3_t.numpy(), - b3=b3_t.numpy(), - metadata_json=np.array(json.dumps(metadata)), - ) - - max_raw_error, max_action_error = validate_numpy_export( - output, - (W1_t, b1_t, W2_t, b2_t, W3_t, b3_t), - scale_mapping, - ) - - print(f"Actor export: {output}") - print(f"Agents: {n_agents}") - print(f"Frames: {training_frames}") - print(f"Actor group: {agent_group}") - print(f"Input: {W1_t.shape[1]}") - print(f"Hidden: {W1_t.shape[0]}, {W2_t.shape[0]}") - print(f"Raw output: {W3_t.shape[0]}") - print("Deterministic action: 2-D TanhNormal deterministic_sample") - print(f"NumPy raw parity max err: {max_raw_error:.3e}") - print(f"Action parity max err: {max_action_error:.3e}") - print("Validation: PASS") - - -if __name__ == "__main__": - main() +source "$HOME/venvs/mappo/bin/activate" +cd "$HOME/BenchMARL" +curl -fL \ + https://learn.arm.com/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/export_mappo_actor.py \ + -o export_mappo_actor.py +chmod +x export_mappo_actor.py ``` {{% notice Security %}} -`config.pkl` uses Python pickle serialization. Run the exporter only on `config.pkl` and checkpoint files produced by training runs that you trust. +`config.pkl` uses Python pickle serialization. Run the exporter only on configuration and checkpoint files produced by a training run you trust. {{% /notice %}} -Make the exporter executable: - -```bash -chmod +x export_mappo_actor.py -``` - ## Export the actor -Make sure `CHECKPOINT` still points to the BenchMARL training checkpoint selected in the previous section: - -```bash -echo "$CHECKPOINT" -``` - -The path should point to the original BenchMARL checkpoint layout: - -```text -/ -├── config.pkl -└── checkpoints/ - └── checkpoint_.pt -``` - -Create the actor output directory: - -```bash -export ACTOR_EXPORT_DIR=$HOME/mappo_actor_exports -``` +Create an output name that includes the checkpoint frame count and a short source checksum. This prevents two different training runs from silently sharing a filename: ```bash +export ACTOR_EXPORT_DIR="$HOME/mappo_actor_exports" mkdir -p "$ACTOR_EXPORT_DIR" -``` -Run the exporter: - -```bash -python export_mappo_actor.py --checkpoint "$CHECKPOINT" --output-dir "$ACTOR_EXPORT_DIR" +CHECKPOINT_STEP="${CHECKPOINT##*checkpoint_}" +CHECKPOINT_STEP="${CHECKPOINT_STEP%.pt}" +CHECKPOINT_SHA="$(sha256sum "$CHECKPOINT" | cut -d' ' -f1)" +export ACTOR_OUTPUT="$ACTOR_EXPORT_DIR/mappo_actor_${CHECKPOINT_AGENTS}agent_${CHECKPOINT_STEP}_${CHECKPOINT_SHA:0:12}.npz" ``` -The exporter reads the agent count and training-frame count from the trusted training artifacts and creates a file with the form: +Run the exporter with the explicit output path: -```text -$HOME/mappo_actor_exports/mappo_actor_agent_.npz +```bash +python export_mappo_actor.py \ + --checkpoint "$CHECKPOINT" \ + --output "$ACTOR_OUTPUT" ``` -For example, a three-agent policy trained for 1,910,000 frames produces: - -```text -$HOME/mappo_actor_exports/mappo_actor_3agent_1910000.npz -``` +The exporter refuses to replace an existing file. Use a different output name for another run, or add `--force` only when you intend to atomically replace the artifact. A successful export ends with output similar to: ```output -Actor export: /home/ubuntu/mappo_actor_exports/mappo_actor_3agent_1910000.npz +Actor export: /home/ubuntu/mappo_actor_exports/mappo_actor_3agent_1910000_a1b2c3d4e5f6.npz +Source SHA-256: a1b2c3d4e5f6... Agents: 3 Frames: 1910000 -Actor group: agents Input: 18 Hidden: 256, 256 Raw output: 4 -Deterministic action: 2-D TanhNormal deterministic_sample NumPy raw parity max err: Action parity max err: Validation: PASS ``` -The `.npz` artifact contains only the actor tensors and metadata: - -```text -W1 (256, 18) -b1 (256,) -W2 (256, 256) -b2 (256,) -W3 (4, 256) -b3 (4,) -metadata_json -``` - -The critic and the remaining BenchMARL training state are not included. - -## Understand the exporter validation - -The exporter performs several checks before and after writing the `.npz` file. - -It verifies the training configuration includes: - -```text -share_policy_params = true -use_tanh_normal = true -collisions = true -observe_all_goals = false -12 LiDAR rays -MLP hidden layers = [256, 256] -hidden activation = torch.nn.Tanh -no normalization layer -``` - -It also checks that the checkpoint contains exactly one actor parameter group with these tensor shapes: - -```text -mlp.params.0.weight -> (256, 18) -mlp.params.0.bias -> (256,) -mlp.params.2.weight -> (256, 256) -mlp.params.2.bias -> (256,) -mlp.params.4.weight -> (4, 256) -mlp.params.4.bias -> (4,) -``` - -After saving the actor, the exporter performs two numerical parity checks: +The exporter validates the following contract: -1. It runs the same test observations through the checkpoint tensors in PyTorch and through the saved NumPy actor, and checks that their raw outputs match. -2. It passes the raw outputs through the same `NormalParamExtractor` and `TanhNormal` deterministic-action semantics used by the MAPPO policy, and checks that the lightweight `tanh(loc)` action matches `TanhNormal.deterministic_sample` for the `[-1, 1]` VMAS action range. +- The navigation configuration includes every deployment-critical setting +- The policy shares actor parameters and uses `TanhNormal` +- The actor has two 256-unit `Tanh` layers and no normalization layer +- The actor input implies exactly 12 LiDAR values +- The serialized NumPy tensors reproduce the extracted PyTorch tensor computation +- The deterministic two-value action matches the TorchRL distribution semantics -These checks validate both the weight extraction and the deterministic action interpretation used by the lightweight actor. +It writes to a temporary file, validates that file, and then moves it atomically to `ACTOR_OUTPUT`. The artifact metadata stores source names and a SHA-256 checksum, but it doesn't expose absolute paths from the training host. ## Inspect the exported artifact -Set the actor path using the most recently created export: +Inspect the arrays and metadata without allowing pickled NumPy objects: ```bash -export ACTOR_OUTPUT=$(find "$ACTOR_EXPORT_DIR" -maxdepth 1 -type f -name 'mappo_actor_*agent_*.npz' -printf '%T@ %p -' | sort -nr | head -1 | cut -d' ' -f2-) -``` - -Display the path: +python - <<'PY' +import json +import os -```bash -echo "$ACTOR_OUTPUT" -``` +import numpy as np -Inspect the tensors: +with np.load(os.environ["ACTOR_OUTPUT"], allow_pickle=False) as actor: + for name in actor.files: + print(name, actor[name].shape, actor[name].dtype) + metadata = json.loads(str(actor["metadata_json"])) -```bash -python -c "import numpy as np; d=np.load('$ACTOR_OUTPUT', allow_pickle=False); [print(k, d[k].shape, d[k].dtype) for k in d.files]" +print(json.dumps(metadata, indent=2)) +PY ``` -Inspect the embedded metadata: +Confirm that the archive contains these arrays: -```bash -python -c "import numpy as np, json; d=np.load('$ACTOR_OUTPUT', allow_pickle=False); print(json.dumps(json.loads(str(d['metadata_json'])), indent=2))" +```output +W1 (256, 18) +b1 (256,) +W2 (256, 256) +b2 (256,) +W3 (4, 256) +b3 (4,) +metadata_json () ``` -The metadata records the source checkpoint name and full path, SHA-256 checksum, training frames, training agent count, model dimensions, action interpretation, observation ordering, LiDAR configuration, and other settings needed to identify the policy interface. - ## Run a standalone inference check -You can also run the exported MLP without BenchMARL or TorchRL. Use an all-zero 18-value observation as a basic smoke test: +Run one inference step with an all-zero observation: ```bash -python -c "import numpy as np; d=np.load('$ACTOR_OUTPUT', allow_pickle=False); obs=np.zeros(18,dtype=np.float32); x=np.tanh(d['W1']@obs+d['b1']); x=np.tanh(d['W2']@x+d['b2']); raw=d['W3']@x+d['b3']; action=np.tanh(raw[:2]); print('Action:',action); print('Shape:',action.shape); print('Finite:',np.all(np.isfinite(action))); print('Within [-1,1]:',np.all(np.abs(action)<=1.0))" +python - <<'PY' +import os + +import numpy as np + +with np.load(os.environ["ACTOR_OUTPUT"], allow_pickle=False) as actor: + observation = np.zeros(18, dtype=np.float32) + hidden1 = np.tanh(actor["W1"] @ observation + actor["b1"]) + hidden2 = np.tanh(actor["W2"] @ hidden1 + actor["b2"]) + raw = actor["W3"] @ hidden2 + actor["b3"] + action = np.tanh(raw[:2]) + +print("Action:", action) +print("Shape:", action.shape) +print("Finite:", np.all(np.isfinite(action))) +print("Within [-1, 1]:", np.all(np.abs(action) <= 1.0)) +PY ``` -A successful check includes: +The expected validation fields are: ```output Shape: (2,) Finite: True -Within [-1,1]: True +Within [-1, 1]: True ``` -You now have two forms of the trained policy: - -```text -Full BenchMARL experiment -checkpoint_.pt + config.pkl - ↓ -Cloud MARL GUI - -Actor-only policy -mappo_actor_agent_.npz - ↓ -Downstream inference runtime -``` +## What you've accomplished -The downstream runtime must reproduce the observation ordering and action interpretation recorded in `metadata_json`. The actor-only artifact does not contain the VMAS environment, critic, or robot-specific sensor and control integration. +You have trained and evaluated a MAPPO navigation policy, exported its shared actor without the training-only state, and validated the deployment interface. Continue with [the Device Connect dashboard Learning Path](/learning-paths/servers-and-cloud-computing/use-mappo-device-connect-dashboard/) to distribute and arm the actor through a simulated device. diff --git a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/_index.md b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/_index.md index 28f64d7100..1afa75d8dc 100644 --- a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/_index.md +++ b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/_index.md @@ -1,25 +1,23 @@ --- -title: Train Multi-Agent Reinforcement Learning policies with MAPPO on Arm cloud +title: Train and export a MAPPO navigation policy on Arm cloud draft: true cascade: draft: true -description: Train a MAPPO navigation policy with BenchMARL and VMAS on an Arm cloud instance, deploy the checkpoint to a visualization GUI, and export an actor-only inference artifact. +description: Train a MAPPO navigation policy with BenchMARL and VMAS on an Arm cloud instance, evaluate it, and export an actor-only inference artifact. minutes_to_complete: 330 who_is_this_for: This Learning Path is for cloud and machine learning developers who want to train and package multi-agent reinforcement learning navigation policies on Arm-based servers. learning_objectives: - - Configure a vectorized BenchMARL and VMAS workload for an Arm cloud instance. - - Train and evaluate a multi-agent navigation policy with MAPPO. - - Package and validate the trained BenchMARL checkpoint for a cloud visualization GUI. - - Extract and validate the shared actor as a smaller inference-only artifact. + - Configure a reproducible BenchMARL and VMAS workload for an Arm cloud instance. + - Train and quantitatively evaluate a multi-agent navigation policy with MAPPO. + - Export and validate the shared actor as a smaller inference-only artifact. prerequisites: - - An Arm64 Ubuntu cloud instance with SSH access, `sudo` privileges, and internet access. + - An Arm-based Ubuntu 24.04 cloud instance with SSH access, `sudo` privileges, and internet access. - Familiarity with Linux, Python, PyTorch, and reinforcement learning concepts such as observations, actions, rewards, and policies. - - A local checkout of the companion MARL GUI containing `tools/deploy_checkpoint.py` and `tools/inspect_checkpoint.py` if you want to complete the GUI deployment section. author: - Sagar Surendran @@ -48,6 +46,10 @@ tools_software_languages: operatingsystems: - Linux further_reading: + - resource: + title: Load a MAPPO policy with the Arm Device Connect dashboard + link: /learning-paths/servers-and-cloud-computing/use-mappo-device-connect-dashboard/ + type: website - resource: title: BenchMARL repository link: https://github.com/facebookresearch/BenchMARL @@ -60,6 +62,10 @@ further_reading: title: The Surprising Effectiveness of PPO in Cooperative Multi-Agent Games link: https://arxiv.org/abs/2103.01955 type: website + - resource: + title: TorchRL documentation + link: https://docs.pytorch.org/rl/stable/ + type: documentation ### FIXED, DO NOT MODIFY # ================================================================================ diff --git a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/export_mappo_actor.py b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/export_mappo_actor.py index 8378dfc7f6..9bdbc89c51 100644 --- a/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/export_mappo_actor.py +++ b/content/learning-paths/servers-and-cloud-computing/train-mappo-navigation-arm-cloud/export_mappo_actor.py @@ -3,7 +3,9 @@ import argparse import hashlib import json +import os import pickle +import tempfile from pathlib import Path import numpy as np @@ -34,6 +36,14 @@ def cfg_get(cfg, key, default=None): return getattr(cfg, key, default) +def cfg_require(cfg, key): + """Return a deployment-critical setting, or stop if it was not recorded.""" + value = cfg_get(cfg, key) + if value is None: + raise SystemExit(f"Task configuration does not contain {key}") + return value + + def class_name(value) -> str: if isinstance(value, type): return f"{value.__module__}.{value.__qualname__}" @@ -81,6 +91,33 @@ def fail_if(condition, message): raise SystemExit(message) +def write_atomic_npz(output, arrays, metadata, torch_weights, scale_mapping): + """Write, validate, and atomically install an actor archive.""" + output.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + dir=output.parent, + prefix=f".{output.name}.", + suffix=".tmp", + ) + temporary = Path(temporary_name) + try: + with os.fdopen(descriptor, "wb") as file: + np.savez_compressed( + file, + **arrays, + metadata_json=np.array(json.dumps(metadata)), + ) + file.flush() + os.fsync(file.fileno()) + + errors = validate_numpy_export(temporary, torch_weights, scale_mapping) + os.replace(temporary, output) + return errors + except BaseException: + temporary.unlink(missing_ok=True) + raise + + def validate_numpy_export(output, torch_weights, scale_mapping): """Validate tensor serialization and deterministic TanhNormal semantics.""" generator = torch.Generator(device="cpu").manual_seed(1234) @@ -116,9 +153,7 @@ def validate_numpy_export(output, torch_weights, scale_mapping): max_action_error = float( torch.max(torch.abs(torchrl_action - lightweight_action)).item() ) - if not torch.allclose( - torchrl_action, lightweight_action, rtol=1e-5, atol=1e-6 - ): + if not torch.allclose(torchrl_action, lightweight_action, rtol=1e-5, atol=1e-6): raise RuntimeError( "tanh(loc) does not match TorchRL TanhNormal deterministic_sample. " f"Maximum action error: {max_action_error}" @@ -136,6 +171,11 @@ def main(): ) parser.add_argument("--checkpoint", required=True) parser.add_argument("--output", help="Exact .npz output path") + parser.add_argument( + "--force", + action="store_true", + help="Atomically replace an existing output file", + ) parser.add_argument( "--output-dir", default=str(Path.home() / "mappo_actor_exports"), @@ -144,9 +184,7 @@ def main(): args = parser.parse_args() checkpoint = Path(args.checkpoint).expanduser().resolve() - requested_output = ( - Path(args.output).expanduser().resolve() if args.output else None - ) + requested_output = Path(args.output).expanduser().resolve() if args.output else None output_dir = Path(args.output_dir).expanduser().resolve() fail_if(not checkpoint.is_file(), f"Checkpoint not found: {checkpoint}") @@ -164,24 +202,18 @@ def main(): _seed = pickle.load(file) experiment_config = pickle.load(file) - n_agents_value = cfg_get(task_config, "n_agents") - fail_if(n_agents_value is None, "Task configuration does not contain n_agents") - n_agents = int(n_agents_value) - - collisions = bool(cfg_get(task_config, "collisions", True)) - observe_all_goals = bool(cfg_get(task_config, "observe_all_goals", False)) - n_lidar_rays = int(cfg_get(task_config, "n_lidar_rays", 12)) + n_agents = int(cfg_require(task_config, "n_agents")) + collisions = bool(cfg_require(task_config, "collisions")) + observe_all_goals = bool(cfg_require(task_config, "observe_all_goals")) + lidar_range = float(cfg_require(task_config, "lidar_range")) + agent_radius = float(cfg_require(task_config, "agent_radius")) + max_steps = int(cfg_require(task_config, "max_steps")) fail_if(not collisions, "This exporter expects navigation collisions=true") fail_if( observe_all_goals, "This exporter expects observe_all_goals=false; the observation layout would differ", ) - fail_if( - n_lidar_rays != 12, - f"This exporter expects 12 LiDAR rays; found {n_lidar_rays}", - ) - fail_if( not bool(getattr(experiment_config, "share_policy_params", False)), "This exporter expects share_policy_params=true", @@ -191,9 +223,12 @@ def main(): "This exporter expects MAPPO use_tanh_normal=true", ) - scale_mapping = str( - getattr(algorithm_config, "scale_mapping", "biased_softplus_1.0") + scale_mapping_value = getattr(algorithm_config, "scale_mapping", None) + fail_if( + scale_mapping_value is None, + "Algorithm configuration does not contain scale_mapping", ) + scale_mapping = str(scale_mapping_value) num_cells = list(getattr(model_config, "num_cells", [])) activation_class = getattr(model_config, "activation_class", None) @@ -229,6 +264,12 @@ def main(): W3_t = find_actor_tensor(actor_state, "mlp.params.4.weight") b3_t = find_actor_tensor(actor_state, "mlp.params.4.bias") + n_lidar_rays = int(W1_t.shape[1]) - 6 + fail_if( + n_lidar_rays != 12, + f"This exporter expects 12 LiDAR inputs; actor implies {n_lidar_rays}", + ) + expected_shapes = { "W1": (256, 18), "b1": (256,), @@ -252,26 +293,39 @@ def main(): f"Expected {expected_shapes}, found {actual_shapes}", ) - training_frames = int(state.get("state", {}).get("total_frames", 0)) + training_state = state.get("state") + fail_if( + not hasattr(training_state, "get"), + "Checkpoint does not contain a readable training state", + ) + training_frames_value = training_state.get("total_frames") + fail_if( + training_frames_value is None, + "Checkpoint training state does not contain total_frames", + ) + training_frames = int(training_frames_value) + fail_if(training_frames <= 0, f"Invalid training frame count: {training_frames}") agent_group = loss_group_name.removeprefix("loss_") + checkpoint_sha256 = sha256_file(checkpoint) + output = ( requested_output if requested_output is not None - else output_dir / f"mappo_actor_{n_agents}agent_{training_frames}.npz" + else output_dir + / f"mappo_actor_{n_agents}agent_{training_frames}_{checkpoint_sha256[:12]}.npz" ) fail_if(output.suffix != ".npz", "Actor output file must use the .npz suffix") - - lidar_range = cfg_get(task_config, "lidar_range", 0.35) - agent_radius = cfg_get(task_config, "agent_radius", 0.1) - max_steps = cfg_get(task_config, "max_steps") + fail_if( + output.exists() and not args.force, + f"Output already exists: {output}. Choose another path or pass --force.", + ) metadata = { "format": "mappo_shared_actor_numpy_v2", "source_checkpoint_name": checkpoint.name, - "source_checkpoint_path": str(checkpoint), - "source_checkpoint_sha256": sha256_file(checkpoint), - "source_config_path": str(config_file), + "source_checkpoint_sha256": checkpoint_sha256, + "source_config_name": config_file.name, "training_frames": training_frames, "training_n_agents": n_agents, "agent_group": agent_group, @@ -314,25 +368,23 @@ def main(): "training_observe_all_goals": observe_all_goals, } - output.parent.mkdir(parents=True, exist_ok=True) - np.savez_compressed( - output, - W1=W1_t.numpy(), - b1=b1_t.numpy(), - W2=W2_t.numpy(), - b2=b2_t.numpy(), - W3=W3_t.numpy(), - b3=b3_t.numpy(), - metadata_json=np.array(json.dumps(metadata)), - ) - - max_raw_error, max_action_error = validate_numpy_export( + max_raw_error, max_action_error = write_atomic_npz( output, + { + "W1": W1_t.numpy(), + "b1": b1_t.numpy(), + "W2": W2_t.numpy(), + "b2": b2_t.numpy(), + "W3": W3_t.numpy(), + "b3": b3_t.numpy(), + }, + metadata, (W1_t, b1_t, W2_t, b2_t, W3_t, b3_t), scale_mapping, ) print(f"Actor export: {output}") + print(f"Source SHA-256: {checkpoint_sha256}") print(f"Agents: {n_agents}") print(f"Frames: {training_frames}") print(f"Actor group: {agent_group}")