{"id":19855136,"url":"https://github.com/leggedrobotics/terra-baselines","last_synced_at":"2025-09-18T14:31:55.557Z","repository":{"id":228911953,"uuid":"647352825","full_name":"leggedrobotics/terra-baselines","owner":"leggedrobotics","description":"Train, visualize, and evaluate RL policies for the Terra environment.","archived":false,"fork":false,"pushed_at":"2025-09-13T13:42:17.000Z","size":27805,"stargazers_count":16,"open_issues_count":1,"forks_count":1,"subscribers_count":11,"default_branch":"master","last_synced_at":"2025-09-13T15:45:29.585Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/leggedrobotics.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2023-05-30T15:38:51.000Z","updated_at":"2025-09-12T23:28:18.000Z","dependencies_parsed_at":"2024-03-23T07:45:16.592Z","dependency_job_id":"f2733de9-c4ef-4a27-994d-58fa1466d7b1","html_url":"https://github.com/leggedrobotics/terra-baselines","commit_stats":null,"previous_names":["leggedrobotics/terra-baselines"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/leggedrobotics/terra-baselines","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/leggedrobotics%2Fterra-baselines","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/leggedrobotics%2Fterra-baselines/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/leggedrobotics%2Fterra-baselines/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/leggedrobotics%2Fterra-baselines/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/leggedrobotics","download_url":"https://codeload.github.com/leggedrobotics/terra-baselines/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/leggedrobotics%2Fterra-baselines/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":275781371,"owners_count":25527359,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-09-18T02:00:09.552Z","response_time":77,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-12T14:11:48.774Z","updated_at":"2025-09-18T14:31:55.547Z","avatar_url":"https://github.com/leggedrobotics.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# 🌍🚀 Terra Baselines - Training, Evals, and Checkpoints for Terra\nTerra Baselines provides a set of tools to train and evaluate RL policies on the [Terra](https://github.com/leggedrobotics/Terra) environment. This implementation allows to train an agent capable of planning earthworks in trenches and foundations environments in less than 1 minute on 8 Nvidia RTX-4090 GPUs.\n\n## Features\n- Train on multiple devices using PPO with `train.py` (based on [XLand-MiniGrid](https://github.com/corl-team/xland-minigrid))\n- Generate metrics for your checkpoint with `eval.py`\n- Visualize rollouts of your checkpoint with `visualize.py`\n- Run a grid search on the hyperparameters with `train_sweep.py` (orchestrated with [wandb](https://wandb.ai/))\n\n## Installation\nClone the repo and install the requirements with\n```\npip install -r requirements.txt\n```\n\nClone [Terra](https://github.com/leggedrobotics/Terra) in a different folder and install it with\n```\npip install -e .\n```\n\nLastly, [install JAX](https://jax.readthedocs.io/en/latest/installation.html).\n\n## Train\nSetup your training job configuring the `TrainConfig`\n``` python\n@dataclass\nclass TrainConfig:\n    name: str\n    num_devices: int = 0\n    project: str = \"excavator\"\n    group: str = \"default\"\n    num_envs_per_device: int = 4096\n    num_steps: int = 32\n    update_epochs: int = 3\n    num_minibatches: int = 32\n    total_timesteps: int = 3_000_000_000\n    lr: float = 3e-4\n    clip_eps: float = 0.5\n    gamma: float = 0.995\n    gae_lambda: float = 0.95\n    ent_coef: float = 0.01\n    vf_coef: float = 5.0\n    max_grad_norm: float = 0.5\n    eval_episodes: int = 100\n    seed: int = 42\n    log_train_interval: int = 1\n    log_eval_interval: int = 50\n    checkpoint_interval: int = 50\n    clip_action_maps = True\n    local_map_normalization_bounds = [-16, 16]\n    loaded_max = 100\n    num_rollouts_eval = 300\n```\nThen, setup the curriculum in `config.py` in Terra (making sure the maps are saved to disk).\n\nRun a training job with\n```\nDATASET_PATH=/path/to/dataset DATASET_SIZE=\u003cnum_maps_per_type\u003e python train.py -d \u003cnum_devices\u003e\n```\nand collect your weights in the `checkpoints/` folder.\n\n## Sweep\n\nYou can run a hyperparameter sweep over reward settings using [Weights \u0026 Biases Sweeps](https://docs.wandb.ai/guides/sweeps). This allows you to efficiently grid search or random search over reward parameters and compare results.\n\n### 1. Define the Sweep\n\nThe sweep configuration is defined in `sweep.py`. It includes a grid over reward parameters such as `existence`, `collision_move`, `move`, etc. The sweep uses the `TrainConfigSweep` dataclass, which extends the standard training config with sweepable reward parameters.\n\n### 2. Create the Sweep\n\nTo create a new sweep on wandb, run:\n```bash\npython sweep.py create\n```\nThis will print a sweep ID (e.g., `abc123xy`). Copy this ID for the next step.\n\n### 3. Launch Agents\n\nYou can launch multiple agents (workers) to run experiments in parallel. Each agent will pick up a different configuration from the sweep and start a training run.\n\nTo launch an agent, run:\n```bash\nwandb agent \u003cSWEEP_ID\u003e\n```\nYou can run this command multiple times (e.g., in different terminals, or as background jobs in a cluster script) to parallelize the sweep.\n\n#### Example: Running Multiple Agents in a Cluster Script\n\nIf you are using a cluster, you can use the provided `sweep_cluster.sh` script. Make sure to set the `SWEEP_ID` variable to your sweep ID:\n```bash\n# In sweep_cluster.sh\nSWEEP_ID=\u003cYOUR_SWEEP_ID\u003e\nwandb agent $SWEEP_ID \u0026\nwandb agent $SWEEP_ID \u0026\nwandb agent $SWEEP_ID \u0026\nwandb agent $SWEEP_ID \u0026\nwait\n```\n\n## Eval\nEvaluate your checkpoint with standard metrics using\n```\nDATASET_PATH=/path/to/dataset DATASET_SIZE=\u003cnum_maps_per_type\u003e python eval.py -run \u003ccheckpoint_path\u003e -n \u003cnum_environments\u003e -steps \u003cnum_steps\u003e\n```\n\n## Visualize\nVisualize the rollout of your policy with\n```\nDATASET_PATH=/path/to/dataset DATASET_SIZE=\u003cnum_maps_per_type\u003e python visualize.py -run \u003ccheckpoint_path\u003e -nx \u003cnum_environments_x\u003e -ny \u003cnum_environments_y\u003e -steps \u003cnum_steps\u003e -o \u003coutput_path.gif\u003e\n```\n\n## Plan Extraction and Analysis\n\nExtract and analyze terrain modification plans from your trained policies using the plan extraction tools.\n\n### Extract Plans\nExtract action maps and terrain modifications from policy rollouts:\n```bash\npython extract_plan.py -policy \u003ccheckpoint_path\u003e -map \u003cmap_metadata_path\u003e -steps \u003cmax_steps\u003e -o \u003coutput_plan.pkl\u003e\n```\n\nExample:\n```bash\npython extract_plan.py -policy checkpoints/tracked-dense.pkl -map plan_maps/foundation -o plan.pkl\n```\n\nThis captures the robot state and terrain modifications at each DO (dig/dump) action, storing:\n- Agent position, orientation, and loaded state\n- Terrain change values and modification masks\n- Traversability information\n\n### Visualize Plans\nCreate visualizations of the extracted terrain modification plans:\n```bash\npython visualize_plan.py \u003cplan.pkl\u003e\n```\n\nThis generates multi-panel plots showing:\n- Terrain modification masks (where digging/dumping occurred)\n- Traversability maps with agent positions\n- Terrain change values and action map evolution\n- Combined overlays of all modifications\n\n## Baselines\nWe train 2 models capable of solving both foundation and trench type of environments. They differentiate themselves based on the type of agent (wheeled or tracked), and the type of curriculum used to train them (dense reward with single level, or sparse reward with curriculum). All models are trained on 64x64 maps and are stored in the `checkpoints/` folder.\n\n| Checkpoint           | Map Type  | $C_r$ | $S_p$ | $S_w$ | $Coverage$ |\n|----------------------|-----------|-------|-------|-------|------------|\n| `tracked-dense.pkl`  |Foundations|97%|5.66 (1.51)|19.06 (2.86)|0.99 (0.04)|\n|                      |Trenches   |94%|7.09 (5.66)|20.57 (5.26)|0.99 (0.10)|\n| `wheeled-dense.pkl`  |Foundations|99%|11.43 (8.96)|22.06 (3.65)|1.00 (0.00)|\n|                      |Trenches   |89%|15.84 (25.10)|21.12 (5.65)|0.96 (0.14)|\n\nWhere we define the metrics from [Terenzi et al](https://arxiv.org/abs/2308.11478):\n\n$$\n\\begin{equation}\n    \\text{Completion Rate}= C_{r} = \\frac{N_{terminated}}{N_{total}}\n\\end{equation}\n$$\n\n$$\n\\begin{equation}\n    \\text{Path Efficiency}=S_{p}=\\sum_{i=0}^{N-1} \\frac{\\left(x_{B_{i+1}}-x_{B_{i}}\\right)}{\\sqrt{A_{d}}}    \n\\end{equation}\n$$\n\n$$\n\\begin{equation}\n    \\text{Workspace Efficiency} = S_{w} = \\frac{N_{w} \\cdot A_{w}}{A_{d}}    \n\\end{equation}\n$$\n\n$$\n\\begin{equation}\n    \\text{Coverage}=\\frac{N_{tiles\\ dug}}{N_{tiles\\ to\\ dig}}    \n\\end{equation}\n$$\n\n### Model Details\nAll the models we train share the same structure. We encode the maps with a CNN, and the agent state and local maps with MLPs. The latent features are concatenated and shared by the two MLP heads of the model (value and action). In total, the model has ~130k parameters counting both value and action weights.\n\n## Policy Rollouts 😄\nHere's a collection of rollouts for the models we trained.\n####  `tracked-dense.pkl`\n![img](assets/tracked-dense.gif)\n#### `wheeled-dense.pkl`\n![img](assets/wheeled-dense.gif)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fleggedrobotics%2Fterra-baselines","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fleggedrobotics%2Fterra-baselines","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fleggedrobotics%2Fterra-baselines/lists"}