add modified sru-navigation-sim (added go2 code)
This commit is contained in:
MattMa
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025 Fan Yang, Per Frivik, Robotic Systems Lab, ETH Zurich
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,22 @@
|
||||
[package]
|
||||
|
||||
# Note: Semantic Versioning is used: https://semver.org/
|
||||
version = "0.1.0"
|
||||
|
||||
# Description
|
||||
title = "IsaacLab Navigation Extension - SRU Project"
|
||||
description = "IsaacLab task extension for SRU visual navigation with hierarchical control, maze terrain generation, and depth-based RL."
|
||||
readme = "docs/README.md"
|
||||
repository = "https://github.com/michaelfyang/sru-pytorch-spatial-learning"
|
||||
category = "robotics"
|
||||
keywords = ["robotics", "rl", "navigation", "visual", "learning", "sru", "depth", "hierarchical", "maze", "isaaclab"]
|
||||
|
||||
[dependencies]
|
||||
"isaaclab" = {}
|
||||
"isaaclab_assets" = {}
|
||||
|
||||
[core]
|
||||
reloadable = false
|
||||
|
||||
[[python.module]]
|
||||
name = "isaaclab_nav_task"
|
||||
@@ -0,0 +1,482 @@
|
||||
# IsaacLab Navigation Extension - SRU Project
|
||||
|
||||
[](https://journals.sagepub.com/home/ijr)
|
||||
[](https://michaelfyang.github.io/sru-project-website/)
|
||||
|
||||
> **📌 Important Note**: This repository contains the **IsaacLab task extension** for the SRU project, providing diverse navigation environments with dynamic obstacle configurations and terrain variations. This repository does **not** include the `rsl_rl` learning module (network architectures, PPO/MDPO training algorithms). See the [project website](https://michaelfyang.github.io/sru-project-website/) for the complete navigation system.
|
||||
|
||||
## Overview
|
||||
|
||||
A standalone, self-contained IsaacLab task extension for visual navigation in Isaac Lab v2.1.1 (Isaac Sim 4.5). This repository provides:
|
||||
|
||||
- **Environment**: Diverse navigation environments in IsaacLab with dynamic obstacle configurations and terrain variations
|
||||
- **Task Definition**: Hierarchical control architecture interface for visual navigation with reinforcement learning
|
||||
- **Simulation**: High-fidelity physics simulation with realistic depth sensor noise
|
||||
|
||||
**Note**: This repository focuses on the simulation environment and task definition. The RL training infrastructure (neural network architectures, PPO/MDPO algorithms) is provided by the separate `rsl_rl` learning module.
|
||||
|
||||
This extension implements a hierarchical control architecture for visual navigation:
|
||||
- **High-level policy**: Learns to output SE2 velocity commands (vx, vy, omega) at 5Hz
|
||||
- **Low-level policy**: Pre-trained locomotion policy that converts velocity commands to joint actions at 50Hz
|
||||
|
||||
The extension is fully self-contained with all necessary robot models, materials, and pre-trained locomotion policies included.
|
||||
|
||||
### What's Included
|
||||
|
||||
- ✅ IsaacLab task extension for visual navigation environments
|
||||
- ✅ Maze terrain generation with curriculum learning
|
||||
- ✅ Self-contained assets: Robot models (USD), locomotion policies, depth encoders
|
||||
- ✅ Multiple robot platforms: B2W (bipedal wheeled) and AoW-D (Anymal on Wheels)
|
||||
- ✅ Observation definitions: Depth images, proprioception, goal commands
|
||||
- ✅ Reward functions: Goal reaching, action smoothing, movement penalties
|
||||
- ✅ Hierarchical action interface: SE2 velocity commands to low-level controllers
|
||||
- ✅ Domain randomization: Camera pose, action scaling, low-pass filters, sensor delays
|
||||
- ✅ Training scripts compatible with RSL-RL (PPO/MDPO algorithms)
|
||||
|
||||
### What's NOT Included
|
||||
|
||||
- ❌ `rsl_rl` learning module (network architectures, PPO/MDPO training algorithms)
|
||||
- ❌ Neural network structures for high-level navigation policy
|
||||
- ❌ On-policy RL training algorithms (PPO/MDPO implementations)
|
||||
|
||||
**Note**: The `rsl_rl` package must be installed separately to train navigation policies. See the Installation section below.
|
||||
|
||||
### Related Projects
|
||||
|
||||
- [sru-pytorch-spatial-learning](https://github.com/michaelfyang/sru-pytorch-spatial-learning) - Core SRU architecture
|
||||
- [SRU Project Website](https://michaelfyang.github.io/sru-project-website/) - Complete navigation system
|
||||
|
||||
## Features
|
||||
|
||||
- **Visual navigation** using depth cameras with realistic noise simulation
|
||||
- **Maze terrain generation** with curriculum learning
|
||||
- **Self-contained assets**: All robot models and locomotion policies included
|
||||
- **Multiple robot platforms**:
|
||||
- **B2W**: Bipedal wheeled robot (with ZedX camera)
|
||||
- **AoW-D**: Anymal on Wheels (with ZedX camera)
|
||||
- **Asymmetric actor-critic** with privileged critic observations
|
||||
- **Curriculum learning** for terrain difficulty progression
|
||||
- **Multiple algorithms**: MDPO and PPO support via RSL-RL
|
||||
- **Domain randomization**: Camera pose, action scaling, low-pass filters, sensor delays
|
||||
|
||||
## Installation
|
||||
|
||||
### Prerequisites
|
||||
|
||||
- Isaac Lab v2.1.1 installed and configured
|
||||
- Isaac Sim 4.5.0
|
||||
- Python 3.10
|
||||
- PyTorch >= 2.5.1
|
||||
|
||||
### Step 1: Clone or Place the Extension
|
||||
|
||||
This extension should be placed in the `source/` directory of your IsaacLab installation:
|
||||
|
||||
```bash
|
||||
# Navigate to your IsaacLab installation
|
||||
cd /path/to/IsaacLab
|
||||
|
||||
# If cloning this repository separately, place it in source/
|
||||
# Your directory structure should look like:
|
||||
# IsaacLab/
|
||||
# ├── source/
|
||||
# │ ├── isaaclab/
|
||||
# │ ├── isaaclab_assets/
|
||||
# │ └── isaaclab_nav_task/ <- This extension
|
||||
```
|
||||
|
||||
### Step 2: Install the Extension
|
||||
|
||||
Install the extension in development mode from the IsaacLab root directory:
|
||||
|
||||
```bash
|
||||
# From IsaacLab root directory
|
||||
./isaaclab.sh -p -m pip install -e source/isaaclab_nav_task
|
||||
|
||||
# Or navigate to the extension directory
|
||||
cd source/isaaclab_nav_task
|
||||
../../isaaclab.sh -p -m pip install -e .
|
||||
```
|
||||
|
||||
### Step 3: Install RSL-RL (Required for Training)
|
||||
|
||||
This extension requires the `rsl_rl` package for training. Install the custom version with MDPO/PPO algorithms:
|
||||
|
||||
```bash
|
||||
# Clone and install custom rsl_rl (if not already installed)
|
||||
cd /path/to/your/workspace
|
||||
git clone https://github.com/leggedrobotics/rsl_rl.git
|
||||
cd rsl_rl
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
### Verify Installation
|
||||
|
||||
Test that the extension is properly installed:
|
||||
|
||||
```bash
|
||||
# From IsaacLab root directory
|
||||
./isaaclab.sh -p -m pip show isaaclab_nav_task
|
||||
|
||||
# List available tasks
|
||||
./isaaclab.sh -p source/isaaclab_nav_task/scripts/train.py --help
|
||||
```
|
||||
|
||||
You should see the task IDs listed (e.g., `Isaac-Nav-MDPO-B2W-v0`, `Isaac-Nav-PPO-AoW-D-v0`, etc.).
|
||||
|
||||
## Available Tasks
|
||||
|
||||
### B2W
|
||||
| Task ID | Description |
|
||||
|---------|-------------|
|
||||
| `Isaac-Nav-MDPO-B2W-v0` | MDPO training |
|
||||
| `Isaac-Nav-PPO-B2W-v0` | PPO training |
|
||||
| `Isaac-Nav-MDPO-B2W-Play-v0` | MDPO playback |
|
||||
| `Isaac-Nav-PPO-B2W-Play-v0` | PPO playback |
|
||||
| `Isaac-Nav-MDPO-B2W-Dev-v0` | MDPO development |
|
||||
| `Isaac-Nav-PPO-B2W-Dev-v0` | PPO development |
|
||||
|
||||
### AoW-D
|
||||
| Task ID | Description |
|
||||
|---------|-------------|
|
||||
| `Isaac-Nav-MDPO-AoW-D-v0` | MDPO training |
|
||||
| `Isaac-Nav-PPO-AoW-D-v0` | PPO training |
|
||||
| `Isaac-Nav-MDPO-AoW-D-Play-v0` | MDPO playback |
|
||||
| `Isaac-Nav-PPO-AoW-D-Play-v0` | PPO playback |
|
||||
| `Isaac-Nav-MDPO-AoW-D-Dev-v0` | MDPO development |
|
||||
| `Isaac-Nav-PPO-AoW-D-Dev-v0` | PPO development |
|
||||
|
||||
## Training
|
||||
|
||||
### Using the standalone training script
|
||||
|
||||
```bash
|
||||
# Train B2W with PPO
|
||||
./isaaclab.sh -p source/isaaclab_nav_task/scripts/train.py \
|
||||
--task Isaac-Nav-PPO-B2W-v0 --num_envs 4096 --headless
|
||||
|
||||
# Train AoW-D with PPO
|
||||
./isaaclab.sh -p source/isaaclab_nav_task/scripts/train.py \
|
||||
--task Isaac-Nav-PPO-AoW-D-v0 --num_envs 4096 --headless
|
||||
|
||||
# Train with custom wandb run name
|
||||
./isaaclab.sh -p source/isaaclab_nav_task/scripts/train.py \
|
||||
--task Isaac-Nav-MDPO-B2W-v0 --num_envs 4096 --headless \
|
||||
--run_name "experiment_v1_with_curriculum"
|
||||
|
||||
# Train with multiple custom parameters
|
||||
./isaaclab.sh -p source/isaaclab_nav_task/scripts/train.py \
|
||||
--task Isaac-Nav-PPO-B2W-v0 --num_envs 2048 --headless \
|
||||
--run_name "large_training_run" --seed 42 --max_iterations 20000
|
||||
```
|
||||
|
||||
### Development/Testing (smaller config with tensorboard)
|
||||
|
||||
The `-Dev-v0` variants use tensorboard logging instead of wandb and have reduced iterations (300 vs 15000) for quick testing:
|
||||
|
||||
```bash
|
||||
# Quick test with small environment count
|
||||
./isaaclab.sh -p source/isaaclab_nav_task/scripts/train.py \
|
||||
--task Isaac-Nav-PPO-B2W-Dev-v0 --num_envs 32 --headless
|
||||
```
|
||||
|
||||
### Using the standard RSL-RL workflow
|
||||
|
||||
```bash
|
||||
# Train with RSL-RL
|
||||
./isaaclab.sh -p source/standalone/workflows/rsl_rl/train.py \
|
||||
--task Isaac-Nav-MDPO-B2W-v0 --num_envs 4096
|
||||
```
|
||||
|
||||
## Playing Trained Policies
|
||||
|
||||
```bash
|
||||
# Play using standalone script
|
||||
./isaaclab.sh -p source/isaaclab_nav_task/scripts/play.py \
|
||||
--task Isaac-Nav-MDPO-B2W-Play-v0 --num_envs 16
|
||||
|
||||
# Play with specific checkpoint
|
||||
./isaaclab.sh -p source/isaaclab_nav_task/scripts/play.py \
|
||||
--task Isaac-Nav-MDPO-B2W-Play-v0 \
|
||||
--checkpoint /path/to/model.pt
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
isaaclab_nav_task/
|
||||
├── config/
|
||||
│ └── extension.toml # Extension metadata
|
||||
├── docs/
|
||||
│ └── README.md # This file
|
||||
├── scripts/
|
||||
│ ├── train.py # Training script
|
||||
│ └── play.py # Playback script
|
||||
├── setup.py # Installation script
|
||||
├── pyproject.toml # Build configuration
|
||||
└── isaaclab_nav_task/
|
||||
├── __init__.py # Extension entry point
|
||||
└── navigation/
|
||||
├── __init__.py
|
||||
├── navigation_env_cfg.py # Base environment config
|
||||
├── assets/ # Robot configurations and data
|
||||
│ ├── __init__.py
|
||||
│ ├── b2w.py # B2W robot config
|
||||
│ ├── aow_d.py # AoW-D robot config
|
||||
│ └── data/ # Self-contained asset directory
|
||||
│ ├── Robots/ # Robot USD models and materials
|
||||
│ │ └── AoW-D/ # AoW-D robot assets
|
||||
│ │ ├── aow_d.usd # Robot USD model
|
||||
│ │ └── Props/ # Materials and textures
|
||||
│ └── Policies/ # Pre-trained models
|
||||
│ ├── depth_encoder/ # VAE depth encoders
|
||||
│ │ └── vae_pretrain_new.pth (ZedX)
|
||||
│ └── locomotion/ # Low-level locomotion policies
|
||||
│ ├── aow_d/ # policy_blind_3_1.pt (1.7 MB)
|
||||
│ └── b2w/ # policy_b2w_new_2.pt (2.0 MB)
|
||||
├── config/
|
||||
│ ├── rl_cfg.py # Base RL configurations
|
||||
│ ├── b2w/
|
||||
│ │ ├── __init__.py # Task registration
|
||||
│ │ ├── navigation_env_cfg.py
|
||||
│ │ └── agents/
|
||||
│ │ └── rsl_rl_cfg.py
|
||||
│ └── aow_d/
|
||||
│ ├── __init__.py
|
||||
│ ├── navigation_env_cfg.py
|
||||
│ └── agents/
|
||||
│ └── rsl_rl_cfg.py
|
||||
├── mdp/
|
||||
│ ├── observations.py # Observation functions (13 functions)
|
||||
│ ├── rewards.py # Reward functions (5 functions)
|
||||
│ ├── terminations.py # Termination conditions (4 functions)
|
||||
│ ├── curriculums.py # Curriculum terms (1 function)
|
||||
│ ├── events.py # Domain randomization events (5 functions)
|
||||
│ ├── depth_utils/ # Depth processing utilities
|
||||
│ │ ├── __init__.py
|
||||
│ │ ├── camera_config.py # Camera configurations (ZedX)
|
||||
│ │ └── depth_noise_encoder.py # VAE-based depth encoder
|
||||
│ └── navigation/
|
||||
│ ├── goal_commands.py
|
||||
│ ├── goal_commands_cfg.py
|
||||
│ └── actions/
|
||||
│ ├── __init__.py
|
||||
│ ├── navigation_se2_actions.py
|
||||
│ └── navigation_se2_actions_cfg.py
|
||||
└── terrains/ # Custom terrain generators
|
||||
├── __init__.py
|
||||
├── hf_terrains_maze.py # Maze terrain generation
|
||||
├── hf_terrains_maze_cfg.py # Maze terrain configs
|
||||
├── maze_config.py # Maze parameters
|
||||
└── patches.py # TerrainImporter patches
|
||||
```
|
||||
|
||||
## Compatibility
|
||||
|
||||
- **Isaac Lab**: v2.1.1
|
||||
- **Isaac Sim**: 4.5.0
|
||||
- **Python**: 3.10
|
||||
- **PyTorch**: >= 2.5.1
|
||||
|
||||
## Self-Contained Assets
|
||||
|
||||
The extension includes all necessary assets and does not depend on external asset repositories:
|
||||
|
||||
### Robot Models (`assets/data/Robots/`)
|
||||
- **AoW-D**: Complete USD model with materials and textures
|
||||
- Used when AoW-D robots are not available in the base `isaaclab_assets`
|
||||
- Includes all necessary Props and material textures (11 baked textures)
|
||||
|
||||
### Pre-trained Policies (`assets/data/Policies/`)
|
||||
|
||||
**Depth Encoders** (`depth_encoder/`):
|
||||
- `vae_pretrain_new.pth`: ZedX camera encoder for B2W and AoW-D
|
||||
- VAE architecture with RegNet backbone + Feature Pyramid Network
|
||||
|
||||
**Locomotion Policies** (`locomotion/`):
|
||||
- `aow_d/policy_blind_3_1.pt` (1.7 MB): AoW-D wheeled locomotion
|
||||
- `b2w/policy_b2w_new_2.pt` (2.0 MB): B2W bipedal wheeled locomotion
|
||||
|
||||
All locomotion policies are pre-trained and loaded by the hierarchical action controller.
|
||||
|
||||
## Key Components
|
||||
|
||||
### Navigation Environment (`navigation_env_cfg.py`)
|
||||
- Defines the scene with terrain, robot, and sensors
|
||||
- Configures observation groups for policy and critic
|
||||
- Sets up reward terms for goal reaching and movement penalties
|
||||
- Configures curriculum for terrain difficulty
|
||||
|
||||
### MDP Components (`mdp/`)
|
||||
|
||||
**Cleaned and optimized** - removed unused functions to improve maintainability:
|
||||
|
||||
- **observations.py** (13 functions): Depth image processing, proprioception, goal direction, delay buffers
|
||||
- **rewards.py** (5 functions): Goal reaching, action smoothing, movement penalties
|
||||
- **terminations.py** (4 functions): Timeout, collision detection, angle limits, goal reaching
|
||||
- **curriculums.py** (1 function): Backward movement penalty scheduling
|
||||
- **events.py** (5 functions): Camera randomization, action scaling, delay buffer management
|
||||
|
||||
### Navigation Actions (`mdp/navigation/actions/`)
|
||||
- Hierarchical action space with SE2 velocity commands
|
||||
- Integration with pre-trained low-level locomotion policies
|
||||
|
||||
### Terrain Generation and Goal Sampling (`terrains/`)
|
||||
|
||||
The extension includes custom maze terrain generators built on Isaac Lab's terrain generation system, providing diverse navigation environments with safe goal and spawn position sampling.
|
||||
|
||||
**Key Features:**
|
||||
- **Four terrain types**: Maze, non-maze/random, stairs, and pits
|
||||
- **Curriculum learning**: 180 terrains organized in 6 difficulty levels
|
||||
- **Safe position sampling**: Separate padding for goals (0.5m) and spawns (0.6m)
|
||||
- **Mesh optimization**: ~80-99% vertex reduction for large-scale training
|
||||
- **Explicit boolean masks**: Pre-computed valid positions for efficient sampling
|
||||
|
||||
**Terrain Configuration:**
|
||||
- Grid: 6 rows (difficulty) × 30 columns (variations) = 180 terrains
|
||||
- Size: 30m × 30m per terrain with 0.1m resolution (300×300 cells)
|
||||
- Proportions: 30% maze, 20% random, 30% stairs, 20% pits
|
||||
|
||||
For detailed documentation on terrain generation, goal/spawn sampling, coordinate systems, and implementation details, see [TERRAIN_AND_GOALS.md](TERRAIN_AND_GOALS.md).
|
||||
|
||||
### Depth Processing (`mdp/depth_utils/`)
|
||||
- **DepthNoise**: Simulates realistic stereo camera noise using disparity-based filtering
|
||||
- **DepthNoiseEncoder**: VAE-based depth encoder using RegNet backbone with Feature Pyramid Network
|
||||
- **Camera Configurations**: Pre-defined configs for different camera types:
|
||||
|
||||
| Camera | Robots | Resolution | Depth Range | Encoder |
|
||||
|--------|--------|------------|-------------|---------|
|
||||
| ZedX | B2W, AoW-D | 64x40 | 0.25-10.0m | `vae_pretrain_new.pth` |
|
||||
|
||||
### Custom Robot Assets (`assets/`)
|
||||
|
||||
Robot configuration modules define robot-specific parameters:
|
||||
|
||||
**B2W** (`b2w.py`):
|
||||
- Actuator configurations (position/velocity control)
|
||||
- Initial joint states
|
||||
- USD asset path (from base `isaaclab_assets`)
|
||||
|
||||
**AoW-D** (`aow_d.py`):
|
||||
- Actuator configurations for wheeled quadruped
|
||||
- Initial joint states
|
||||
- USD asset path (from local `assets/data/Robots/AoW-D/`)
|
||||
- Uses local robot model when not available in base assets
|
||||
|
||||
Both configurations integrate seamlessly with the hierarchical navigation controller and pre-trained locomotion policies.
|
||||
|
||||
## Docker and Cluster Setup
|
||||
|
||||
### Docker Modifications
|
||||
|
||||
The Dockerfile includes:
|
||||
1. **Custom RSL-RL**: Installs custom `rsl_rl` package in editable mode
|
||||
2. **Git safe directories**: Prevents ownership errors in containers
|
||||
|
||||
### Quick Start Workflow
|
||||
|
||||
```bash
|
||||
# 1. Build Docker image
|
||||
./docker/container.sh start --suffix nav
|
||||
|
||||
# 2. Push to cluster (converts to Singularity automatically)
|
||||
./docker/cluster/cluster_interface.sh push base-nav
|
||||
|
||||
# 3. Submit training job
|
||||
./docker/cluster/cluster_interface.sh job base-nav \
|
||||
"--task Isaac-Nav-PPO-B2W-v0" \
|
||||
"--num_envs 2048" \
|
||||
"--max_iterations 10000" \
|
||||
"--headless"
|
||||
|
||||
# 4. Monitor job
|
||||
squeue -u $USER
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
**Step 1**: Create `.env.base-nav` profile in `docker/` directory:
|
||||
```bash
|
||||
cp docker/.env.base docker/.env.base-nav
|
||||
```
|
||||
|
||||
**Step 2**: Configure `docker/cluster/.env.cluster` before deployment:
|
||||
- Set `CLUSTER_PYTHON_EXECUTABLE=source/isaaclab_nav_task/scripts/train.py`
|
||||
- Add cluster credentials and paths
|
||||
|
||||
**Step 3**: Add cluster-specific module loads in `docker/cluster/submit_job_slurm.sh`:
|
||||
```bash
|
||||
module load eth_proxy # Required for network access on ETH cluster
|
||||
```
|
||||
|
||||
See the [IsaacLab cluster guide](https://isaac-sim.github.io/IsaacLab/main/source/deployment/cluster.html#cluster-guide) for details.
|
||||
|
||||
### Training Examples
|
||||
|
||||
```bash
|
||||
# B2W with MDPO training (10k iterations)
|
||||
./docker/cluster/cluster_interface.sh job base-nav \
|
||||
"--task Isaac-Nav-MDPO-B2W-v0" \
|
||||
"--num_envs 2048" \
|
||||
"--max_iterations 10000" \
|
||||
"--headless"
|
||||
|
||||
# B2W with custom run name
|
||||
./docker/cluster/cluster_interface.sh job base-nav \
|
||||
"--task Isaac-Nav-MDPO-B2W-v0" \
|
||||
"--num_envs 2048" \
|
||||
"--max_iterations 10000" \
|
||||
"--run_name experiment_v1_b2w" \
|
||||
"--headless"
|
||||
|
||||
# AoW-D with MDPO training (10k iterations)
|
||||
./docker/cluster/cluster_interface.sh job base-nav \
|
||||
"--task Isaac-Nav-MDPO-AoW-D-v0" \
|
||||
"--num_envs 2048" \
|
||||
"--max_iterations 10000" \
|
||||
"--headless"
|
||||
|
||||
# Quick dev test with PPO training (300 iters, tensorboard)
|
||||
./docker/cluster/cluster_interface.sh job base-nav \
|
||||
"--task Isaac-Nav-PPO-B2W-Dev-v0" \
|
||||
"--num_envs 32" \
|
||||
"--headless"
|
||||
```
|
||||
|
||||
### Troubleshooting
|
||||
|
||||
**Git ownership errors**: Rebuild Docker image (includes fix) or run in container:
|
||||
```bash
|
||||
git config --global --add safe.directory '*'
|
||||
```
|
||||
|
||||
**Memory issues**: Reduce `--num_envs` or increase `#SBATCH --mem-per-cpu`
|
||||
|
||||
## License
|
||||
|
||||
MIT License - See [LICENSE](../LICENSE) file for details
|
||||
|
||||
Copyright (c) 2025 Fan Yang, Per Frivik, Robotic Systems Lab, ETH Zurich
|
||||
|
||||
## Citation
|
||||
|
||||
If you use this codebase in your research, please cite:
|
||||
|
||||
```bibtex
|
||||
@article{yang2025sru,
|
||||
author = {Yang, Fan and Frivik, Per and Hoeller, David and Wang, Chen and Cadena, Cesar and Hutter, Marco},
|
||||
title = {Spatially-enhanced recurrent memory for long-range mapless navigation via end-to-end reinforcement learning},
|
||||
journal = {The International Journal of Robotics Research},
|
||||
year = {2025},
|
||||
doi = {10.1177/02783649251401926},
|
||||
url = {https://doi.org/10.1177/02783649251401926}
|
||||
}
|
||||
```
|
||||
|
||||
## Contact
|
||||
|
||||
**Authors**:
|
||||
- Fan Yang (fanyang1@ethz.ch)
|
||||
- Per Frivik (pfrivik@ethz.ch)
|
||||
|
||||
**Affiliation**: Robotic Systems Lab, ETH Zurich
|
||||
@@ -0,0 +1,340 @@
|
||||
# Terrain Generation and Goal Sampling
|
||||
|
||||
This document provides detailed information about the terrain generation system and goal/spawn position sampling used in the SRU navigation simulation.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Terrain Generation System](#terrain-generation-system)
|
||||
- [Architecture Overview](#architecture-overview)
|
||||
- [Key Files](#key-files)
|
||||
- [Mesh Optimization](#mesh-optimization)
|
||||
- [Terrain Data Flow](#terrain-data-flow)
|
||||
- [Maze Terrain Types](#maze-terrain-types)
|
||||
- [Safety Padding](#safety-padding)
|
||||
- [Terrain Configuration](#terrain-configuration)
|
||||
- [Curriculum Learning](#curriculum-learning)
|
||||
- [Goal and Spawn Sampling](#goal-and-spawn-sampling)
|
||||
- [Architecture](#architecture)
|
||||
- [Key Features](#key-features)
|
||||
- [Coordinate System](#coordinate-system)
|
||||
- [Terrain Index Mapping](#terrain-index-mapping)
|
||||
- [Position Sampling](#position-sampling)
|
||||
- [Robot Spawn Height](#robot-spawn-height)
|
||||
|
||||
## Terrain Generation System
|
||||
|
||||
The extension includes custom maze terrain generators built on Isaac Lab's terrain generation system.
|
||||
|
||||
### Architecture Overview
|
||||
|
||||
```
|
||||
Terrain Generation Flow:
|
||||
┌────────────────────────────────────────────────────────────────────┐
|
||||
│ 1. HfMazeTerrainCfg │
|
||||
│ └─► maze_terrain() generates: │
|
||||
│ - heights: Height field for physics/rendering │
|
||||
│ - valid_mask: Valid goal positions (GOAL_PADDING=5 cells) │
|
||||
│ - spawn_mask: Valid spawn positions (SPAWN_PADDING=6 cells)│
|
||||
│ - platform_mask: Elevated platforms for curriculum │
|
||||
│ │
|
||||
│ 2. TerrainGenerator (patched) │
|
||||
│ └─► Collects height field data from all sub-terrains │
|
||||
│ └─► Concatenates into single tensors per attribute │
|
||||
│ │
|
||||
│ 3. TerrainImporter (patched) │
|
||||
│ └─► Stores on self._height_field_* attributes │
|
||||
│ │
|
||||
│ 4. RobotNavigationGoalCommand │
|
||||
│ └─► Reads from env.scene.terrain._height_field_* │
|
||||
│ └─► Creates PositionSampler with both masks │
|
||||
└────────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Key Files
|
||||
|
||||
| File | Purpose |
|
||||
|------|---------|
|
||||
| `hf_terrains_maze.py` | Terrain generation with explicit valid_mask/spawn_mask |
|
||||
| `hf_terrains_maze_cfg.py` | Configuration dataclass (lines 41-51: mask storage attributes) |
|
||||
| `terrain_constants.py` | Constants: PADDING (5/6 cells), HEIGHTS (0/200/300), VERTICAL_SCALE (0.005) |
|
||||
| `patches.py` | Monkey-patches TerrainGenerator/TerrainImporter for mask storage |
|
||||
| `maze_config.py` | MAZE_TERRAIN_CFG with sub-terrain configurations |
|
||||
|
||||
### Mesh Optimization
|
||||
|
||||
The extension includes automatic mesh optimization that significantly reduces GPU memory usage when training with many environments. This is especially important for large-scale RL training (4096+ environments).
|
||||
|
||||
**How it works:**
|
||||
- Uses hierarchical block-based approach (20x20 → 10x10 → 5x5 blocks)
|
||||
- Flat terrain regions are simplified to just 2 triangles instead of full mesh detail
|
||||
- Non-flat regions recursively subdivide until 5x5 blocks, then generate detailed mesh
|
||||
- Applied automatically via monkey-patching when the extension is imported
|
||||
|
||||
**Memory Reduction:**
|
||||
| Terrain Type | Vertex Reduction |
|
||||
|--------------|------------------|
|
||||
| Flat terrain | ~99% |
|
||||
| Maze-like | ~89% |
|
||||
| Pits terrain | ~80% |
|
||||
| Mixed terrain | ~79% |
|
||||
|
||||
This optimization is transparent - it produces visually identical terrains while dramatically reducing the mesh vertex count. The patches are applied before any terrain generation occurs, ensuring all height-field terrains benefit from the optimization.
|
||||
|
||||
### Terrain Data Flow
|
||||
|
||||
The terrain system uses **explicit boolean masks** instead of height-based classification:
|
||||
|
||||
```python
|
||||
# During terrain generation (hf_terrains_maze.py)
|
||||
terrain = TerrainData.create(width, height) # Lines 57-63
|
||||
|
||||
# Mark obstacles as invalid
|
||||
terrain.set_obstacle(x_start, x_end, y_start, y_end, wall_height) # Lines 65-73
|
||||
|
||||
# Apply padding and create masks
|
||||
terrain.apply_padding(PADDING.GOAL_PADDING) # Lines 92-97: 5 cells = 0.5m for goals
|
||||
spawn_mask = terrain.create_spawn_mask(PADDING.SPAWN_PADDING) # Lines 99-104: 6 cells = 0.6m
|
||||
|
||||
# Store on config for patches to pick up
|
||||
cfg.height_field_visual = heights # For Z-lookup (terrain height)
|
||||
cfg.height_field_valid_mask = valid_mask # For goal sampling (5 cell padding)
|
||||
cfg.height_field_spawn_mask = spawn_mask # For spawn sampling (6 cell padding)
|
||||
cfg.height_field_platform_mask = platform_mask # For curriculum learning
|
||||
```
|
||||
|
||||
These masks are set in the maze terrain function and picked up by the patches system (patches.py:45-66), which stores them on `TerrainImporter` for access via `env.scene.terrain._height_field_*`.
|
||||
|
||||
### Maze Terrain Types
|
||||
|
||||
Four terrain types are available via `HfMazeTerrainCfg`:
|
||||
|
||||
1. **Maze** (`non_maze_terrain=False, stairs=False`)
|
||||
- DFS-generated maze with configurable wall openings
|
||||
- Random obstacle shapes (pillars, bars, crosses, blocks)
|
||||
- Optional stairs integration (`add_stairs_to_maze=True`)
|
||||
|
||||
2. **Non-Maze/Random** (`non_maze_terrain=True`)
|
||||
- Random obstacle placement (~15-35% coverage based on difficulty)
|
||||
- Good for testing navigation without maze structure
|
||||
|
||||
3. **Stairs** (`stairs=True`)
|
||||
- 3x3 stair/platform structures with 4 cardinal stairways
|
||||
- Elevated platforms marked for curriculum learning
|
||||
- Tests robot climbing capabilities
|
||||
|
||||
4. **Pits** (`dynamic_obstacles=True`)
|
||||
- Pit rows with bridge crossings
|
||||
- Mix of pit (60%) and wall (40%) obstacles
|
||||
- Tests navigation over negative obstacles
|
||||
|
||||
### Safety Padding
|
||||
|
||||
Two padding levels ensure safe robot placement:
|
||||
|
||||
| Padding Type | Cells | Meters | Purpose |
|
||||
|--------------|-------|--------|---------|
|
||||
| `GOAL_PADDING` | 5 | 0.5m | Goal positions (robot just needs to reach) |
|
||||
| `SPAWN_PADDING` | 6 | 0.6m | Spawn positions (accounts for robot body) |
|
||||
|
||||
The larger spawn padding accounts for:
|
||||
- Robot body dimensions (~0.5m × 0.3m for quadrupeds)
|
||||
- Random yaw orientation (diagonal ~0.58m requires ~0.3m clearance)
|
||||
- Platform edge safety (prevent falling when spawning near stairs)
|
||||
- Controller startup behavior
|
||||
|
||||
These values are defined in `terrain_constants.py` as `PADDING.GOAL_PADDING` and `PADDING.SPAWN_PADDING`.
|
||||
|
||||
### Terrain Configuration
|
||||
|
||||
The main terrain configuration is defined in [maze_config.py](../source/isaaclab_nav_task/isaaclab_nav_task/navigation/terrains/maze_config.py):
|
||||
|
||||
```python
|
||||
MAZE_TERRAIN_CFG = TerrainGeneratorCfg(
|
||||
size=(30.0, 30.0), # 30m × 30m per terrain tile
|
||||
border_width=30.0, # Border around entire grid (not per-tile)
|
||||
num_rows=6, # 6 difficulty levels
|
||||
num_cols=30, # 30 terrain variations
|
||||
horizontal_scale=0.1, # 0.1m per height field cell (300×300 cells)
|
||||
vertical_scale=0.005, # Height conversion: value * 0.005 = meters
|
||||
curriculum=False, # Random terrain, not progressive
|
||||
difficulty_range=(0.5, 1.0), # Difficulty sampling range
|
||||
sub_terrains={
|
||||
"maze": HfMazeTerrainCfg(proportion=0.3, open_probability=0.9,
|
||||
random_wall_ratio=0.5, add_stairs_to_maze=True),
|
||||
"non_maze": HfMazeTerrainCfg(proportion=0.2, non_maze_terrain=True,
|
||||
random_wall_ratio=1.0),
|
||||
"stairs": HfMazeTerrainCfg(proportion=0.3, stairs=True,
|
||||
randomize_wall=False),
|
||||
"pits": HfMazeTerrainCfg(proportion=0.2, dynamic_obstacles=True,
|
||||
random_wall_ratio=1.0),
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
**Global Parameters:**
|
||||
- Total: **180 terrains** (6 rows × 30 cols), each 30m × 30m
|
||||
- Height field: **300×300 cells per terrain** (0.1m resolution)
|
||||
- `curriculum=False`: Random assignment across all 180 terrains
|
||||
|
||||
### Curriculum Learning
|
||||
|
||||
Terrains are organized in a grid with difficulty varying by row:
|
||||
- **Rows** (`terrain_levels`): Difficulty levels (0.0 to 1.0)
|
||||
- **Columns** (`terrain_types`): Different terrain types
|
||||
|
||||
```
|
||||
Difficulty
|
||||
1.0 | [Hard Maze] [Random Obs] [Tall Stairs] [Deep Pits] ...
|
||||
0.8 | [Med Maze] [Med Obs] [Med Stairs] [Med Pits] ...
|
||||
0.5 | [Easy Maze] [Few Obs] [Low Stairs] [Shallow] ...
|
||||
0.0 | [Flat] [Flat] [Flat] [Flat] ...
|
||||
└──────────────────────────────────────────────────────
|
||||
maze non_maze stairs pits
|
||||
```
|
||||
|
||||
## Goal and Spawn Sampling
|
||||
|
||||
The goal command generator ([RobotNavigationGoalCommand](../source/isaaclab_nav_task/isaaclab_nav_task/navigation/mdp/navigation/goal_commands.py)) handles sampling valid goal and spawn positions from maze terrains using pre-computed boolean masks.
|
||||
|
||||
### Architecture
|
||||
|
||||
```
|
||||
┌────────────────────────────────────────────────────────────────────┐
|
||||
│ RobotNavigationGoalCommand (goal_commands.py:329-766) │
|
||||
│ └─► _initialize_position_sampling() (once, lines 415-471) │
|
||||
│ └─► Creates PositionSampler with: │
|
||||
│ - heights: Z-lookup for terrain height │
|
||||
│ - valid_mask: Goal positions (5 cells = 0.5m padding) │
|
||||
│ - spawn_mask: Spawn positions (6 cells = 0.6m padding) │
|
||||
│ - platform_mask: Curriculum learning targets │
|
||||
│ - platform_repeat_count: 10 (default, line 74) │
|
||||
│ - border_width: 0.0 (sub-terrain default, line 457) │
|
||||
│ │
|
||||
│ └─► _resample_command(env_ids) (each reset, lines 499-553) │
|
||||
│ └─► sample(): Goal from valid_mask (line 519) │
|
||||
│ └─► sample_spawn(): Spawn from spawn_mask (line 521) │
|
||||
│ └─► Convert local → world coordinates (lines 524-547) │
|
||||
└────────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Key Features
|
||||
|
||||
- **Pre-computed masks**: `valid_mask` and `spawn_mask` generated during terrain creation
|
||||
- **Separate padding**: Goals (5 cells = 0.5m) vs spawns (6 cells = 0.6m) for robot body clearance
|
||||
- **Platform repetition**: Stair platforms repeated 10x in sampling for curriculum learning (goal_commands.py:74, 154)
|
||||
- **Efficient lookup**: Pre-built position tables enable O(1) random sampling
|
||||
- **Coordinate conversion**: Handles mesh border offset and centering transform
|
||||
- **Dynamic border calculation**: Border pixels computed from terrain config (goal_commands.py:107)
|
||||
|
||||
### Coordinate System
|
||||
|
||||
The terrain mesh uses a coordinate system with:
|
||||
- **Border pixels**: Dynamically computed as `int(border_width / horizontal_scale) + 1`
|
||||
- Sub-terrain configs default to `border_width=0.0` (from `HfTerrainBaseCfg`)
|
||||
- With default values: `border_pixels = int(0.0 / 0.1) + 1 = 1`
|
||||
- **Centering transform**: Mesh is centered at origin by `-terrain_size/2`
|
||||
|
||||
```python
|
||||
# Converting valid_mask index to local coordinates (goal_commands.py:283-284):
|
||||
local_x = (x_idx + border_pixels) * horizontal_scale - terrain_size/2
|
||||
local_y = (y_idx + border_pixels) * horizontal_scale - terrain_size/2
|
||||
|
||||
# Example: terrain_size=30m, horizontal_scale=0.1m, border_width=0.0
|
||||
# border_pixels = 1
|
||||
# valid_mask[0, 0] → local position: (0.1 - 15, 0.1 - 15) = (-14.9, -14.9)
|
||||
```
|
||||
|
||||
The `border_pixels` calculation is done in the `PositionSampler.__init__()` (goal_commands.py:107) to match the `@height_field_to_mesh` decorator behavior.
|
||||
|
||||
### Terrain Index Mapping
|
||||
|
||||
The terrain index formula depends on the generation mode:
|
||||
|
||||
| Mode | Formula | Description |
|
||||
|------|---------|-------------|
|
||||
| `curriculum=True` | `level + type * num_rows` | Column-major (iterate rows first) |
|
||||
| `curriculum=False` | `level * num_cols + type` | Row-major (iterate cols first) |
|
||||
|
||||
```python
|
||||
# In goal_commands.py:
|
||||
def _get_terrain_indices(self, env_ids):
|
||||
terrain = self.env.scene.terrain
|
||||
levels = terrain.terrain_levels[env_ids] # row
|
||||
types = terrain.terrain_types[env_ids] # col
|
||||
|
||||
if terrain_cfg.curriculum:
|
||||
return levels + types * num_rows # column-major
|
||||
else:
|
||||
return levels * num_cols + types # row-major
|
||||
```
|
||||
|
||||
### Position Sampling
|
||||
|
||||
**PositionSampler** (goal_commands.py:51-290) provides two sampling methods:
|
||||
|
||||
```python
|
||||
class PositionSampler:
|
||||
def sample(terrain_indices) -> (x, y, z):
|
||||
"""Sample GOAL positions from valid_mask.
|
||||
Uses platform repetition for curriculum learning.
|
||||
Default platform_repeat_count=10 (goal_commands.py:74)"""
|
||||
|
||||
def sample_spawn(terrain_indices) -> (x, y, z):
|
||||
"""Sample SPAWN positions from spawn_mask.
|
||||
Larger padding (6 cells vs 5 cells) for robot body with random orientation."""
|
||||
```
|
||||
|
||||
**Implementation Details:**
|
||||
- Pre-computes position tables during initialization (goal_commands.py:115-201)
|
||||
- Platform positions are repeated 10x in goal sampling for curriculum learning
|
||||
- Uses efficient O(1) random sampling from pre-built position tables
|
||||
- Falls back to `valid_mask` if `spawn_mask` is not provided (goal_commands.py:99)
|
||||
|
||||
**Sampling flow during episode reset:**
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────────┐
|
||||
│ _resample_command(env_ids) (goal_commands.py:499-553) │
|
||||
├─────────────────────────────────────────────────────────────────┤
|
||||
│ 1. Get terrain indices for each environment (line 516) │
|
||||
│ - terrain_levels[env_ids] → row (difficulty) │
|
||||
│ - terrain_types[env_ids] → col (terrain type) │
|
||||
│ - Apply curriculum/random index formula (lines 473-493) │
|
||||
│ │
|
||||
│ 2. Sample goal position from valid_mask (line 519) │
|
||||
│ - Random sample from pre-computed goal position table │
|
||||
│ - Platform positions repeated 10x for curriculum weighting │
|
||||
│ - Uses GOAL_PADDING = 5 cells = 0.5m │
|
||||
│ │
|
||||
│ 3. Sample spawn position from spawn_mask (line 521) │
|
||||
│ - Random sample from pre-computed spawn position table │
|
||||
│ - Larger padding (SPAWN_PADDING = 6 cells = 0.6m) │
|
||||
│ - Ensures robot body clearance with random orientation │
|
||||
│ │
|
||||
│ 4. Convert to world coordinates (lines 524-547) │
|
||||
│ - Add terrain_origins[level, type] offset │
|
||||
│ - Goal: Add random height offset (0.2-0.8m) for marker │
|
||||
│ - Spawn: Add spawn_offset = 0.05m (NOT 0.5m!) │
|
||||
│ (robot's default_root_state already has standing height) │
|
||||
│ │
|
||||
│ 5. Update environment origins (lines 540-542) │
|
||||
│ - env.scene.terrain.env_origins[env_ids] = spawn position │
|
||||
│ - Robot will be reset to this position │
|
||||
└─────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Robot Spawn Height
|
||||
|
||||
The spawn height offset accounts for the robot's standing height:
|
||||
|
||||
```python
|
||||
# In _resample_command() (goal_commands.py:537-542):
|
||||
spawn_offset = 0.05 # Small offset to prevent clipping into terrain
|
||||
|
||||
# Note: robot's default_root_state already includes standing height (~0.5m)
|
||||
terrain.env_origins[env_ids, 0] = terrain_origins[:, 0] + spawn_x
|
||||
terrain.env_origins[env_ids, 1] = terrain_origins[:, 1] + spawn_y
|
||||
terrain.env_origins[env_ids, 2] = spawn_z + spawn_offset # Just 5cm above terrain
|
||||
```
|
||||
|
||||
**Key Change:** The implementation uses only a small 5cm offset because the robot's `default_root_state` configuration already includes the proper standing height (~0.5m). This ensures the robot spawns at the correct height without double-counting the base height.
|
||||
@@ -0,0 +1,65 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""
|
||||
Isaac Lab Navigation Tasks extension.
|
||||
|
||||
This extension provides navigation tasks for robot learning with visual perception.
|
||||
|
||||
On import, this module:
|
||||
1. Applies monkey-patches to Isaac Lab terrain system for height field storage
|
||||
2. Registers maze terrain types
|
||||
3. Registers navigation task environments
|
||||
"""
|
||||
|
||||
import os
|
||||
import toml
|
||||
|
||||
# Conveniences to other module directories via relative paths
|
||||
ISAACLAB_NAV_TASKS_EXT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../"))
|
||||
"""Path to the extension source directory."""
|
||||
|
||||
ISAACLAB_NAV_TASKS_METADATA = toml.load(os.path.join(ISAACLAB_NAV_TASKS_EXT_DIR, "config", "extension.toml"))
|
||||
"""Extension metadata dictionary parsed from the extension.toml file."""
|
||||
|
||||
# Configure the module-level variables
|
||||
__version__ = ISAACLAB_NAV_TASKS_METADATA["package"]["version"]
|
||||
|
||||
##
|
||||
# Apply patches to Isaac Lab terrain system.
|
||||
# This must be done before any terrain generation happens.
|
||||
# IMPORTANT: Import patches directly from the module file, NOT through
|
||||
# terrains/__init__.py, to avoid triggering terrain imports before patching.
|
||||
##
|
||||
|
||||
# Direct import from patches module to avoid loading terrains/__init__.py
|
||||
import importlib.util
|
||||
import os as _os
|
||||
_patches_path = _os.path.join(_os.path.dirname(__file__), "terrains", "patches.py")
|
||||
_spec = importlib.util.spec_from_file_location("patches", _patches_path)
|
||||
_patches_module = importlib.util.module_from_spec(_spec)
|
||||
_spec.loader.exec_module(_patches_module)
|
||||
_patches_module.apply_terrain_patches()
|
||||
del _patches_path, _spec, _patches_module
|
||||
|
||||
##
|
||||
# Export terrain types for convenient access.
|
||||
##
|
||||
|
||||
from .terrains import (
|
||||
HfMazeTerrainCfg,
|
||||
MAZE_TERRAIN_CFG,
|
||||
)
|
||||
|
||||
##
|
||||
# Register Gym environments.
|
||||
##
|
||||
|
||||
from isaaclab_tasks.utils import import_packages
|
||||
|
||||
# The blacklist is used to prevent importing configs from sub-packages
|
||||
_BLACKLIST_PKGS = ["utils", "terrains"]
|
||||
# Import all configs in this package
|
||||
import_packages(__name__, _BLACKLIST_PKGS)
|
||||
@@ -0,0 +1,12 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Navigation task environments for Isaac Lab."""
|
||||
|
||||
from .navigation_env import NavigationEnv
|
||||
from .navigation_env_cfg import *
|
||||
|
||||
# Import robot-specific configurations
|
||||
from .config import *
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Custom robot configurations and assets for navigation tasks."""
|
||||
|
||||
import os
|
||||
|
||||
# Path to the local data directory containing robots and policies
|
||||
ISAACLAB_NAV_TASKS_ASSETS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
|
||||
"""Path to the navigation tasks assets data directory."""
|
||||
|
||||
from .b2w import *
|
||||
from .aow_d import *
|
||||
from .go2 import *
|
||||
|
||||
__all__ = ["ISAACLAB_NAV_TASKS_ASSETS_DIR", "B2W_CFG", "ANYMAL_D_ON_WHEELS_CFG", "GO2_CFG"]
|
||||
@@ -0,0 +1,70 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Configuration for the ANYbotics robots with wheels.
|
||||
|
||||
The following configuration parameters are available:
|
||||
|
||||
* :obj:`ANYMAL_D_ON_WHEELS_CFG`: The ANYmal-D on wheels with ImplicitActuatorCfg.
|
||||
|
||||
"""
|
||||
|
||||
from . import ISAACLAB_NAV_TASKS_ASSETS_DIR
|
||||
|
||||
import isaaclab.sim as sim_utils
|
||||
from isaaclab.actuators import ImplicitActuatorCfg
|
||||
from isaaclab.assets.articulation import ArticulationCfg
|
||||
|
||||
__all__ = ["ANYMAL_D_ON_WHEELS_CFG"]
|
||||
|
||||
|
||||
ANYMAL_D_ON_WHEELS_CFG = ArticulationCfg(
|
||||
spawn=sim_utils.UsdFileCfg(
|
||||
usd_path=f"{ISAACLAB_NAV_TASKS_ASSETS_DIR}/Robots/AoW-D/aow_d.usd",
|
||||
activate_contact_sensors=True,
|
||||
rigid_props=sim_utils.RigidBodyPropertiesCfg(
|
||||
disable_gravity=False,
|
||||
retain_accelerations=False,
|
||||
linear_damping=0.0,
|
||||
angular_damping=0.0,
|
||||
max_linear_velocity=None,
|
||||
max_angular_velocity=None,
|
||||
max_depenetration_velocity=1.0,
|
||||
enable_gyroscopic_forces=True,
|
||||
),
|
||||
articulation_props=sim_utils.ArticulationRootPropertiesCfg(
|
||||
enabled_self_collisions=True, solver_position_iteration_count=4, solver_velocity_iteration_count=0
|
||||
),
|
||||
),
|
||||
init_state=ArticulationCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.65),
|
||||
joint_pos={
|
||||
".*HAA": 0.0, # all HAA
|
||||
".*F_HFE": 0.4, # both front HFE
|
||||
".*H_HFE": -0.4, # both hind HFE
|
||||
".*F_KFE": -0.8, # both front KFE
|
||||
".*H_KFE": 0.8, # both hind KFE
|
||||
".*WHEEL": 0.0, # all WHEEL
|
||||
},
|
||||
),
|
||||
actuators={
|
||||
"legs": ImplicitActuatorCfg(
|
||||
joint_names_expr=[".*HAA", ".*HFE", ".*KFE"],
|
||||
effort_limit_sim=80.0, # From URDF: HAA, HFE, KFE joints all have 80.0 Nm effort limit
|
||||
velocity_limit_sim=8.5, # From URDF: HAA, HFE, KFE joints all have 8.5 rad/s velocity limit
|
||||
stiffness={".*": 100.0}, # Typical stiffness value for legged robots
|
||||
damping={".*": 3.5}, # Typical damping value for legged robots
|
||||
),
|
||||
"wheels": ImplicitActuatorCfg(
|
||||
joint_names_expr=[".*WHEEL"],
|
||||
effort_limit_sim=28.0, # From URDF: wheel joints have 28.0 Nm effort limit
|
||||
velocity_limit_sim=200.0, # From URDF: wheel joints have 200.0 rad/s velocity limit
|
||||
stiffness={".*": 0.0}, # Wheels typically have no stiffness (direct velocity control)
|
||||
damping={".*": 5.0}, # From original configuration
|
||||
),
|
||||
},
|
||||
soft_joint_pos_limit_factor=0.95,
|
||||
)
|
||||
"""Configuration of wheelified ANYmal-D using ImplicitActuatorCfg."""
|
||||
@@ -0,0 +1,80 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Configuration for the B2W (Unitree B2W) robot.
|
||||
|
||||
The following configuration parameters are available:
|
||||
|
||||
* :obj:`B2W_CFG`: The B2W robot with wheels.
|
||||
|
||||
Reference:
|
||||
The B2W is a bipedal wheeled robot.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import isaaclab.sim as sim_utils
|
||||
from isaaclab.actuators import ImplicitActuatorCfg
|
||||
from isaaclab.assets.articulation import ArticulationCfg
|
||||
|
||||
# Local assets directory for this extension
|
||||
_ASSETS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
|
||||
|
||||
__all__ = ["B2W_CFG"]
|
||||
|
||||
|
||||
B2W_CFG = ArticulationCfg(
|
||||
spawn=sim_utils.UsdFileCfg(
|
||||
usd_path=f"{_ASSETS_DIR}/Robots/B2W/b2w_rsl.usd",
|
||||
activate_contact_sensors=True,
|
||||
rigid_props=sim_utils.RigidBodyPropertiesCfg(
|
||||
disable_gravity=False,
|
||||
retain_accelerations=False,
|
||||
linear_damping=0.0,
|
||||
angular_damping=0.0,
|
||||
max_linear_velocity=None,
|
||||
max_angular_velocity=None,
|
||||
max_depenetration_velocity=1.0,
|
||||
enable_gyroscopic_forces=True,
|
||||
),
|
||||
articulation_props=sim_utils.ArticulationRootPropertiesCfg(
|
||||
enabled_self_collisions=True, solver_position_iteration_count=4, solver_velocity_iteration_count=0
|
||||
),
|
||||
),
|
||||
init_state=ArticulationCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.75),
|
||||
joint_pos={
|
||||
".*hip_joint": 0.0,
|
||||
".*thigh_joint": 0.4,
|
||||
".*foot_joint": 0.0,
|
||||
".*calf_joint": -1.3,
|
||||
},
|
||||
),
|
||||
actuators={
|
||||
"legs": ImplicitActuatorCfg(
|
||||
joint_names_expr=[".*hip_joint", ".*thigh_joint"],
|
||||
effort_limit_sim=200.0,
|
||||
velocity_limit_sim=23.0,
|
||||
stiffness={".*": 100.0},
|
||||
damping={".*": 3.5},
|
||||
),
|
||||
"legs_calf": ImplicitActuatorCfg(
|
||||
joint_names_expr=[".*calf_joint"],
|
||||
effort_limit_sim=320.0,
|
||||
velocity_limit_sim=14.0,
|
||||
stiffness={".*": 100.0},
|
||||
damping={".*": 3.5},
|
||||
),
|
||||
"wheels": ImplicitActuatorCfg(
|
||||
joint_names_expr=[".*foot_joint"],
|
||||
effort_limit_sim=20.0,
|
||||
velocity_limit_sim=50.0,
|
||||
stiffness={".*": 0.0},
|
||||
damping={".*": 3.0},
|
||||
),
|
||||
},
|
||||
soft_joint_pos_limit_factor=0.95,
|
||||
)
|
||||
"""Configuration of B2W robot using ImplicitActuatorCfg."""
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
|
After Width: | Height: | Size: 111 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 449 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 175 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 157 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 223 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 225 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 201 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 140 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 200 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 302 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 114 KiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 97 KiB |
BIN
Binary file not shown.
BIN
Binary file not shown.
+17
@@ -0,0 +1,17 @@
|
||||
asset_path: /home/passutte/Developer/leggedrobotics/rsl_assets/Alma-D-Wheels-No-Arm/urdf/alma_manual_wheels.urdf
|
||||
usd_dir: /home/passutte/Developer/leggedrobotics/orbit/source/extensions/omni.isaac.orbit_assets/data/Robots/RSL-ETHZ/AoW-D
|
||||
usd_file_name: aow_d.usd
|
||||
force_usd_conversion: true
|
||||
make_instanceable: true
|
||||
import_inertia_tensor: true
|
||||
fix_base: false
|
||||
merge_fixed_joints: true
|
||||
self_collision: false
|
||||
default_drive_type: none
|
||||
default_drive_stiffness: 0.0
|
||||
default_drive_damping: 0.0
|
||||
link_density: 0.0
|
||||
convex_decompose_mesh: false
|
||||
##
|
||||
# Generated by UrdfConverter on 2024-06-08 at 15:27:00.
|
||||
##
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
+5
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
print("hi")
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
+18
@@ -0,0 +1,18 @@
|
||||
asset_path: /home/fanyang1/Desktop/catkin_ws/src/b2w_description/urdf/b2w_description_rsl.urdf
|
||||
usd_dir: /home/fanyang1/IsaacLab-Internal/source/standalone/tools
|
||||
usd_file_name: b2w_rsl.usd
|
||||
force_usd_conversion: true
|
||||
make_instanceable: true
|
||||
import_inertia_tensor: true
|
||||
fix_base: false
|
||||
merge_fixed_joints: true
|
||||
self_collision: false
|
||||
default_drive_type: none
|
||||
override_joint_dynamics: false
|
||||
default_drive_stiffness: 0.0
|
||||
default_drive_damping: 0.0
|
||||
link_density: 0.0
|
||||
convex_decompose_mesh: false
|
||||
##
|
||||
# Generated by UrdfConverter on 2024-11-04 at 15:36:52.
|
||||
##
|
||||
BIN
Binary file not shown.
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Configuration for the Unitree Go2 robot (reuses IsaacLab built-in USD).
|
||||
|
||||
* :obj:`GO2_CFG`: Unitree Go2 quadruped (12 leg joints, no wheels).
|
||||
"""
|
||||
|
||||
from isaaclab_assets.robots.unitree import UNITREE_GO2_CFG
|
||||
|
||||
__all__ = ["GO2_CFG"]
|
||||
|
||||
# Reuse the upstream IsaacLab Go2 articulation cfg verbatim. The locomotion
|
||||
# checkpoint we converted (policy_go2_jit.pt) was trained on this exact USD
|
||||
# (Isaac-Velocity-Flat-Unitree-Go2-v0), so joint order and joint count match.
|
||||
GO2_CFG = UNITREE_GO2_CFG.copy()
|
||||
"""Configuration of Unitree Go2 robot for navigation tasks."""
|
||||
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Robot-specific configuration modules for navigation tasks."""
|
||||
|
||||
from .b2w import *
|
||||
from .aow_d import *
|
||||
from .go2 import *
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import gymnasium as gym
|
||||
|
||||
from . import agents, navigation_env_cfg
|
||||
|
||||
##
|
||||
# Register Gym environments.
|
||||
##
|
||||
|
||||
##############################################################################################################
|
||||
# MDPO
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-MDPO-AoW-D-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.AowDNavigationEnvCfg,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.AowDNavMDPORunnerCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-MDPO-AoW-D-Play-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.AowDNavigationEnvCfg_PLAY,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.AowDNavMDPORunnerCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-MDPO-AoW-D-Dev-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.AowDNavigationEnvCfg_DEV,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.AowDNavMDPORunnerDevCfg,
|
||||
},
|
||||
)
|
||||
|
||||
######################################################################################
|
||||
# PPO
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-AoW-D-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.AowDNavigationEnvCfg,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.AowDNavPPORunnerCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-AoW-D-Play-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.AowDNavigationEnvCfg_PLAY,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.AowDNavPPORunnerCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-AoW-D-Dev-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.AowDNavigationEnvCfg_DEV,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.AowDNavPPORunnerDevCfg,
|
||||
},
|
||||
)
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Agent configurations for AoW-D navigation tasks."""
|
||||
|
||||
from . import rsl_rl_cfg
|
||||
+119
@@ -0,0 +1,119 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""RSL-RL agent configurations for AoW-D navigation tasks."""
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
|
||||
from isaaclab_nav_task.navigation.config.rl_cfg import (
|
||||
RslRlOnPolicyRunnerCfg,
|
||||
RslRlPpoActorCriticCfg,
|
||||
RslRlPpoAlgorithmCfg,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class AowDNavMDPORunnerCfg(RslRlOnPolicyRunnerCfg):
|
||||
"""MDPO runner configuration for AoW-D navigation."""
|
||||
|
||||
num_steps_per_env = 16
|
||||
max_iterations = 15000
|
||||
save_interval = 500
|
||||
logger = "wandb"
|
||||
seed = 50
|
||||
wandb_project = "isaaclab_nav_aow_d"
|
||||
experiment_name = "aow_d_navigation_mdpo"
|
||||
empirical_normalization = False
|
||||
reward_shifting_value = 0.05
|
||||
policy = RslRlPpoActorCriticCfg(
|
||||
class_name="ActorCriticSRU",
|
||||
init_noise_std=1.0,
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
activation="elu",
|
||||
rnn_hidden_size=512,
|
||||
rnn_type="lstm_sru",
|
||||
rnn_num_layers=1,
|
||||
dropout=0.2,
|
||||
num_cameras=1,
|
||||
image_input_dims=(64, 5, 8),
|
||||
)
|
||||
algorithm = RslRlPpoAlgorithmCfg(
|
||||
class_name="MDPO",
|
||||
value_loss_coef=0.02,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.2,
|
||||
value_clip_param=0.2,
|
||||
entropy_coef=0.00375,
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=1.0e-3,
|
||||
schedule="adaptive",
|
||||
gamma=0.995,
|
||||
lam=0.95,
|
||||
desired_kl=0.01,
|
||||
max_grad_norm=1.0,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class AowDNavMDPORunnerDevCfg(AowDNavMDPORunnerCfg):
|
||||
"""Development configuration for MDPO with reduced iterations."""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.max_iterations = 300
|
||||
self.experiment_name = "aow_d_navigation_mdpo_dev"
|
||||
self.logger = "tensorboard"
|
||||
|
||||
|
||||
@configclass
|
||||
class AowDNavPPORunnerCfg(RslRlOnPolicyRunnerCfg):
|
||||
"""PPO runner configuration for AoW-D navigation."""
|
||||
|
||||
num_steps_per_env = 24
|
||||
max_iterations = 15000
|
||||
save_interval = 500
|
||||
logger = "wandb"
|
||||
wandb_project = "isaaclab_nav_aow_d"
|
||||
experiment_name = "aow_d_navigation_ppo"
|
||||
empirical_normalization = False
|
||||
policy = RslRlPpoActorCriticCfg(
|
||||
class_name="ActorCriticSRU",
|
||||
init_noise_std=1.0,
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
activation="elu",
|
||||
rnn_hidden_size=256,
|
||||
rnn_type="lstm_sru",
|
||||
num_cameras=1,
|
||||
image_input_dims=(64, 5, 8),
|
||||
)
|
||||
algorithm = RslRlPpoAlgorithmCfg(
|
||||
class_name="PPO",
|
||||
value_loss_coef=0.1,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.2,
|
||||
entropy_coef=0.005,
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=1.0e-3,
|
||||
schedule="adaptive",
|
||||
gamma=0.995,
|
||||
lam=0.95,
|
||||
desired_kl=0.01,
|
||||
max_grad_norm=1.0,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class AowDNavPPORunnerDevCfg(AowDNavPPORunnerCfg):
|
||||
"""Development configuration for PPO with reduced iterations."""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.max_iterations = 300
|
||||
self.experiment_name = "aow_d_navigation_ppo_dev"
|
||||
self.logger = "tensorboard"
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""AoW-D (Anymal on Wheels) specific configuration for navigation environment."""
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
|
||||
from isaaclab_nav_task.navigation.navigation_env_cfg import NavigationEnvCfg
|
||||
import isaaclab_nav_task.navigation.mdp as mdp
|
||||
|
||||
from isaaclab_nav_task.navigation.assets import ANYMAL_D_ON_WHEELS_CFG # isort: skip
|
||||
|
||||
|
||||
@configclass
|
||||
class AowDNavigationEnvCfg(NavigationEnvCfg):
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
|
||||
from isaaclab_nav_task.navigation.mdp.observations import initialize_depth_noise_generator
|
||||
from isaaclab_nav_task.navigation.mdp.depth_utils.camera_config import get_camera_config
|
||||
|
||||
initialize_depth_noise_generator(robot_name="aow_d", use_jit_precompiled=False)
|
||||
|
||||
camera_config = get_camera_config("aow_d")
|
||||
CAMERA_RESOLUTION = camera_config.resolution
|
||||
|
||||
self.scene.robot = ANYMAL_D_ON_WHEELS_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
|
||||
|
||||
self.scene.terrain.max_init_terrain_level = 10
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [0.5, 1.0]
|
||||
self.scene.terrain.terrain_generator.curriculum = False
|
||||
|
||||
@configclass
|
||||
class AowDNavigationEnvCfg_DEV(AowDNavigationEnvCfg):
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.scene.terrain.terrain_generator.num_rows = 1
|
||||
self.scene.terrain.terrain_generator.num_cols = 30
|
||||
self.scene.terrain.max_init_terrain_level = 10
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [0.4, 0.8]
|
||||
self.scene.terrain.terrain_generator.curriculum = True
|
||||
self.scene.num_envs = 1
|
||||
|
||||
@configclass
|
||||
class AowDNavigationEnvCfg_PLAY(AowDNavigationEnvCfg):
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
|
||||
self.scene.num_envs = 20
|
||||
self.scene.env_spacing = 2.5
|
||||
self.scene.terrain.max_init_terrain_level = None
|
||||
if self.scene.terrain.terrain_generator is not None:
|
||||
self.scene.terrain.terrain_generator.num_rows = 2
|
||||
self.scene.terrain.terrain_generator.num_cols = 2
|
||||
|
||||
self.observations.policy.enable_corruption = False
|
||||
self.events.base_external_force_torque = None
|
||||
self.events.push_robot = None
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import gymnasium as gym
|
||||
|
||||
from . import agents, navigation_env_cfg
|
||||
|
||||
##
|
||||
# Register Gym environments.
|
||||
##
|
||||
|
||||
##############################################################################################################
|
||||
# MDPO
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-MDPO-B2W-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.B2WNavigationEnvCfg,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.B2WNavMDPORunnerCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-MDPO-B2W-Play-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.B2WNavigationEnvCfg_PLAY,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.B2WNavMDPORunnerCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-MDPO-B2W-Dev-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.B2WNavigationEnvCfg_DEV,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.B2WNavMDPORunnerDevCfg,
|
||||
},
|
||||
)
|
||||
|
||||
######################################################################################
|
||||
# PPO
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-B2W-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.B2WNavigationEnvCfg,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.B2WNavPPORunnerCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-B2W-Play-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.B2WNavigationEnvCfg_PLAY,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.B2WNavPPORunnerCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-B2W-Dev-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.B2WNavigationEnvCfg_DEV,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.B2WNavPPORunnerDevCfg,
|
||||
},
|
||||
)
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Agent configurations for B2W navigation tasks."""
|
||||
|
||||
from . import rsl_rl_cfg
|
||||
+126
@@ -0,0 +1,126 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""RSL-RL agent configurations for B2W navigation tasks."""
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
|
||||
from isaaclab_nav_task.navigation.config.rl_cfg import (
|
||||
RslRlOnPolicyRunnerCfg,
|
||||
RslRlPpoActorCriticCfg,
|
||||
RslRlPpoAlgorithmCfg,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class B2WNavMDPORunnerCfg(RslRlOnPolicyRunnerCfg):
|
||||
"""MDPO runner configuration for B2W navigation."""
|
||||
|
||||
num_steps_per_env = 16
|
||||
max_iterations = 15000
|
||||
save_interval = 500
|
||||
logger = "wandb"
|
||||
seed = 60
|
||||
wandb_project = "isaaclab_nav_b2w"
|
||||
experiment_name = "b2w_navigation_mdpo"
|
||||
empirical_normalization = False
|
||||
reward_shifting_value = 0.05
|
||||
policy = RslRlPpoActorCriticCfg(
|
||||
class_name="ActorCriticSRU",
|
||||
init_noise_std=1.0,
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
activation="elu",
|
||||
rnn_hidden_size=512,
|
||||
rnn_type="lstm_sru",
|
||||
rnn_num_layers=1,
|
||||
dropout=0.2,
|
||||
num_cameras=1,
|
||||
image_input_dims=(64, 5, 8), # depth image: 64 channels * 5 * 8 = 2560
|
||||
height_input_dims=(64, 7, 7), # encoded height_scan_critic: 64*7*7 = 3136
|
||||
)
|
||||
algorithm = RslRlPpoAlgorithmCfg(
|
||||
class_name="MDPO",
|
||||
value_loss_coef=0.02,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.2,
|
||||
value_clip_param=0.2,
|
||||
entropy_coef=0.00375,
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=1.0e-3,
|
||||
schedule="exponential",
|
||||
gamma=0.999,
|
||||
lam=0.95,
|
||||
desired_kl=0.01,
|
||||
max_grad_norm=1.0,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class B2WNavMDPORunnerDevCfg(B2WNavMDPORunnerCfg):
|
||||
"""Development configuration for MDPO with reduced iterations."""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.max_iterations = 300
|
||||
self.experiment_name = "b2w_navigation_mdpo_dev"
|
||||
self.logger = "tensorboard"
|
||||
|
||||
|
||||
@configclass
|
||||
class B2WNavPPORunnerCfg(RslRlOnPolicyRunnerCfg):
|
||||
"""PPO runner configuration for B2W navigation."""
|
||||
|
||||
num_steps_per_env = 16
|
||||
max_iterations = 15000
|
||||
save_interval = 500
|
||||
logger = "wandb"
|
||||
seed = 60
|
||||
wandb_project = "isaaclab_nav_b2w"
|
||||
experiment_name = "b2w_navigation_ppo"
|
||||
empirical_normalization = False
|
||||
reward_shifting_value = 0.05
|
||||
policy = RslRlPpoActorCriticCfg(
|
||||
class_name="ActorCriticSRU",
|
||||
init_noise_std=1.0,
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
activation="elu",
|
||||
rnn_hidden_size=512,
|
||||
rnn_type="lstm_sru",
|
||||
rnn_num_layers=1,
|
||||
dropout=0.2,
|
||||
num_cameras=1,
|
||||
image_input_dims=(64, 5, 8), # depth image: 64 channels * 5 * 8 = 2560
|
||||
height_input_dims=(64, 7, 7), # encoded height_scan_critic: 64*7*7 = 3136
|
||||
)
|
||||
algorithm = RslRlPpoAlgorithmCfg(
|
||||
class_name="PPO",
|
||||
value_loss_coef=0.02,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.2,
|
||||
value_clip_param=0.2,
|
||||
entropy_coef=0.00375,
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=1.0e-3,
|
||||
schedule="adaptive",
|
||||
gamma=0.995,
|
||||
lam=0.95,
|
||||
desired_kl=0.01,
|
||||
max_grad_norm=1.0,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class B2WNavPPORunnerDevCfg(B2WNavPPORunnerCfg):
|
||||
"""Development configuration for PPO with reduced iterations."""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.max_iterations = 300
|
||||
self.experiment_name = "b2w_navigation_ppo_dev"
|
||||
self.logger = "tensorboard"
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""B2W specific configuration for navigation environment."""
|
||||
|
||||
import os
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
|
||||
from isaaclab_nav_task.navigation.navigation_env_cfg import NavigationEnvCfg
|
||||
import isaaclab_nav_task.navigation.mdp as mdp
|
||||
|
||||
from isaaclab_nav_task.navigation.assets import B2W_CFG, ISAACLAB_NAV_TASKS_ASSETS_DIR # isort: skip
|
||||
|
||||
|
||||
LEG_JOINT_NAMES = [".*hip_joint", ".*thigh_joint", ".*calf_joint"]
|
||||
WHEEL_JOINT_NAMES = [".*foot_joint"]
|
||||
|
||||
@configclass
|
||||
class B2WNavigationEnvCfg(NavigationEnvCfg):
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
|
||||
from isaaclab_nav_task.navigation.mdp.observations import initialize_depth_noise_generator
|
||||
from isaaclab_nav_task.navigation.mdp.depth_utils.camera_config import get_camera_config
|
||||
|
||||
initialize_depth_noise_generator(robot_name="b2w", use_jit_precompiled=False)
|
||||
|
||||
camera_config = get_camera_config("b2w")
|
||||
CAMERA_RESOLUTION = camera_config.resolution
|
||||
|
||||
self.scene.robot = B2W_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
|
||||
|
||||
self.scene.raycast_camera.prim_path = "{ENV_REGEX_NS}/Robot/base_link"
|
||||
self.scene.raycast_camera.offset.pos = (0.387, 0.0, 0.28)
|
||||
self.scene.height_scanner_critic.prim_path = "{ENV_REGEX_NS}/Robot/base_link"
|
||||
|
||||
self.terminations.base_contact.params = {"sensor_cfg": SceneEntityCfg("contact_forces", body_names=["base_link", ".*hip", ".*thigh"]), "threshold": 1.0}
|
||||
|
||||
self.actions.velocity_command.low_level_position_action = mdp.JointPositionActionCfg(asset_name="robot", joint_names=[".*hip_joint", ".*thigh_joint", ".*calf_joint"], scale=0.5, use_default_offset=True)
|
||||
self.actions.velocity_command.low_level_velocity_action = mdp.JointVelocityActionCfg(asset_name="robot", joint_names=[".*foot_joint"], scale=5.0, use_default_offset=True)
|
||||
self.actions.velocity_command.low_level_policy_file = os.path.join(ISAACLAB_NAV_TASKS_ASSETS_DIR, "Policies", "locomotion", "b2w", "policy_b2w_new_2.pt")
|
||||
|
||||
self.rewards.joint_acc_l2_joint.params = {"asset_cfg": SceneEntityCfg("robot", joint_names=LEG_JOINT_NAMES+WHEEL_JOINT_NAMES)}
|
||||
|
||||
self.terminations.base_contact.params = {"sensor_cfg": SceneEntityCfg("contact_forces", body_names=["base_link", ".*hip", ".*thigh"]), "threshold": 1.0}
|
||||
|
||||
self.events.randomize_low_pass_filter_alpha.params = {
|
||||
"alpha_range": (0.1, 0.6),
|
||||
"action_term": "velocity_command",
|
||||
"per_dimension": True,
|
||||
"alpha_range_vx": (0.1, 0.6),
|
||||
"alpha_range_vy": (0.1, 0.6),
|
||||
"alpha_range_omega": (0.1, 0.6),
|
||||
}
|
||||
|
||||
self.scene.terrain.max_init_terrain_level = 10
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [0.5, 1.0]
|
||||
self.scene.terrain.terrain_generator.curriculum = False
|
||||
|
||||
@configclass
|
||||
class B2WNavigationEnvCfg_DEV(B2WNavigationEnvCfg):
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.scene.terrain.terrain_generator.num_rows = 2
|
||||
self.scene.terrain.terrain_generator.num_cols = 30
|
||||
self.scene.terrain.max_init_terrain_level = 10
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [0.5, 1.0]
|
||||
self.scene.terrain.terrain_generator.curriculum = False
|
||||
|
||||
@configclass
|
||||
class B2WNavigationEnvCfg_PLAY(B2WNavigationEnvCfg):
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
|
||||
self.scene.num_envs = 20
|
||||
self.scene.env_spacing = 2.5
|
||||
self.scene.terrain.max_init_terrain_level = None
|
||||
if self.scene.terrain.terrain_generator is not None:
|
||||
self.scene.terrain.terrain_generator.num_rows = 2
|
||||
self.scene.terrain.terrain_generator.num_cols = 2
|
||||
|
||||
# ---- PLAY_DIFFICULTY="lo,hi": difficulty band ----
|
||||
lo, hi = 0.5, 1.0
|
||||
_play_diff = os.environ.get("PLAY_DIFFICULTY", "").strip()
|
||||
if _play_diff:
|
||||
_parts = _play_diff.replace(" ", "").split(",")
|
||||
lo, hi = float(_parts[0]), float(_parts[1])
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [lo, hi]
|
||||
self.scene.terrain.terrain_generator.curriculum = False
|
||||
print(f"[B2W PLAY] terrain difficulty_range = [{lo}, {hi}] (curriculum off)")
|
||||
|
||||
sub_terrains = self.scene.terrain.terrain_generator.sub_terrains
|
||||
|
||||
# PLAY_MAZE_ONLY=1: 100% maze sub-terrain
|
||||
_maze_only = os.environ.get("PLAY_MAZE_ONLY", "").strip().lower() in ("1", "true", "yes")
|
||||
if _maze_only and "maze" in sub_terrains:
|
||||
for name in list(sub_terrains.keys()):
|
||||
sub_terrains[name].proportion = 1.0 if name == "maze" else 0.0
|
||||
print(f"[B2W PLAY] PLAY_MAZE_ONLY=1 -> 100% maze sub-terrain")
|
||||
|
||||
# PLAY_SUBTERRAIN_MIX="maze=1,non_maze=0,pits=0"
|
||||
_mix = os.environ.get("PLAY_SUBTERRAIN_MIX", "").strip()
|
||||
if _mix:
|
||||
overrides = {}
|
||||
for part in _mix.split(","):
|
||||
if "=" not in part:
|
||||
continue
|
||||
k, v = part.split("=", 1)
|
||||
overrides[k.strip()] = float(v.strip())
|
||||
for name, prop in overrides.items():
|
||||
if name in sub_terrains:
|
||||
sub_terrains[name].proportion = prop
|
||||
final = {n: sub_terrains[n].proportion for n in sub_terrains}
|
||||
print(f"[B2W PLAY] PLAY_SUBTERRAIN_MIX -> {final}")
|
||||
|
||||
# PLAY_CELL_SIZE: meters per maze cell
|
||||
_play_cell = os.environ.get("PLAY_CELL_SIZE", "").strip()
|
||||
if _play_cell:
|
||||
new_cell = float(_play_cell)
|
||||
for cfg_sub in sub_terrains.values():
|
||||
cfg_sub.cell_size = new_cell
|
||||
tile_m = 15 * new_cell
|
||||
self.scene.terrain.terrain_generator.size = (tile_m, tile_m)
|
||||
print(f"[B2W PLAY] PLAY_CELL_SIZE={new_cell}m -> tile={tile_m}m")
|
||||
|
||||
# PLAY_CLASSIC_MAZE=1: paper-style uniform walls
|
||||
_classic = os.environ.get("PLAY_CLASSIC_MAZE", "").strip().lower() in ("1", "true", "yes")
|
||||
if _classic and "maze" in sub_terrains:
|
||||
sub_terrains["maze"].randomize_wall = False
|
||||
sub_terrains["maze"].random_wall_ratio = 0.0
|
||||
print(f"[B2W PLAY] PLAY_CLASSIC_MAZE=1 -> classic uniform walls in maze")
|
||||
|
||||
self.observations.policy.enable_corruption = False
|
||||
self.events.base_external_force_torque = None
|
||||
self.events.push_robot = None
|
||||
@@ -0,0 +1,58 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import gymnasium as gym
|
||||
|
||||
from . import agents, navigation_env_cfg
|
||||
|
||||
##
|
||||
# Register Gym environments.
|
||||
##
|
||||
|
||||
##############################################################################################################
|
||||
# PPO
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-Go2-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.Go2NavigationEnvCfg,
|
||||
# MixedCfg is the cold-start tuned variant of the Dev hyperparams;
|
||||
# the original B2W-FT ``Go2NavPPORunnerCfg`` collapses exploration on a
|
||||
# from-scratch run (init_std=0.5, entropy=0.001, clip=0.1).
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.Go2NavPPORunnerMixedCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-Go2-Play-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.Go2NavigationEnvCfg_PLAY,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.Go2NavPPORunnerCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-Go2-Dev-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.Go2NavigationEnvCfg_DEV,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.Go2NavPPORunnerDevCfg,
|
||||
},
|
||||
)
|
||||
|
||||
gym.register(
|
||||
id="Isaac-Nav-PPO-Go2-PureMaze-v0",
|
||||
entry_point="isaaclab_nav_task.navigation:NavigationEnv",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": navigation_env_cfg.Go2NavigationEnvCfg_PureMaze,
|
||||
"rsl_rl_cfg_entry_point": agents.rsl_rl_cfg.Go2NavPPORunnerPureMazeCfg,
|
||||
},
|
||||
)
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Agent configurations for Go2 navigation tasks."""
|
||||
|
||||
from . import rsl_rl_cfg
|
||||
+186
@@ -0,0 +1,186 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""RSL-RL agent configurations for Unitree Go2 navigation tasks.
|
||||
|
||||
The policy block is intentionally identical to ``B2WNavPPORunnerCfg`` so that a
|
||||
B2W checkpoint can be warm-started via ``--resume`` / ``--load_run`` with
|
||||
``strict=True`` state-dict loading.
|
||||
"""
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
|
||||
from isaaclab_nav_task.navigation.config.rl_cfg import (
|
||||
RslRlOnPolicyRunnerCfg,
|
||||
RslRlPpoActorCriticCfg,
|
||||
RslRlPpoAlgorithmCfg,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class Go2NavPPORunnerCfg(RslRlOnPolicyRunnerCfg):
|
||||
"""PPO runner configuration for Go2 navigation, tuned for fine-tuning from a B2W checkpoint."""
|
||||
|
||||
num_steps_per_env = 16
|
||||
max_iterations = 5000
|
||||
save_interval = 200
|
||||
# Use TensorBoard by default; wandb requires WANDB_API_KEY and an interactive
|
||||
# tty for first-time login, neither of which is available inside the headless
|
||||
# docker container. All downstream cfgs (Dev / PureMaze) already override to
|
||||
# tensorboard; the base now matches so the mixed-terrain task
|
||||
# ``Isaac-Nav-PPO-Go2-v0`` no longer crashes on cold start.
|
||||
logger = "tensorboard"
|
||||
seed = 42
|
||||
wandb_project = "isaaclab_nav_go2"
|
||||
experiment_name = "go2_navigation_ppo_ft_from_b2w"
|
||||
empirical_normalization = False
|
||||
reward_shifting_value = 0.05
|
||||
|
||||
# IMPORTANT: must match the B2W policy block byte-for-byte so that
|
||||
# ActorCriticSRU.load_state_dict(..., strict=True) succeeds.
|
||||
policy = RslRlPpoActorCriticCfg(
|
||||
class_name="ActorCriticSRU",
|
||||
init_noise_std=0.5, # ↓ vs 1.0: warm-start, less exploration noise
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
activation="elu",
|
||||
rnn_hidden_size=512,
|
||||
rnn_type="lstm_sru",
|
||||
rnn_num_layers=1,
|
||||
dropout=0.2,
|
||||
num_cameras=1,
|
||||
image_input_dims=(64, 5, 8),
|
||||
height_input_dims=(64, 7, 7),
|
||||
)
|
||||
|
||||
# Conservative fine-tuning hyperparameters: smaller LR, tighter clip,
|
||||
# smaller KL target, lower grad-norm cap.
|
||||
algorithm = RslRlPpoAlgorithmCfg(
|
||||
class_name="PPO",
|
||||
value_loss_coef=0.05,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.1,
|
||||
value_clip_param=0.1,
|
||||
entropy_coef=0.001,
|
||||
num_learning_epochs=3,
|
||||
num_mini_batches=4,
|
||||
learning_rate=3.0e-4,
|
||||
schedule="adaptive",
|
||||
gamma=0.995,
|
||||
lam=0.95,
|
||||
desired_kl=0.005,
|
||||
max_grad_norm=0.5,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class Go2NavPPORunnerDevCfg(Go2NavPPORunnerCfg):
|
||||
"""Dev cfg used by the smoke-test task (Isaac-Nav-PPO-Go2-Dev-v0).
|
||||
|
||||
Unlike the production ``Go2NavPPORunnerCfg`` (whose hyperparameters are
|
||||
deliberately conservative for warm-starting from a B2W checkpoint), the dev
|
||||
cfg is meant for *from-scratch* Go2 training with the Odin1 camera. The
|
||||
conservative FT settings (entropy_coef=0.001, desired_kl=0.005, clip=0.1,
|
||||
lr=3e-4) caused exploration to collapse (noise_std plateaued at ~0.157) and
|
||||
success to stall at ~0.5 by 23k iters. Here we re-open exploration and the
|
||||
learning rate so the policy can escape that local optimum.
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.max_iterations = 300
|
||||
self.experiment_name = "go2_navigation_ppo_dev"
|
||||
self.logger = "tensorboard"
|
||||
|
||||
# ---- Phase 1 (from-scratch, 0->2000): re-open exploration ----
|
||||
# init_noise_std only takes effect on a from-scratch run (a resumed
|
||||
# run loads the std from the checkpoint, which was ~0.76 after Phase 1).
|
||||
# Phase 1 broke the 0.5 plateau (success ~0.65, noise_std held at 0.76),
|
||||
# but two issues appeared: (a) adaptive LR floored at 1e-5 because
|
||||
# desired_kl=0.01 was too tight against clip=0.2/entropy=0.005, throttling
|
||||
# late updates; (b) Loss/value_function spiked to ~127 around iter 1600.
|
||||
#
|
||||
# ---- Phase 2 (resume 2000->5000): refine & stabilise ----
|
||||
# Relax desired_kl so adaptive LR can climb off the 1e-5 floor, anneal
|
||||
# exploration down, and tighten clip to curb the value-loss instability.
|
||||
#
|
||||
# ---- Phase 3 (resume 7000->14000, difficulty [0.2, 0.6]): ----
|
||||
# Kept Phase 2 hyperparams. Result: success 0.93, base_contact -71%,
|
||||
# tip-over -76%. Confirmed depth is used and avoidance generalises when
|
||||
# the task forces it.
|
||||
#
|
||||
# ---- Phase 4 (resume 14000->24000, difficulty [0.3, 0.8]): ----
|
||||
# noise_std fell to ~0.20 in Phase 3 (close to collapse). The Phase 4
|
||||
# difficulty jump (0.6 -> 0.8) demands fresh exploration on unseen
|
||||
# terrain, so re-open entropy a notch. Result: success climbed 0 -> 0.88
|
||||
# (peak 0.91 @ iter 20519). But noise_std overshot to 0.33 by iter 23999,
|
||||
# LR hit 1e-5 floor again, and success retreated from 0.905 -> 0.876.
|
||||
#
|
||||
# ---- Phase 5 (resume 24000->?, difficulty [0.3, 0.8]): refine ----
|
||||
# Anneal entropy back to let noise_std settle ~0.22-0.25 so the policy
|
||||
# can stop trembling and finish the last refinement past 0.91 success.
|
||||
self.policy.init_noise_std = 1.0 # only matters on a fresh run
|
||||
self.algorithm.entropy_coef = 0.003 # 0.005 -> 0.003: anneal exploration for Phase 5 refinement
|
||||
self.algorithm.desired_kl = 0.02 # 0.01 -> 0.02: let adaptive LR recover off the 1e-5 floor
|
||||
self.algorithm.clip_param = 0.15 # 0.2 -> 0.15: smaller, more stable policy updates
|
||||
self.algorithm.value_clip_param = 0.15 # 0.2 -> 0.15: match, curb value-loss spikes
|
||||
self.algorithm.learning_rate = 1.0e-3 # adaptive ceiling reference
|
||||
self.algorithm.num_learning_epochs = 5 # keep: more updates per batch
|
||||
self.algorithm.max_grad_norm = 0.8 # 1.0 -> 0.8: tighten to damp the value-fn spike
|
||||
|
||||
|
||||
@configclass
|
||||
class Go2NavPPORunnerMixedCfg(Go2NavPPORunnerDevCfg):
|
||||
"""Cold-start hyperparameters for mixed-terrain training (maze + non_maze
|
||||
+ pits) on the full-size ``Isaac-Nav-PPO-Go2-v0`` task.
|
||||
|
||||
Inherits from the Dev cfg (which already overrides the B2W-FT defaults to
|
||||
something usable from scratch), but rolls the *exploration-heavy* knobs
|
||||
BACK to the Phase 1/2 values that were proven to break the 0.5 plateau in
|
||||
the dev curriculum, rather than the Phase 5 refinement settings the Dev cfg
|
||||
currently sits at.
|
||||
|
||||
Distilled lessons from the dev curriculum (see comments in
|
||||
``Go2NavPPORunnerDevCfg``):
|
||||
* Phase 1 (entropy=0.005, init_std=1.0): broke success 0.5 -> 0.65 plateau
|
||||
but kl=0.01 was too tight (LR floored at 1e-5).
|
||||
* Phase 2 (kl=0.02, clip=0.15, epochs=5): unflored the LR and tamed the
|
||||
Loss/value_function spike.
|
||||
* Phase 5 (entropy=0.003): refinement-only, kills exploration too early
|
||||
for a from-scratch run.
|
||||
|
||||
Cold-start recipe = Phase 1 exploration + Phase 2 stability:
|
||||
entropy=0.005, init_std=1.0, kl=0.02, clip=0.15, lr=1e-3, epochs=5,
|
||||
grad_norm=0.8.
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
# Independent experiment dir keeps mixed runs out of the dev TB tree.
|
||||
self.experiment_name = "go2_navigation_ppo_mixed"
|
||||
|
||||
# Override Dev's Phase 5 refinement back to cold-start values.
|
||||
self.policy.init_noise_std = 1.0 # Phase 1: re-open exploration
|
||||
self.algorithm.entropy_coef = 0.005 # Phase 1/4: cold-start exploration
|
||||
self.algorithm.desired_kl = 0.02 # Phase 2: avoid 1e-5 LR floor
|
||||
self.algorithm.clip_param = 0.15 # Phase 2: avoid value-loss spikes
|
||||
self.algorithm.value_clip_param = 0.15
|
||||
self.algorithm.learning_rate = 1.0e-3 # adaptive ceiling
|
||||
self.algorithm.num_learning_epochs = 5
|
||||
self.algorithm.max_grad_norm = 0.8
|
||||
|
||||
|
||||
@configclass
|
||||
class Go2NavPPORunnerPureMazeCfg(Go2NavPPORunnerDevCfg):
|
||||
"""Hyperparameters for the pure-maze training task.
|
||||
|
||||
Inherits the Phase 5 dev hyperparams (they shipped a strong checkpoint at
|
||||
iter 27800 on the mixed terrain). Only the log directory differs so
|
||||
pure-maze runs don't mix with the dev curriculum's TensorBoard history.
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.experiment_name = "go2_navigation_ppo_puremaze"
|
||||
+445
@@ -0,0 +1,445 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Unitree Go2 specific configuration for navigation environment."""
|
||||
|
||||
import os
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
|
||||
from isaaclab_nav_task.navigation.navigation_env_cfg import NavigationEnvCfg
|
||||
import isaaclab_nav_task.navigation.mdp as mdp
|
||||
|
||||
from isaaclab_nav_task.navigation.assets import GO2_CFG, ISAACLAB_NAV_TASKS_ASSETS_DIR # isort: skip
|
||||
|
||||
|
||||
# Go2 has 12 leg joints (no wheels). Joint names use the underscore convention
|
||||
# from the IsaacLab Go2 USD: FL_hip_joint, FL_thigh_joint, FL_calf_joint, ...
|
||||
LEG_JOINT_NAMES = [".*_hip_joint", ".*_thigh_joint", ".*_calf_joint"]
|
||||
|
||||
|
||||
@configclass
|
||||
class Go2NavigationEnvCfg(NavigationEnvCfg):
|
||||
"""Navigation env for Unitree Go2 (legs-only quadruped).
|
||||
|
||||
The locomotion checkpoint expected at
|
||||
``assets/data/Policies/locomotion/go2/policy_go2_jit.pt`` was converted from
|
||||
a rsl_rl checkpoint trained on Isaac-Velocity-Flat-Unitree-Go2-v0:
|
||||
- input : 48-dim observation (matches sru-navigation-sim's
|
||||
LowLevelPolicyCfg)
|
||||
- output : 12-dim joint position deltas (one per leg joint)
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
|
||||
from isaaclab_nav_task.navigation.mdp.observations import initialize_depth_noise_generator
|
||||
from isaaclab_nav_task.navigation.mdp.depth_utils.camera_config import get_camera_config
|
||||
|
||||
initialize_depth_noise_generator(robot_name="go2", use_jit_precompiled=False)
|
||||
camera_config = get_camera_config("go2")
|
||||
# Camera resolution is consumed implicitly through the encoder; sru-nav
|
||||
# always feeds (64, 40) into the VAE regardless of robot.
|
||||
_ = camera_config
|
||||
|
||||
# ---- Robot ----
|
||||
self.scene.robot = GO2_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
|
||||
|
||||
# Go2's body link is named "base" in the IsaacLab USD (not "base_link"
|
||||
# like B2W).
|
||||
self.scene.raycast_camera.prim_path = "{ENV_REGEX_NS}/Robot/base"
|
||||
# Real-robot Odin1 mount on Go2: 25.8 cm forward, 15.4 cm above body
|
||||
# center, camera looking straight ahead (no pitch).
|
||||
# NOTE: The extrinsic Tcl from the Odin1 calibration is camera-from-lidar,
|
||||
# not camera-from-base. Until we have T_base_lidar (lidar pose in Go2 base
|
||||
# frame), keep this measured mount offset rather than chaining Tcl.
|
||||
self.scene.raycast_camera.offset.pos = (0.258, 0.0, 0.154)
|
||||
self.scene.raycast_camera.offset.rot = (1.0, 0.0, 0.0, 0.0)
|
||||
self.scene.height_scanner_critic.prim_path = "{ENV_REGEX_NS}/Robot/base"
|
||||
|
||||
# ---- Camera intrinsics: Odin1 (LiDAR-aligned depth) ----
|
||||
# Real Odin1 @ 1600x1296:
|
||||
# fx = 737.357, fy = 737.292, cx = 794.372, cy = 666.259
|
||||
# hFOV ~94.67 deg, vFOV ~82.65 deg
|
||||
# To preserve the FOV while keeping the upstream downsample_factor=3
|
||||
# (raycast at 192x120 then downsample to the VAE-required 64x40), we
|
||||
# scale the intrinsics linearly:
|
||||
# fx_192 = 737.357 * 192/1600 = 88.48 px
|
||||
# fy_120 = 737.292 * 120/1296 = 68.27 px
|
||||
# cx_192 = 794.372 * 192/1600 = 95.32 px
|
||||
# cy_120 = 666.259 * 120/1296 = 61.69 px
|
||||
# Distortion (k2..k7) is ignored: raycast uses an ideal pinhole. The
|
||||
# real-side pipeline must undistort before feeding the network.
|
||||
# max_distance raised slightly because Odin1's wider FOV picks up more
|
||||
# nearby clutter; keep 11 m to match upstream.
|
||||
from isaaclab.sensors import patterns
|
||||
self.scene.raycast_camera.max_distance = 11.0
|
||||
self.scene.raycast_camera.pattern_cfg = patterns.PinholeCameraPatternCfg.from_ros_camera_info(
|
||||
fx=88.48,
|
||||
fy=68.27,
|
||||
cx=95.32,
|
||||
cy=61.69,
|
||||
width=192,
|
||||
height=120,
|
||||
downsample_factor=3, # 192x120 -> 64x40 to match VAE input
|
||||
)
|
||||
|
||||
# ---- Goal placement ----
|
||||
# The default goal height offset (0.2-0.8 m above ground) is tuned for
|
||||
# B2W's body height (~0.5-0.6 m). Go2 stands at ~0.3 m, so lower the
|
||||
# range to (0.1, 0.4) m to keep the goal marker near the Go2 body.
|
||||
# NOTE: success/termination is purely horizontal (xy), so this is mostly
|
||||
# a visual fix + removes the small constant z-error in the 3D
|
||||
# reach_goal_xy_soft reward term.
|
||||
self.commands.robot_goal.goal_height_offset_range = (0.1, 0.4)
|
||||
|
||||
# ---- Termination ----
|
||||
# Penalize/terminate if the body, hips, thighs, or head links hit
|
||||
# something. Go2's main body is "base"; the front "head" is split into
|
||||
# two extra rigid links ("Head_upper" / "Head_lower" in the upstream
|
||||
# Unitree Go2 USD) which stick forward ~10cm past the base. Without
|
||||
# those head links in the watch-list, GUI playback shows the snout
|
||||
# visually crashing into walls while the base contact sensor stays
|
||||
# silent and the episode does not terminate. Including them closes
|
||||
# that loophole.
|
||||
self.terminations.base_contact.params = {
|
||||
"sensor_cfg": SceneEntityCfg(
|
||||
"contact_forces",
|
||||
body_names=["base", ".*_hip", ".*_thigh", "Head_upper", "Head_lower"],
|
||||
),
|
||||
"threshold": 1.0,
|
||||
}
|
||||
|
||||
# Diagnostic-only: tighten `terrain_fall` so it actually fires when the
|
||||
# robot drops into the pit terrain. The base default of -2.0m never
|
||||
# triggered for Go2 (standing base_z ~+0.4m, typical pit depth 0.75-2.25m
|
||||
# -> body z bottoms out around -1.1m, well above -2.0m), making
|
||||
# `Episode_Termination/terrain_fall` permanently 0 across phases 1-5.
|
||||
# -0.3m means: body must be ~0.7m below normal standing height to fire,
|
||||
# which cleanly captures real pit falls without flagging walking dips.
|
||||
# NOTE: `terrain_fall` is `time_out=True`, so it does NOT trigger the
|
||||
# -50 `episode_termination` penalty -- this is purely a metric fix.
|
||||
self.terminations.terrain_fall.params = {"fall_height_threshold": -0.3}
|
||||
|
||||
# ---- Action interface (legs only, no wheels) ----
|
||||
self.actions.velocity_command.low_level_position_action = mdp.JointPositionActionCfg(
|
||||
asset_name="robot",
|
||||
joint_names=LEG_JOINT_NAMES,
|
||||
scale=0.25, # MUST match the scale used to train policy_go2_jit.pt
|
||||
use_default_offset=True,
|
||||
)
|
||||
# Disable the velocity-action branch entirely (patched in
|
||||
# navigation_se2_actions.py to be optional).
|
||||
self.actions.velocity_command.low_level_velocity_action = None
|
||||
self.actions.velocity_command.low_level_policy_file = os.path.join(
|
||||
ISAACLAB_NAV_TASKS_ASSETS_DIR,
|
||||
"Policies", "locomotion", "go2", "policy_go2_jit.pt",
|
||||
)
|
||||
# SE2 command space scale: Go2 cannot match B2W's wheeled top speed,
|
||||
# so we shrink the SRU command range to what Go2 can actually track.
|
||||
# The high-level SRU policy was trained on B2W with [1.0, 1.0, 1.0]; we
|
||||
# shrink each axis so the same network output produces feasible Go2
|
||||
# commands.
|
||||
self.actions.velocity_command.scale = [0.6, 0.3, 0.7]
|
||||
|
||||
# ---- Reward shaping ----
|
||||
self.rewards.joint_acc_l2_joint.params = {
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=LEG_JOINT_NAMES),
|
||||
}
|
||||
# Keep the B2W default termination penalty (-50.0). G0-G2 confirmed the
|
||||
# Go2 locomotion JIT is stable enough that we do not need the relaxed
|
||||
# -25 hedge anymore; restoring -50 makes wall hits / falls genuinely
|
||||
# costly so the policy stops trading collisions for goal reward.
|
||||
# (Only fires on `base_contact` + `large_pitch_angle`; `time_out`,
|
||||
# `at_goal`, `terrain_fall` are all marked time_out=True and don't
|
||||
# trigger this term.)
|
||||
self.rewards.episode_termination.weight = -50.0
|
||||
|
||||
# Forward-facing depth camera only senses obstacles ahead of the robot.
|
||||
# In G2.0 the agent learned to walk *backwards* (vx < 0) to hide obstacles
|
||||
# from the camera and bypass forward-collision termination. Enable the
|
||||
# backward-movement penalty (default weight is 0.0 in the base cfg).
|
||||
# NOTE: the base CurriculumCfg disables this penalty after 500 global
|
||||
# steps; if backward-walking re-emerges late in training, either delete
|
||||
# ``curriculum.disable_backward_penalty`` or lift its ``disable_after_steps``.
|
||||
self.rewards.backward_movement_penalty.weight = -1.0
|
||||
|
||||
# ---- Domain randomization ----
|
||||
# Go2 responds slower than wheels; widen the low-pass filter range.
|
||||
self.events.randomize_low_pass_filter_alpha.params = {
|
||||
"alpha_range": (0.3, 0.7),
|
||||
"action_term": "velocity_command",
|
||||
"per_dimension": True,
|
||||
"alpha_range_vx": (0.3, 0.7),
|
||||
"alpha_range_vy": (0.3, 0.7),
|
||||
"alpha_range_omega": (0.3, 0.7),
|
||||
}
|
||||
|
||||
# ---- Terrain ----
|
||||
# The Go2 locomotion .pt is flat-trained, so start with easier terrain.
|
||||
# Defaults assume RESUMING a non-trivial ckpt (max_init level 5, band
|
||||
# [0.3, 0.8]). For a COLD START set GO2_DIFFICULTY="0.0,0.4" (and the
|
||||
# env var GO2_INIT_LEVEL=0 if you want everyone to start at the easiest
|
||||
# row), otherwise the policy gets dropped into mid-difficulty terrain
|
||||
# and stalls. Curriculum stays ON in both cases so envs auto-promote
|
||||
# toward the high end as they succeed.
|
||||
lo, hi = 0.3, 0.8
|
||||
_diff = os.environ.get("GO2_DIFFICULTY", "").strip()
|
||||
if _diff:
|
||||
_parts = _diff.replace(" ", "").split(",")
|
||||
lo, hi = float(_parts[0]), float(_parts[1])
|
||||
init_level = 5
|
||||
_lvl = os.environ.get("GO2_INIT_LEVEL", "").strip()
|
||||
if _lvl:
|
||||
init_level = int(_lvl)
|
||||
self.scene.terrain.max_init_terrain_level = init_level
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [lo, hi]
|
||||
self.scene.terrain.terrain_generator.curriculum = True
|
||||
print(f"[Go2 Mixed] difficulty=[{lo}, {hi}], max_init_terrain_level={init_level} (curriculum on)")
|
||||
|
||||
# Drop the dedicated stairs sub-terrain for Go2: the locomotion JIT is
|
||||
# trained on flat ground only, climbing stairs is unreliable, and the
|
||||
# platform_mask on stair tops would otherwise bias goal sampling
|
||||
# towards stair platforms (see PositionSampler.platform_repeat_count=10
|
||||
# in mdp/navigation/goal_commands.py).
|
||||
# NOTE: ``add_stairs_to_maze`` flag exists in HfMazeTerrainCfg but is
|
||||
# never consumed by hf_terrains_maze.py, so the only real stair source
|
||||
# is the ``"stairs"`` sub-terrain. We rebalance the remaining 3 to keep
|
||||
# roughly the original maze:non_maze:pits ratio (3:2:2 -> ~0.43/0.29/0.29).
|
||||
sub_terrains = self.scene.terrain.terrain_generator.sub_terrains
|
||||
if "stairs" in sub_terrains:
|
||||
del sub_terrains["stairs"]
|
||||
if "maze" in sub_terrains:
|
||||
sub_terrains["maze"].proportion = 0.43
|
||||
if "non_maze" in sub_terrains:
|
||||
sub_terrains["non_maze"].proportion = 0.29
|
||||
if "pits" in sub_terrains:
|
||||
sub_terrains["pits"].proportion = 0.28
|
||||
|
||||
# ---- Shrink tile size for 12 GB VRAM ----
|
||||
# Defaults: 30m x 30m tiles with horizontal_scale=0.1 -> 300x300 cells
|
||||
# per tile heightfield, x 6 rows x 30 cols x 3 sub-terrains = ~16M
|
||||
# heightfield cells + 3 sub-terrain collision meshes. On a 12 GB 5070
|
||||
# that OOMs above ~768 envs. Halve cell_size to 1.0m (matches what the
|
||||
# PureMaze runs used, where 1024 envs ran stably): tile shrinks to
|
||||
# 15m x 15m, heightfield drops 4x to ~4M cells. Override via env var
|
||||
# GO2_CELL_SIZE if you want to go back to the paper-default 2.0m.
|
||||
_cell = os.environ.get("GO2_CELL_SIZE", "1.0").strip()
|
||||
if _cell:
|
||||
new_cell = float(_cell)
|
||||
for cfg_sub in sub_terrains.values():
|
||||
cfg_sub.cell_size = new_cell
|
||||
# grid_size stays (15, 15) -> tile_size = 15 * cell_size meters.
|
||||
tile_m = 15 * new_cell
|
||||
self.scene.terrain.terrain_generator.size = (tile_m, tile_m)
|
||||
print(f"[Go2 Mixed] cell_size={new_cell}m -> tile={tile_m}m "
|
||||
f"(set GO2_CELL_SIZE=2.0 to restore paper defaults)")
|
||||
|
||||
|
||||
@configclass
|
||||
class Go2NavigationEnvCfg_DEV(Go2NavigationEnvCfg):
|
||||
"""Development configuration with smaller terrain and lower difficulty."""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.scene.terrain.terrain_generator.num_rows = 2
|
||||
self.scene.terrain.terrain_generator.num_cols = 30
|
||||
# The Go2 locomotion JIT is flat-trained. Starting at difficulty 0.3-0.8
|
||||
# with no curriculum produced ~16% tip-overs (large_pitch_angle) and
|
||||
# stalled success at ~0.5. Start from the easiest terrain and let the
|
||||
# curriculum ramp difficulty up as the policy succeeds.
|
||||
#
|
||||
# Phase 1/2 trained on [0.0, 0.4] (success ~0.86, depth confirmed used
|
||||
# via ablation). Phase 3 raises the band so the avoidance behaviour
|
||||
# generalises to denser obstacles / deeper pits. Control the band per
|
||||
# run with DEV_DIFFICULTY="lo,hi" (default stays [0.0, 0.4]); e.g.
|
||||
# DEV_DIFFICULTY="0.2,0.6" (recommended Phase 3 bridge)
|
||||
# DEV_DIFFICULTY="0.3,0.8" (Phase 4, near paper level)
|
||||
# max_init_terrain_level stays 0 so a resumed policy re-enters at the
|
||||
# easy end of the new band and the curriculum re-ramps from there.
|
||||
lo, hi = 0.0, 0.4
|
||||
_dev_diff = os.environ.get("DEV_DIFFICULTY", "").strip()
|
||||
if _dev_diff:
|
||||
_parts = _dev_diff.replace(" ", "").split(",")
|
||||
lo, hi = float(_parts[0]), float(_parts[1])
|
||||
self.scene.terrain.max_init_terrain_level = 0
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [lo, hi]
|
||||
self.scene.terrain.terrain_generator.curriculum = True
|
||||
print(f"[Go2 DEV] terrain difficulty_range = [{lo}, {hi}] (curriculum on)")
|
||||
|
||||
|
||||
@configclass
|
||||
class Go2NavigationEnvCfg_PureMaze(Go2NavigationEnvCfg):
|
||||
"""Pure-maze training configuration.
|
||||
|
||||
All tiles are paper-style mazes:
|
||||
* 100% ``maze`` sub-terrain (no ``non_maze`` / ``pits`` sub-terrains).
|
||||
* Walls are uniform full-height rectangles (no randomized pillar/bar/cross
|
||||
obstacles inserted into the maze pattern).
|
||||
* Curriculum on, ramps difficulty from open to dense corridors.
|
||||
|
||||
Env-var overrides:
|
||||
* ``PUREMAZE_DIFFICULTY="lo,hi"``: difficulty band (default ``"0.0,1.0"``;
|
||||
higher band -> denser corridors / more dead-ends).
|
||||
* ``PUREMAZE_CELL_SIZE``: meters per maze cell (default 2.0). Lowering
|
||||
tightens corridors; tile size scales accordingly so total terrain area
|
||||
stays the same per tile in cells.
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
|
||||
sub_terrains = self.scene.terrain.terrain_generator.sub_terrains
|
||||
# Lock to 100% maze, paper-style walls.
|
||||
for name in list(sub_terrains.keys()):
|
||||
sub_terrains[name].proportion = 1.0 if name == "maze" else 0.0
|
||||
if "maze" in sub_terrains:
|
||||
sub_terrains["maze"].randomize_wall = False
|
||||
sub_terrains["maze"].random_wall_ratio = 0.0
|
||||
|
||||
# Difficulty range (curriculum on, ramps from easy to hard).
|
||||
lo, hi = 0.0, 1.0
|
||||
_diff = os.environ.get("PUREMAZE_DIFFICULTY", "").strip()
|
||||
if _diff:
|
||||
_parts = _diff.replace(" ", "").split(",")
|
||||
lo, hi = float(_parts[0]), float(_parts[1])
|
||||
self.scene.terrain.max_init_terrain_level = 0
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [lo, hi]
|
||||
self.scene.terrain.terrain_generator.curriculum = True
|
||||
|
||||
# Cell size (corridor width).
|
||||
_cell = os.environ.get("PUREMAZE_CELL_SIZE", "").strip()
|
||||
if _cell:
|
||||
new_cell = float(_cell)
|
||||
for cfg_sub in sub_terrains.values():
|
||||
cfg_sub.cell_size = new_cell
|
||||
tile_m = 15 * new_cell
|
||||
self.scene.terrain.terrain_generator.size = (tile_m, tile_m)
|
||||
print(f"[Go2 PureMaze] cell_size={new_cell}m -> tile={tile_m}m, "
|
||||
f"corridor ~{new_cell * 0.6:.2f}-{new_cell * 0.9:.2f}m")
|
||||
|
||||
# Terrain grid dimensions. Defaults are compact (3×10=30 tiles) to save
|
||||
# VRAM; the original 6×30=180 tiles wasted mesh memory when most tiles
|
||||
# were never visited (especially with few envs / play). Override via
|
||||
# PUREMAZE_NUM_ROWS / PUREMAZE_NUM_COLS for large-scale training.
|
||||
_rows = int(os.environ.get("PUREMAZE_NUM_ROWS", "3"))
|
||||
_cols = int(os.environ.get("PUREMAZE_NUM_COLS", "10"))
|
||||
self.scene.terrain.terrain_generator.num_rows = _rows
|
||||
self.scene.terrain.terrain_generator.num_cols = _cols
|
||||
|
||||
print(f"[Go2 PureMaze] 100% maze sub-terrain, classic walls, "
|
||||
f"difficulty=[{lo}, {hi}] (curriculum on), "
|
||||
f"grid={_rows}x{_cols}={_rows*_cols} tiles")
|
||||
|
||||
|
||||
@configclass
|
||||
class Go2NavigationEnvCfg_PLAY(Go2NavigationEnvCfg):
|
||||
"""Evaluation/visualization configuration."""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
|
||||
self.scene.num_envs = 20
|
||||
self.scene.env_spacing = 2.5
|
||||
self.scene.terrain.max_init_terrain_level = None
|
||||
if self.scene.terrain.terrain_generator is not None:
|
||||
self.scene.terrain.terrain_generator.num_rows = 2
|
||||
self.scene.terrain.terrain_generator.num_cols = 2
|
||||
|
||||
# ---- Evaluate on the SAME difficulty the policy was trained on ----
|
||||
# The base Go2NavigationEnvCfg uses difficulty_range=[0.3, 0.8], but
|
||||
# the Dev curriculum (what model_*.pt was actually trained on) ran on
|
||||
# [0.0, 0.4]. Replaying on [0.3, 0.8] shows the policy unseen-hard
|
||||
# terrain and makes it look like it "can't avoid" obstacles. Default
|
||||
# the PLAY difficulty to the training band for an honest read; allow
|
||||
# an explicit stress-test override via PLAY_DIFFICULTY="lo,hi".
|
||||
lo, hi = 0.0, 0.4
|
||||
_play_diff = os.environ.get("PLAY_DIFFICULTY", "").strip()
|
||||
if _play_diff:
|
||||
_parts = _play_diff.replace(" ", "").split(",")
|
||||
lo, hi = float(_parts[0]), float(_parts[1])
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [lo, hi]
|
||||
# Fixed difficulty band for evaluation (no curriculum ramping).
|
||||
self.scene.terrain.terrain_generator.curriculum = False
|
||||
print(f"[Go2 PLAY] terrain difficulty_range = [{lo}, {hi}] (curriculum off)")
|
||||
|
||||
# ---- Stress-test knobs for narrow-passage / dense-scene eval ----
|
||||
# PLAY_CELL_SIZE: meters per maze cell (default 2.0). Lowering this
|
||||
# tightens corridor width without changing tile size:
|
||||
# 2.0m -> 1.7m-2.0m open corridors (default; very wide for Go2)
|
||||
# 1.5m -> 1.0m-1.3m corridors (Go2 fits comfortably with margin)
|
||||
# 1.2m -> 0.7m-0.9m corridors (tight, real perception test)
|
||||
# 1.0m -> 0.5m-0.7m corridors (Go2 body width ~0.3m -> very tight)
|
||||
# NOTE: grid_size stays (15,15) so tile size shrinks proportionally
|
||||
# (30m -> 15m at cell_size=1.0). To preserve outer terrain footprint,
|
||||
# the tile size is rescaled below.
|
||||
_play_cell = os.environ.get("PLAY_CELL_SIZE", "").strip()
|
||||
# PLAY_MAZE_ONLY=1: force 100% maze sub-terrain (drop pits/non_maze)
|
||||
# so every tile is a corridor navigation challenge -- no easy open
|
||||
# tiles diluting the sample.
|
||||
_maze_only = os.environ.get("PLAY_MAZE_ONLY", "").strip().lower() in ("1", "true", "yes")
|
||||
|
||||
sub_terrains = self.scene.terrain.terrain_generator.sub_terrains
|
||||
if _maze_only and "maze" in sub_terrains:
|
||||
for name in list(sub_terrains.keys()):
|
||||
sub_terrains[name].proportion = 1.0 if name == "maze" else 0.0
|
||||
print(f"[Go2 PLAY] PLAY_MAZE_ONLY=1 -> 100% maze sub-terrain")
|
||||
|
||||
# PLAY_SUBTERRAIN_MIX="maze=0.5,non_maze=0.3,pits=0.2"
|
||||
# Override any subset of sub-terrain proportions. Names not listed
|
||||
# are left at their training-cfg default. Values do NOT need to sum
|
||||
# to 1.0 (IsaacLab normalizes them internally before sampling).
|
||||
# Setting a value to 0 effectively disables that sub-terrain.
|
||||
_mix = os.environ.get("PLAY_SUBTERRAIN_MIX", "").strip()
|
||||
if _mix:
|
||||
overrides = {}
|
||||
for part in _mix.split(","):
|
||||
if "=" not in part:
|
||||
continue
|
||||
k, v = part.split("=", 1)
|
||||
overrides[k.strip()] = float(v.strip())
|
||||
for name, prop in overrides.items():
|
||||
if name in sub_terrains:
|
||||
sub_terrains[name].proportion = prop
|
||||
else:
|
||||
print(f"[Go2 PLAY] WARNING: unknown sub-terrain '{name}' "
|
||||
f"(known: {list(sub_terrains.keys())})")
|
||||
final = {n: sub_terrains[n].proportion for n in sub_terrains}
|
||||
print(f"[Go2 PLAY] PLAY_SUBTERRAIN_MIX -> {final}")
|
||||
|
||||
if _play_cell:
|
||||
new_cell = float(_play_cell)
|
||||
for cfg_sub in sub_terrains.values():
|
||||
cfg_sub.cell_size = new_cell
|
||||
# Keep grid_size=(15,15) -> tile size = 15 * cell_size meters.
|
||||
# Match outer tile size so the scene grid stays consistent.
|
||||
tile_m = 15 * new_cell
|
||||
self.scene.terrain.terrain_generator.size = (tile_m, tile_m)
|
||||
print(f"[Go2 PLAY] PLAY_CELL_SIZE={new_cell}m -> tile={tile_m}m, "
|
||||
f"corridor ~{new_cell * 0.6:.2f}-{new_cell * 0.9:.2f}m")
|
||||
|
||||
# PLAY_CLASSIC_MAZE=1: reproduce paper-style mazes with uniform
|
||||
# rectangular walls and clear corridors (no randomized pillar / bar
|
||||
# / cross obstacles inserted into the maze pattern). Affects only
|
||||
# the `maze` sub-terrain since `non_maze` and `pits` rely on the
|
||||
# randomized shapes for their gameplay.
|
||||
_classic = os.environ.get("PLAY_CLASSIC_MAZE", "").strip().lower() in ("1", "true", "yes")
|
||||
if _classic and "maze" in sub_terrains:
|
||||
sub_terrains["maze"].randomize_wall = False
|
||||
sub_terrains["maze"].random_wall_ratio = 0.0
|
||||
print(f"[Go2 PLAY] PLAY_CLASSIC_MAZE=1 -> maze walls are full-height "
|
||||
f"rectangles (no randomized obstacle shapes)")
|
||||
|
||||
self.observations.policy.enable_corruption = False
|
||||
self.events.base_external_force_torque = None
|
||||
self.events.push_robot = None
|
||||
|
||||
# Visualize the depth-camera ray hits + height scanner grid in the GUI.
|
||||
# The RayCasterCamera draws a point per ray that hits geometry; the
|
||||
# density / pattern of these points is the effective FOV / resolution.
|
||||
self.scene.raycast_camera.debug_vis = True
|
||||
if self.scene.height_scanner_critic is not None:
|
||||
self.scene.height_scanner_critic.debug_vis = True
|
||||
@@ -0,0 +1,172 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Custom RSL-RL configuration classes for navigation tasks.
|
||||
|
||||
These config classes replace the standard Isaac Lab RL configs to support
|
||||
custom network architectures used in navigation tasks with depth camera inputs.
|
||||
"""
|
||||
|
||||
from dataclasses import MISSING
|
||||
from typing import Literal, Optional
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
|
||||
|
||||
@configclass
|
||||
class RslRlPpoActorCriticCfg:
|
||||
"""Configuration for the PPO actor-critic networks with navigation extensions."""
|
||||
|
||||
class_name: str = "ActorCritic"
|
||||
"""The policy class name. Default is ActorCritic."""
|
||||
|
||||
init_noise_std: float = MISSING
|
||||
"""The initial noise standard deviation for the policy."""
|
||||
|
||||
actor_hidden_dims: list[int] = MISSING
|
||||
"""The hidden dimensions of the actor network."""
|
||||
|
||||
critic_hidden_dims: list[int] = MISSING
|
||||
"""The hidden dimensions of the critic network."""
|
||||
|
||||
activation: str = MISSING
|
||||
"""The activation function for the actor and critic networks."""
|
||||
|
||||
rnn_type: str = "lstm"
|
||||
"""The type of RNN to use."""
|
||||
|
||||
rnn_hidden_size: int = 256
|
||||
"""The hidden size of the RNN."""
|
||||
|
||||
rnn_num_layers: int = 1
|
||||
"""The number of layers in the RNN."""
|
||||
|
||||
dropout: float = 0.0
|
||||
"""The dropout rate for the first layer of the actor and critic networks."""
|
||||
|
||||
# Visual inputs
|
||||
num_cameras: int = 1
|
||||
"""Number of depth cameras encoded into the observation (1 or 2)."""
|
||||
|
||||
image_input_dims: tuple[int, int, int] = (64, 5, 8)
|
||||
"""Encoded depth feature shape as (C, H, W)."""
|
||||
|
||||
height_input_dims: tuple[int, int, int] = (64, 7, 7)
|
||||
"""Encoded height scan feature shape as (C, H, W). Default is (64, 7, 7) for 64*7*7=3136 features."""
|
||||
|
||||
|
||||
@configclass
|
||||
class RslRlPpoAlgorithmCfg:
|
||||
"""Configuration for the PPO algorithm."""
|
||||
|
||||
class_name: str = MISSING
|
||||
"""The algorithm class name. Default is PPO."""
|
||||
|
||||
value_loss_coef: float = MISSING
|
||||
"""The coefficient for the value loss."""
|
||||
|
||||
use_clipped_value_loss: bool = MISSING
|
||||
"""Whether to use clipped value loss."""
|
||||
|
||||
clip_param: float = MISSING
|
||||
"""The clipping parameter for the policy."""
|
||||
|
||||
value_clip_param: float = 0.2
|
||||
"""The value clipping parameter. Default is 0.2."""
|
||||
|
||||
entropy_coef: float = MISSING
|
||||
"""The coefficient for the entropy loss."""
|
||||
|
||||
num_learning_epochs: int = MISSING
|
||||
"""The number of learning epochs per update."""
|
||||
|
||||
num_mini_batches: int = MISSING
|
||||
"""The number of mini-batches per update."""
|
||||
|
||||
learning_rate: float = MISSING
|
||||
"""The learning rate for the policy."""
|
||||
|
||||
schedule: str = MISSING
|
||||
"""The learning rate schedule."""
|
||||
|
||||
gamma: float = MISSING
|
||||
"""The discount factor."""
|
||||
|
||||
lam: float = MISSING
|
||||
"""The lambda parameter for Generalized Advantage Estimation (GAE)."""
|
||||
|
||||
desired_kl: float = MISSING
|
||||
"""The desired KL divergence."""
|
||||
|
||||
max_grad_norm: float = MISSING
|
||||
"""The maximum gradient norm."""
|
||||
|
||||
|
||||
@configclass
|
||||
class RslRlOnPolicyRunnerCfg:
|
||||
"""Configuration of the runner for on-policy algorithms."""
|
||||
|
||||
seed: Optional[int] = 42
|
||||
"""The seed for the experiment. Default is 42."""
|
||||
|
||||
device: str = "cuda:0"
|
||||
"""The device for the rl-agent. Default is cuda:0."""
|
||||
|
||||
num_steps_per_env: int = MISSING
|
||||
"""The number of steps per environment per update."""
|
||||
|
||||
max_iterations: int = MISSING
|
||||
"""The maximum number of iterations."""
|
||||
|
||||
empirical_normalization: bool = MISSING
|
||||
"""Whether to use empirical normalization."""
|
||||
|
||||
policy: RslRlPpoActorCriticCfg = MISSING
|
||||
"""The policy configuration."""
|
||||
|
||||
algorithm: RslRlPpoAlgorithmCfg = MISSING
|
||||
"""The algorithm configuration."""
|
||||
|
||||
reward_shifting_value: float = 0.0
|
||||
"""The value to shift the reward by. Default is 0.0."""
|
||||
|
||||
##
|
||||
# Checkpointing parameters
|
||||
##
|
||||
|
||||
save_interval: int = MISSING
|
||||
"""The number of iterations between saves."""
|
||||
|
||||
experiment_name: str = MISSING
|
||||
"""The experiment name."""
|
||||
|
||||
run_name: str = ""
|
||||
"""The run name. Default is empty string."""
|
||||
|
||||
##
|
||||
# Logging parameters
|
||||
##
|
||||
|
||||
logger: Literal["tensorboard", "neptune", "wandb"] = "tensorboard"
|
||||
"""The logger to use. Default is tensorboard."""
|
||||
|
||||
neptune_project: str = "isaaclab"
|
||||
"""The neptune project name. Default is "isaaclab"."""
|
||||
|
||||
wandb_project: str = "isaaclab"
|
||||
"""The wandb project name. Default is "isaaclab"."""
|
||||
|
||||
##
|
||||
# Loading parameters
|
||||
##
|
||||
|
||||
resume: bool = False
|
||||
"""Whether to resume. Default is False."""
|
||||
|
||||
load_run: str = ".*"
|
||||
"""The run directory to load. Default is ".*" (all)."""
|
||||
|
||||
load_checkpoint: str = "model_.*.pt"
|
||||
"""The checkpoint file to load. Default is "model_.*.pt" (all)."""
|
||||
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""MDP components for navigation tasks.
|
||||
|
||||
This module provides observations, rewards, terminations, curriculums,
|
||||
and actions specific to navigation tasks.
|
||||
"""
|
||||
|
||||
from isaaclab.envs.mdp import * # noqa: F401, F403
|
||||
|
||||
from .curriculums import * # noqa: F401, F403
|
||||
from .events import * # noqa: F401, F403
|
||||
from .rewards import * # noqa: F401, F403
|
||||
from .terminations import * # noqa: F401, F403
|
||||
from .observations import * # noqa: F401, F403
|
||||
from .navigation import * # noqa: F401, F403
|
||||
from .delay_manager import ObservationDelayManager, ObservationDelayManagerCfg, DelayBufferState # noqa: F401
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Curriculum functions for navigation tasks.
|
||||
|
||||
The functions can be passed to the :class:`isaaclab.managers.CurriculumTermCfg` object to enable
|
||||
the curriculum introduced by the function.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedRLEnv
|
||||
|
||||
|
||||
def disable_backward_penalty_after_steps(
|
||||
env: ManagerBasedRLEnv,
|
||||
env_ids: Sequence[int],
|
||||
term_name: str = "backward_movement_penalty",
|
||||
num_steps: int = 1000
|
||||
) -> torch.Tensor:
|
||||
"""Curriculum that disables the backward movement penalty after a certain number of steps.
|
||||
|
||||
This helps with early training by preventing backward movement, but removes the constraint
|
||||
later to allow more natural movement patterns.
|
||||
|
||||
Args:
|
||||
env: The learning environment.
|
||||
env_ids: Not used since all environments are affected.
|
||||
term_name: The name of the backward movement penalty term.
|
||||
num_steps: The number of steps after which the penalty should be disabled.
|
||||
|
||||
Returns:
|
||||
Current step counter as float for logging purposes.
|
||||
"""
|
||||
if env.common_step_counter > num_steps:
|
||||
# Check if the term exists and has a non-zero weight
|
||||
if hasattr(env.reward_manager, 'get_term_cfg'):
|
||||
try:
|
||||
term_cfg = env.reward_manager.get_term_cfg(term_name)
|
||||
if term_cfg.weight != 0.0:
|
||||
# Disable the penalty by setting weight to 0
|
||||
term_cfg.weight = 0.0
|
||||
env.reward_manager.set_term_cfg(term_name, term_cfg)
|
||||
print(f"Disabled backward movement penalty at step {env.common_step_counter}")
|
||||
except KeyError:
|
||||
# Term doesn't exist, which is fine
|
||||
pass
|
||||
|
||||
return torch.tensor(float(env.common_step_counter))
|
||||
@@ -0,0 +1,106 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Custom noise configurations for navigation tasks."""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
from isaaclab.utils.math import quat_from_euler_xyz, transform_points
|
||||
from isaaclab.utils.noise import NoiseCfg
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def goal_to_xyz(goal: torch.Tensor) -> torch.Tensor:
|
||||
"""Convert goal in (direction, log distance) format to (x, y, z) coordinates."""
|
||||
direction = goal[:, :3] # Extract directional vector
|
||||
distance = torch.exp(goal[:, 3]) - 1 # Convert log distance back to distance
|
||||
return direction * distance.unsqueeze(-1) # Scale direction by distance
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def xyz_to_goal(xyz: torch.Tensor) -> torch.Tensor:
|
||||
"""Convert (x, y, z) coordinates to (direction, log distance) format."""
|
||||
distance = torch.norm(xyz, dim=1, keepdim=True) + 1e-6 # Calculate distance
|
||||
direction = xyz / distance # Normalize to get directional vector
|
||||
distance = torch.log(1 + distance) # Convert distance to log distance
|
||||
return torch.cat([direction, distance], dim=1) # Concatenate direction and log distance
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def delta_transformation_noise(data: torch.Tensor, cfg: "DeltaTransformationNoiseCfg") -> torch.Tensor:
|
||||
"""Delta transformation noise involving random rotation and translation.
|
||||
|
||||
The output is returned in a new tensor (instead of modifying 'data' in-place).
|
||||
"""
|
||||
# Ensure the data has the shape (..., 3) for 3D coordinates, or (direction, log distance)
|
||||
if data.shape[-1] not in (3, 4):
|
||||
raise ValueError(
|
||||
"Data must have shape (..., 3) for 3D coordinates, or (..., 4) for (direction, log distance)."
|
||||
)
|
||||
|
||||
# Determine whether data is in (direction, log distance) format
|
||||
unit_vec = data.shape[-1] == 4
|
||||
|
||||
# Convert to (x, y, z) if needed
|
||||
if unit_vec:
|
||||
coordinate = goal_to_xyz(data)
|
||||
else:
|
||||
coordinate = data
|
||||
|
||||
# Generate small rotation noise (Rx, Ry, Rz) using uniform distribution
|
||||
rotation_noise = torch.empty((data.shape[0], 3), device=data.device).uniform_(-cfg.rotation, cfg.rotation)
|
||||
|
||||
# Convert Euler angles (Rx, Ry, Rz) to quaternions
|
||||
quat_noise = quat_from_euler_xyz(rotation_noise[:, 0], rotation_noise[:, 1], rotation_noise[:, 2])
|
||||
|
||||
# Generate small translation noise (Tx, Ty, Tz) using uniform distribution
|
||||
translation_noise = torch.empty((data.shape[0], 3), device=data.device).uniform_(-cfg.translation, cfg.translation)
|
||||
|
||||
# Apply random rotation + translation
|
||||
transformed_data = transform_points(coordinate.unsqueeze(1), translation_noise, quat_noise).squeeze(1)
|
||||
|
||||
# If input was in goal format, convert back
|
||||
if unit_vec:
|
||||
transformed_data = xyz_to_goal(transformed_data)
|
||||
|
||||
# Randomly apply the noise based on the probability
|
||||
random_mask = torch.rand(data.shape[0], device=data.device) < cfg.noise_prob
|
||||
|
||||
# Use torch.where to select transformed values on the mask
|
||||
random_mask_expanded = random_mask.unsqueeze(-1)
|
||||
output = torch.where(random_mask_expanded, transformed_data, data)
|
||||
|
||||
if cfg.remove_dist:
|
||||
if unit_vec:
|
||||
# Zero out the distance component by directly assigning 0
|
||||
output[..., 3:] = 0
|
||||
else:
|
||||
# Normalize the vector by its norm (avoiding division by zero)
|
||||
norm = torch.norm(output, dim=-1, keepdim=True) + 1e-6
|
||||
output = output / norm
|
||||
|
||||
return output
|
||||
|
||||
|
||||
@configclass
|
||||
class DeltaTransformationNoiseCfg(NoiseCfg):
|
||||
"""Configuration for small delta transformation noise involving rotation and translation."""
|
||||
|
||||
func = delta_transformation_noise
|
||||
|
||||
rotation: float = 0.1
|
||||
"""The maximum rotation angle in radians. Defaults to 0.1."""
|
||||
translation: float = 0.5
|
||||
"""The maximum translation in meters. Defaults to 0.5."""
|
||||
noise_prob: float = 0.25
|
||||
"""The probability of applying the noise. Defaults to 0.25."""
|
||||
remove_dist: bool = False
|
||||
"""Whether to remove the distance from the output. Defaults to False."""
|
||||
@@ -0,0 +1,272 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Standalone observation delay manager for simulating sensor delays."""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
|
||||
from isaaclab.utils import DelayBuffer, configclass
|
||||
|
||||
|
||||
@configclass
|
||||
class ObservationDelayManagerCfg:
|
||||
"""Configuration for the observation delay manager.
|
||||
|
||||
This configuration controls the simulation of sensor delays, where each environment
|
||||
can have a different random delay sampled from [0, max_delay] at episode reset.
|
||||
"""
|
||||
|
||||
enabled: bool = True
|
||||
"""Whether to enable delay simulation for observations."""
|
||||
|
||||
max_delay_lin_vel: int = 2
|
||||
"""Maximum delay in timesteps for linear velocity observations."""
|
||||
|
||||
max_delay_ang_vel: int = 2
|
||||
"""Maximum delay in timesteps for angular velocity observations."""
|
||||
|
||||
max_delay_projected_gravity: int = 2
|
||||
"""Maximum delay in timesteps for projected gravity observations."""
|
||||
|
||||
max_delay_target_position: int = 2
|
||||
"""Maximum delay in timesteps for target position/goal observations."""
|
||||
|
||||
max_delay_depth: int = 2
|
||||
"""Maximum delay in timesteps for depth image observations."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class DelayBufferState:
|
||||
"""Container for a delay buffer and its associated time lags.
|
||||
|
||||
This encapsulates the state needed for delayed observations:
|
||||
- The buffer storing historical data
|
||||
- Per-environment time lags (delays)
|
||||
- Maximum delay configured for this buffer
|
||||
"""
|
||||
|
||||
buffer: DelayBuffer
|
||||
time_lags: torch.Tensor
|
||||
max_delay: int
|
||||
|
||||
def compute(self, data: torch.Tensor) -> torch.Tensor:
|
||||
"""Append data and return delayed version."""
|
||||
return self.buffer.compute(data)
|
||||
|
||||
def reset(self, env_ids: Sequence[int] | None = None):
|
||||
"""Reset buffer for specified environments."""
|
||||
self.buffer.reset(env_ids)
|
||||
|
||||
def randomize_lags(self, env_ids: torch.Tensor, device: str):
|
||||
"""Randomize time lags for specified environments."""
|
||||
new_lags = torch.randint(0, self.max_delay + 1, (len(env_ids),), dtype=torch.int, device=device)
|
||||
self.time_lags[env_ids] = new_lags
|
||||
env_ids_list = env_ids.tolist() if env_ids is not None else None
|
||||
self.buffer.set_time_lag(new_lags, env_ids_list)
|
||||
|
||||
def set_lags(self, time_lags: torch.Tensor, env_ids: torch.Tensor | None = None):
|
||||
"""Set time lags for specified environments."""
|
||||
if env_ids is None:
|
||||
self.time_lags[:] = time_lags
|
||||
self.buffer.set_time_lag(time_lags)
|
||||
else:
|
||||
self.time_lags[env_ids] = time_lags
|
||||
self.buffer.set_time_lag(time_lags, env_ids.tolist())
|
||||
|
||||
|
||||
class ObservationDelayManager:
|
||||
"""Standalone manager for simulating observation delays.
|
||||
|
||||
This manager can be added to any environment to simulate sensor delays
|
||||
without coupling to a specific action term. Each observation type can
|
||||
have its own delay buffer with per-environment random delays.
|
||||
|
||||
Usage:
|
||||
# In environment config or __init__:
|
||||
self.delay_manager = ObservationDelayManager(cfg, num_envs, device)
|
||||
|
||||
# In observation functions:
|
||||
delayed_vel = env.delay_manager.compute_delayed_lin_vel(current_vel)
|
||||
|
||||
# On episode reset:
|
||||
env.delay_manager.reset(env_ids)
|
||||
env.delay_manager.randomize_lags(env_ids)
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: ObservationDelayManagerCfg, num_envs: int, device: str):
|
||||
"""Initialize the observation delay manager.
|
||||
|
||||
Args:
|
||||
cfg: Configuration for delay buffers.
|
||||
num_envs: Number of parallel environments.
|
||||
device: Device for tensor operations.
|
||||
"""
|
||||
self.cfg = cfg
|
||||
self.num_envs = num_envs
|
||||
self.device = device
|
||||
self._delay_enabled = cfg.enabled
|
||||
|
||||
self._init_delay_buffers()
|
||||
|
||||
def _create_delay_buffer_state(self, max_delay: int) -> DelayBufferState:
|
||||
"""Create a DelayBufferState with initialized buffer and random time lags."""
|
||||
buffer = DelayBuffer(max_delay, batch_size=self.num_envs, device=self.device)
|
||||
time_lags = torch.randint(0, max_delay + 1, (self.num_envs,), dtype=torch.int, device=self.device)
|
||||
buffer.set_time_lag(time_lags)
|
||||
return DelayBufferState(buffer=buffer, time_lags=time_lags, max_delay=max_delay)
|
||||
|
||||
def _init_delay_buffers(self):
|
||||
"""Initialize all delay buffers based on configuration."""
|
||||
if not self.cfg.enabled:
|
||||
# No delays - set all delay states to None
|
||||
self._delay_lin_vel: DelayBufferState | None = None
|
||||
self._delay_ang_vel: DelayBufferState | None = None
|
||||
self._delay_projected_gravity: DelayBufferState | None = None
|
||||
self._delay_target_position: DelayBufferState | None = None
|
||||
self._delay_depth_buffers: dict[str, DelayBufferState] = {}
|
||||
self._max_delay_depth = 0
|
||||
return
|
||||
|
||||
# Initialize delay buffer states for each observation type
|
||||
self._delay_lin_vel = self._create_delay_buffer_state(self.cfg.max_delay_lin_vel)
|
||||
self._delay_ang_vel = self._create_delay_buffer_state(self.cfg.max_delay_ang_vel)
|
||||
self._delay_projected_gravity = self._create_delay_buffer_state(self.cfg.max_delay_projected_gravity)
|
||||
self._delay_target_position = self._create_delay_buffer_state(self.cfg.max_delay_target_position)
|
||||
# Depth buffers are created dynamically per camera
|
||||
self._delay_depth_buffers: dict[str, DelayBufferState] = {}
|
||||
self._max_delay_depth = self.cfg.max_delay_depth
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
"""Whether delay simulation is enabled."""
|
||||
return self._delay_enabled
|
||||
|
||||
def get_or_create_depth_delay_buffer(self, camera_name: str) -> DelayBufferState:
|
||||
"""Get or create a delay buffer for a specific depth camera.
|
||||
|
||||
Args:
|
||||
camera_name: Name of the depth camera sensor.
|
||||
|
||||
Returns:
|
||||
DelayBufferState for the specified camera.
|
||||
"""
|
||||
if camera_name not in self._delay_depth_buffers:
|
||||
self._delay_depth_buffers[camera_name] = self._create_delay_buffer_state(self._max_delay_depth)
|
||||
return self._delay_depth_buffers[camera_name]
|
||||
|
||||
def reset(self, env_ids: torch.Tensor):
|
||||
"""Reset all delay buffers for specified environments.
|
||||
|
||||
Args:
|
||||
env_ids: Environment indices to reset.
|
||||
"""
|
||||
env_ids_list = env_ids.tolist() if env_ids is not None else None
|
||||
|
||||
if self._delay_lin_vel is not None:
|
||||
self._delay_lin_vel.reset(env_ids_list)
|
||||
if self._delay_ang_vel is not None:
|
||||
self._delay_ang_vel.reset(env_ids_list)
|
||||
if self._delay_projected_gravity is not None:
|
||||
self._delay_projected_gravity.reset(env_ids_list)
|
||||
if self._delay_target_position is not None:
|
||||
self._delay_target_position.reset(env_ids_list)
|
||||
|
||||
# Reset all depth buffers
|
||||
for depth_buffer in self._delay_depth_buffers.values():
|
||||
depth_buffer.reset(env_ids_list)
|
||||
|
||||
def randomize_lags(self, env_ids: torch.Tensor):
|
||||
"""Randomize time lags for all delay buffers.
|
||||
|
||||
Args:
|
||||
env_ids: Environment indices to randomize.
|
||||
"""
|
||||
if self._delay_lin_vel is not None:
|
||||
self._delay_lin_vel.randomize_lags(env_ids, self.device)
|
||||
if self._delay_ang_vel is not None:
|
||||
self._delay_ang_vel.randomize_lags(env_ids, self.device)
|
||||
if self._delay_projected_gravity is not None:
|
||||
self._delay_projected_gravity.randomize_lags(env_ids, self.device)
|
||||
if self._delay_target_position is not None:
|
||||
self._delay_target_position.randomize_lags(env_ids, self.device)
|
||||
|
||||
# Randomize all depth buffers
|
||||
for depth_buffer in self._delay_depth_buffers.values():
|
||||
depth_buffer.randomize_lags(env_ids, self.device)
|
||||
|
||||
# ============================================================================
|
||||
# Delayed observation computation methods
|
||||
# ============================================================================
|
||||
|
||||
def compute_delayed_lin_vel(self, lin_vel: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute delayed linear velocity observation.
|
||||
|
||||
Args:
|
||||
lin_vel: Current linear velocity tensor of shape (num_envs, 3).
|
||||
|
||||
Returns:
|
||||
Delayed linear velocity, or original if delays disabled.
|
||||
"""
|
||||
if self._delay_lin_vel is None:
|
||||
return lin_vel
|
||||
return self._delay_lin_vel.compute(lin_vel)
|
||||
|
||||
def compute_delayed_ang_vel(self, ang_vel: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute delayed angular velocity observation.
|
||||
|
||||
Args:
|
||||
ang_vel: Current angular velocity tensor of shape (num_envs, 3).
|
||||
|
||||
Returns:
|
||||
Delayed angular velocity, or original if delays disabled.
|
||||
"""
|
||||
if self._delay_ang_vel is None:
|
||||
return ang_vel
|
||||
return self._delay_ang_vel.compute(ang_vel)
|
||||
|
||||
def compute_delayed_projected_gravity(self, projected_gravity: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute delayed projected gravity observation.
|
||||
|
||||
Args:
|
||||
projected_gravity: Current projected gravity tensor of shape (num_envs, 3).
|
||||
|
||||
Returns:
|
||||
Delayed projected gravity, or original if delays disabled.
|
||||
"""
|
||||
if self._delay_projected_gravity is None:
|
||||
return projected_gravity
|
||||
return self._delay_projected_gravity.compute(projected_gravity)
|
||||
|
||||
def compute_delayed_target_position(self, target_position: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute delayed target position observation.
|
||||
|
||||
Args:
|
||||
target_position: Current target position tensor.
|
||||
|
||||
Returns:
|
||||
Delayed target position, or original if delays disabled.
|
||||
"""
|
||||
if self._delay_target_position is None:
|
||||
return target_position
|
||||
return self._delay_target_position.compute(target_position)
|
||||
|
||||
def compute_delayed_depth(self, depth_features: torch.Tensor, camera_name: str) -> torch.Tensor:
|
||||
"""Compute delayed depth observation for a specific camera.
|
||||
|
||||
Args:
|
||||
depth_features: Encoded depth features of shape (num_envs, feature_dim).
|
||||
camera_name: Name of the camera sensor.
|
||||
|
||||
Returns:
|
||||
Delayed depth features, or original if delays disabled.
|
||||
"""
|
||||
if not self._delay_enabled:
|
||||
return depth_features
|
||||
delay_buffer = self.get_or_create_depth_delay_buffer(camera_name)
|
||||
return delay_buffer.compute(depth_features)
|
||||
+332
@@ -0,0 +1,332 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
"""Depth noise encoder module for visual navigation.
|
||||
|
||||
This module provides a VAE-based depth encoder with realistic stereo depth noise simulation.
|
||||
"""
|
||||
import os
|
||||
from collections import OrderedDict
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torchvision.models import regnet_x_400mf
|
||||
from torchvision.ops import Conv2dNormActivation, FeaturePyramidNetwork
|
||||
|
||||
from isaaclab_assets import ISAACLAB_ASSETS_DATA_DIR
|
||||
|
||||
from .camera_config import CameraConfig
|
||||
|
||||
# Default camera configuration (ZedX Camera for b2w and aow d)
|
||||
DEFAULT_CAMERA_CONFIG = CameraConfig(
|
||||
focal_length=25.0,
|
||||
baseline=0.12,
|
||||
min_depth=0.25,
|
||||
max_depth=10.0,
|
||||
depth_encoder_path=os.path.join(str(ISAACLAB_ASSETS_DATA_DIR), "Policies", "RSL-ETHZ/AoW_d/depth_encoder", "vae_pretrain_fuse.pth")
|
||||
)
|
||||
|
||||
class VAESampler(nn.Module):
|
||||
def __init__(self, input_dim, latent_dim):
|
||||
super(VAESampler, self).__init__()
|
||||
self.input_dim = input_dim
|
||||
self.latent_dim = latent_dim
|
||||
self.conv = Conv2dNormActivation(input_dim, latent_dim, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
|
||||
# Convolutional Layers for 2D mean and logvar
|
||||
self.mean_layers = nn.Sequential(
|
||||
Conv2dNormActivation(latent_dim, latent_dim, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
nn.Conv2d(latent_dim, latent_dim, kernel_size=1, stride=1, padding=0)
|
||||
)
|
||||
|
||||
self.logvar_layers = nn.Sequential(
|
||||
Conv2dNormActivation(latent_dim, latent_dim, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
nn.Conv2d(latent_dim, latent_dim, kernel_size=1, stride=1, padding=0)
|
||||
)
|
||||
|
||||
def reparameterize(self, mu, logvar):
|
||||
std = torch.exp(0.5 * logvar)
|
||||
eps = torch.randn_like(std)
|
||||
return mu + eps * std
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.mean_layers(x)
|
||||
return x
|
||||
|
||||
class EncoderFPN(nn.Module):
|
||||
def __init__(self, in_channel, out_channel):
|
||||
super(EncoderFPN, self).__init__()
|
||||
encoder = regnet_x_400mf(weights=None)
|
||||
# Remove classification head from the encoder
|
||||
encoder = nn.Sequential(*list(encoder.children())[:-2])
|
||||
# Modify the first layer to accept the number of channels in the input image
|
||||
encoder[0][0] = nn.Conv2d(in_channel, 32, kernel_size=3, stride=2, padding=1, bias=False)
|
||||
self.enc = encoder[0]
|
||||
self.enc_1 = encoder[1][:2]
|
||||
self.enc_2 = encoder[1][2]
|
||||
self.enc_3 = encoder[1][3]
|
||||
|
||||
# Feature Pyramid Network
|
||||
self.fpn = FeaturePyramidNetwork([64, 160, 400], out_channel)
|
||||
|
||||
def forward(self, x):
|
||||
raise NotImplementedError("Forward method must be implemented in the child class.")
|
||||
|
||||
|
||||
class DepthEncoder(EncoderFPN):
|
||||
"""Depth image encoder using Feature Pyramid Network."""
|
||||
def __init__(self, out_channel):
|
||||
super(DepthEncoder, self).__init__(1, out_channel)
|
||||
|
||||
def forward(self, x):
|
||||
# check if depth has channel dimension
|
||||
if x.dim() == 3:
|
||||
x = x.unsqueeze(1)
|
||||
|
||||
out = OrderedDict()
|
||||
x = self.enc(x)
|
||||
out['feat1'] = self.enc_1(x)
|
||||
out['feat2'] = self.enc_2(out['feat1'])
|
||||
out['feat3'] = self.enc_3(out['feat2'])
|
||||
|
||||
out = self.fpn(out)
|
||||
|
||||
return out['feat1']
|
||||
|
||||
|
||||
class VAEDecoder(nn.Module):
|
||||
def __init__(self, input_dim, out_dim):
|
||||
super(VAEDecoder, self).__init__()
|
||||
self.input_dim = input_dim
|
||||
self.conv = Conv2dNormActivation(input_dim, input_dim, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
|
||||
self.decoder = nn.Sequential(
|
||||
# Layer 0
|
||||
nn.ConvTranspose2d(input_dim, input_dim, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
# Layer 1
|
||||
nn.ConvTranspose2d(input_dim, input_dim, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
# Layer 2
|
||||
nn.ConvTranspose2d(input_dim, input_dim, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
# Final Output Layer
|
||||
nn.Conv2d(input_dim, out_dim, kernel_size=3, stride=1, padding=1)
|
||||
) # Expand the spatial dimensions by a factor of 2**4=16
|
||||
|
||||
def forward(self, z):
|
||||
z = self.conv(z)
|
||||
img = self.decoder(z)
|
||||
return img
|
||||
|
||||
class DepthDecoder(VAEDecoder):
|
||||
def __init__(self, input_dim):
|
||||
super(DepthDecoder, self).__init__(input_dim, 1)
|
||||
|
||||
class VAENet(nn.Module):
|
||||
def __init__(self, latent_dim):
|
||||
super(VAENet, self).__init__()
|
||||
self.depth_encoder = DepthEncoder(latent_dim)
|
||||
|
||||
self.vae_sampler = VAESampler(latent_dim, latent_dim)
|
||||
|
||||
self.depth_decoder = DepthDecoder(latent_dim)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.depth_encoder(x)
|
||||
x = self.vae_sampler(x)
|
||||
return x
|
||||
|
||||
def decode(self, z):
|
||||
return self.depth_decoder(z)
|
||||
|
||||
class DepthNoise(torch.nn.Module):
|
||||
def __init__(self,
|
||||
focal_length,
|
||||
baseline,
|
||||
min_depth,
|
||||
max_depth,
|
||||
filter_size=3,
|
||||
inlier_thred_range=(0.01, 0.05),
|
||||
prob_range=(0.4, 0.6),
|
||||
invalid_disp=1e7
|
||||
):
|
||||
"""
|
||||
A Simply PyTorch module to add realistic noise to depth images.
|
||||
|
||||
Args:
|
||||
focal_length (float): Focal length of the camera (in pixels).
|
||||
baseline (float): Baseline distance between stereo cameras (in meters).
|
||||
min_depth (float): Minimum depth value after clamping.
|
||||
max_depth (float): Maximum depth value after clamping.
|
||||
filter_size (int): Kernel size for local mean disparity computation. (tuning based on image resolution)
|
||||
inlier_thred_range (tuple): Threshold range for normalized disparity differences.
|
||||
prob_range (tuple): Probability range for matching pixels.
|
||||
invalid_disp (float): Invalid disparity
|
||||
|
||||
"""
|
||||
super().__init__()
|
||||
self.focal_length = focal_length
|
||||
self.baseline = baseline
|
||||
self.min_depth = min_depth
|
||||
self.max_depth = max_depth
|
||||
self.invalid_disp = invalid_disp
|
||||
self.inlier_thred_range = inlier_thred_range
|
||||
self.prob_range = prob_range
|
||||
self.filter_size = filter_size
|
||||
|
||||
weights, substitutes = self._compute_weights(filter_size)
|
||||
self.register_buffer('weights', weights.view(1, 1, filter_size, filter_size))
|
||||
self.register_buffer('substitutes', substitutes.view(1, 1, filter_size, filter_size))
|
||||
|
||||
|
||||
def _compute_weights(self, filter_size):
|
||||
"""
|
||||
Compute weights and substitutes for disparity filtering.
|
||||
|
||||
Args:
|
||||
filter_size (int): Kernel size for local mean disparity computation.
|
||||
"""
|
||||
center = filter_size // 2
|
||||
idx = torch.arange(filter_size) - center
|
||||
x_filter, y_filter = torch.meshgrid(idx, idx, indexing='ij')
|
||||
sqr_radius = x_filter ** 2 + y_filter ** 2
|
||||
sqrt_radius = torch.sqrt(sqr_radius)
|
||||
weights = 1 / torch.where(sqr_radius == 0, torch.ones_like(sqrt_radius), sqrt_radius)
|
||||
weights = weights / weights.sum()
|
||||
fill_weights = 1 / (1 + sqrt_radius)
|
||||
fill_weights = torch.where(sqr_radius > filter_size, -1.0, fill_weights)
|
||||
substitutes = (fill_weights > 0).float()
|
||||
|
||||
return weights, substitutes
|
||||
|
||||
def filter_disparity(self, disparity):
|
||||
"""
|
||||
Filter the disparity map using local mean disparity.
|
||||
|
||||
Args:
|
||||
disparity (torch.Tensor): Input disparity map tensor of shape (B, C, H, W).
|
||||
"""
|
||||
B, _, H, W = disparity.shape
|
||||
device = disparity.device
|
||||
center = self.filter_size // 2
|
||||
|
||||
output_disparity = torch.full_like(disparity, self.invalid_disp)
|
||||
|
||||
prob = torch.rand(B, 1, 1, 1, device=device) * (self.prob_range[1] - self.prob_range[0]) + self.prob_range[0]
|
||||
random_mask = (torch.rand(B, 1, H, W, device=device) < prob)
|
||||
|
||||
# Compute mean disparity
|
||||
weighted_disparity = F.conv2d(disparity, self.weights, padding=center)
|
||||
|
||||
# Compute differences
|
||||
differences = torch.abs(disparity - weighted_disparity)
|
||||
|
||||
# Normalize differences based on current image statistics for consistent thresholding
|
||||
differences_flat = differences.view(B, -1) # Flatten spatial dimensions
|
||||
mean_diff = torch.mean(differences_flat, dim=1, keepdim=True)
|
||||
std_diff = torch.std(differences_flat, dim=1, keepdim=True) + 1e-6 # Add epsilon to avoid division by zero
|
||||
|
||||
# Normalize differences: (diff - mean) / std, then shift to [0, 1] range approximately
|
||||
normalized_differences_flat = (differences_flat - mean_diff) / std_diff
|
||||
normalized_differences = normalized_differences_flat.view_as(differences)
|
||||
|
||||
# Use parameter-based threshold on normalized differences
|
||||
threshold = torch.rand(B, 1, 1, 1, device=device) * (self.inlier_thred_range[1] - self.inlier_thred_range[0]) + self.inlier_thred_range[0]
|
||||
update_mask = (normalized_differences < threshold) & random_mask
|
||||
|
||||
# Compute output value: round with 1/32 precision
|
||||
disparity = torch.round(disparity * 32.0) / 32.0
|
||||
|
||||
# Update output disparity
|
||||
output_disparity = torch.where(update_mask, disparity, output_disparity)
|
||||
|
||||
# Apply substitutes to fill neighboring pixels
|
||||
filled_values = F.conv2d(update_mask.float() * disparity, self.substitutes, padding=center)
|
||||
counts = F.conv2d(update_mask.float(), self.substitutes, padding=center) + 1e-9
|
||||
average_filled_values = filled_values / counts
|
||||
output_disparity = torch.where(counts >= 1, average_filled_values, output_disparity)
|
||||
|
||||
return output_disparity
|
||||
|
||||
def forward(self, depth, add_noise: bool) -> torch.Tensor:
|
||||
# correct input shape
|
||||
if len(depth.shape) == 3:
|
||||
depth = depth.unsqueeze(1) # add channel dimension
|
||||
|
||||
# check dimension (B, 1, H, W)
|
||||
assert depth.shape[1] == 1, "Input depth tensor must have shape (B, 1, H, W)."
|
||||
assert len(depth.shape) == 4, "Input depth tensor must have shape (B, 1, H, W)."
|
||||
|
||||
if add_noise:
|
||||
# Clamp the depth values
|
||||
depth = torch.clamp(depth, min=1. / self.invalid_disp)
|
||||
|
||||
# Step 1: Convert depth to disparity
|
||||
disparity = self.focal_length * self.baseline / depth
|
||||
|
||||
# Step 2: Filter the disparity map
|
||||
filtered_disparity = self.filter_disparity(disparity)
|
||||
|
||||
# Step 3: Recompute depth from disparity
|
||||
depth = self.focal_length * self.baseline / filtered_disparity
|
||||
|
||||
# Step 4: Clamp the depth values
|
||||
depth[depth < self.min_depth] = 0.0
|
||||
|
||||
# Step 5: Set invalid depth values to 0.0 (values outside valid range are not measurable)
|
||||
depth[depth > self.max_depth] = 0.0
|
||||
|
||||
return depth
|
||||
|
||||
|
||||
class DepthNoiseEncoder(torch.nn.Module):
|
||||
def __init__(self,
|
||||
feature_dim,
|
||||
camera_config: CameraConfig = None
|
||||
):
|
||||
"""
|
||||
A Simply PyTorch module to add realistic noise to depth images.
|
||||
|
||||
Args:
|
||||
feature_dim (int): Number of output channels from the encoder.
|
||||
camera_config (CameraConfig, optional): Camera configuration parameters.
|
||||
If None, uses DEFAULT_CAMERA_CONFIG.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Use provided config or default
|
||||
if camera_config is None:
|
||||
camera_config = DEFAULT_CAMERA_CONFIG
|
||||
|
||||
self.camera_config = camera_config
|
||||
self.depth_noise = DepthNoise(
|
||||
focal_length=camera_config.focal_length,
|
||||
baseline=camera_config.baseline,
|
||||
min_depth=camera_config.min_depth,
|
||||
max_depth=camera_config.max_depth
|
||||
)
|
||||
|
||||
self.encoder = VAENet(feature_dim)
|
||||
|
||||
try:
|
||||
self.encoder.load_state_dict(torch.load(camera_config.depth_encoder_path, weights_only=True), strict=True)
|
||||
print('\033[92m' + f'Successfully loaded pre-trained weights from {camera_config.depth_encoder_path}' + '\033[0m')
|
||||
except Exception as e:
|
||||
print('\033[91m' + f'Failed to load pre-trained weights: {e}' + '\033[0m')
|
||||
|
||||
def forward(self, depth: torch.Tensor, add_noise: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
noisy_depth = self.depth_noise(depth, add_noise)
|
||||
encoded_depth = self.encoder(noisy_depth) # (B, 1, H, W) -> (B, C, H//8, W//8)
|
||||
return encoded_depth, noisy_depth
|
||||
|
||||
@torch.jit.export
|
||||
def decode(self, z: torch.Tensor) -> torch.Tensor:
|
||||
return self.encoder.decode(z)
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Depth noise and encoding utilities for navigation tasks."""
|
||||
|
||||
from .camera_config import (
|
||||
CameraConfig,
|
||||
ZEDX_CAMERA_CONFIG,
|
||||
DEFAULT_CAMERA_CONFIG,
|
||||
ROBOT_CAMERA_CONFIGS,
|
||||
get_camera_config,
|
||||
)
|
||||
from .depth_noise_encoder import DepthNoiseEncoder, DepthNoise
|
||||
|
||||
__all__ = [
|
||||
"CameraConfig",
|
||||
"ZEDX_CAMERA_CONFIG",
|
||||
"DEFAULT_CAMERA_CONFIG",
|
||||
"ROBOT_CAMERA_CONFIGS",
|
||||
"get_camera_config",
|
||||
"DepthNoiseEncoder",
|
||||
"DepthNoise",
|
||||
]
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Camera configuration parameters for different robots and camera types.
|
||||
|
||||
This module provides camera-specific parameters for depth noise generation and encoding.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
|
||||
# Local assets directory for this extension
|
||||
# Path: depth_utils -> mdp -> navigation -> assets/data
|
||||
_ASSETS_DIR = os.path.normpath(
|
||||
os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "assets", "data")
|
||||
)
|
||||
|
||||
|
||||
def _get_encoder_path(model_filename: str) -> str:
|
||||
"""Helper function to construct encoder model path.
|
||||
|
||||
Args:
|
||||
model_filename: Name of the encoder model file (e.g., 'vae_pretrain_new.pth')
|
||||
|
||||
Returns:
|
||||
Full path to the encoder model file
|
||||
"""
|
||||
return os.path.join(_ASSETS_DIR, "Policies", "depth_encoder", model_filename)
|
||||
|
||||
|
||||
@configclass
|
||||
class CameraConfig:
|
||||
"""Configuration class for camera parameters.
|
||||
|
||||
This class contains all camera-specific parameters needed for depth noise generation
|
||||
and depth encoder initialization.
|
||||
"""
|
||||
|
||||
# Camera intrinsic parameters
|
||||
focal_length: float = 25.0
|
||||
baseline: float = 0.12
|
||||
|
||||
# Depth range parameters
|
||||
min_depth: float = 0.25
|
||||
max_depth: float = 10.0
|
||||
|
||||
# Camera resolution (width, height)
|
||||
resolution: tuple[int, int] = (53, 30)
|
||||
|
||||
# Depth encoder model path
|
||||
depth_encoder_path: str = ""
|
||||
|
||||
def __post_init__(self):
|
||||
"""Post-initialization to set default encoder path if not provided."""
|
||||
if not self.depth_encoder_path:
|
||||
self.depth_encoder_path = _get_encoder_path("vae_pretrain_fuse.pth")
|
||||
|
||||
|
||||
# Predefined camera configurations
|
||||
ZEDX_CAMERA_CONFIG = CameraConfig(
|
||||
focal_length=25.0,
|
||||
baseline=0.12,
|
||||
min_depth=0.25,
|
||||
max_depth=10.0,
|
||||
resolution=(64, 40),
|
||||
depth_encoder_path=_get_encoder_path("vae_pretrain_new.pth"),
|
||||
)
|
||||
"""Configuration for ZedX camera (used with b2w and aow_d robots)."""
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Odin1 (LiDAR-aligned depth on Unitree Go2)
|
||||
# -----------------------------------------------------------------------------
|
||||
# Real intrinsics @ native 1600x1296 (from ROS calibration):
|
||||
# fx = 737.357 px, fy = 737.292 px, cx = 794.372 px, cy = 666.259 px
|
||||
# hFOV = 2*atan(W / 2fx) = 94.67 deg
|
||||
# vFOV = 2*atan(H / 2fy) = 82.65 deg
|
||||
#
|
||||
# Scaled to VAE-fixed 64x40 input (keep FOV identical):
|
||||
# fx_sim = 737.357 * 64/1600 = 29.49 px
|
||||
# fy_sim = 737.292 * 40/1296 = 22.76 px
|
||||
# cx_sim ~= 31.77 px (close to width/2)
|
||||
# cy_sim ~= 20.56 px (close to height/2)
|
||||
#
|
||||
# Noise model caveat: the upstream DepthNoise module simulates *stereo*
|
||||
# disparity-quantization noise (sigma_d proportional to d^2 / (fx * baseline)).
|
||||
# Odin1 depth comes from LiDAR-to-image alignment, where the real noise is
|
||||
# closer to a constant ~1-2 cm regardless of range. We keep the stereo-style
|
||||
# noise as a coarse proxy (set focal_length=fx_sim, baseline=0.05 m) but if
|
||||
# sim2real depth fidelity becomes a problem, swap in a constant-sigma noise
|
||||
# model later.
|
||||
ODIN1_CAMERA_CONFIG = CameraConfig(
|
||||
focal_length=29.49, # pixels (NOT mm; matches sim pinhole fx)
|
||||
baseline=0.05, # fudge: makes per-meter quantization noise sane
|
||||
min_depth=0.25,
|
||||
max_depth=10.0,
|
||||
resolution=(64, 40), # do not change; VAE encoder input is fixed
|
||||
depth_encoder_path=_get_encoder_path("vae_pretrain_new.pth"),
|
||||
)
|
||||
"""Configuration for Odin1 LiDAR-aligned depth (used with Go2)."""
|
||||
|
||||
|
||||
# Default camera configuration
|
||||
DEFAULT_CAMERA_CONFIG = ZEDX_CAMERA_CONFIG
|
||||
"""Default camera configuration (ZedX camera settings)."""
|
||||
|
||||
# Robot-to-camera mapping. Go2 now uses Odin1 (LiDAR-aligned depth). b2w and
|
||||
# aow_d keep the original ZedX profile so the upstream paper results stay
|
||||
# reproducible.
|
||||
ROBOT_CAMERA_CONFIGS = {
|
||||
"b2w": ZEDX_CAMERA_CONFIG,
|
||||
"aow_d": ZEDX_CAMERA_CONFIG,
|
||||
"go2": ODIN1_CAMERA_CONFIG,
|
||||
}
|
||||
"""Dictionary mapping robot names to their camera configurations."""
|
||||
|
||||
|
||||
def get_camera_config(robot_name: str, use_default_fallback: bool = False) -> CameraConfig:
|
||||
"""Get camera configuration for a specific robot.
|
||||
|
||||
Args:
|
||||
robot_name: Name of the robot (e.g., 'b2w', 'aow_d')
|
||||
use_default_fallback: If True, return DEFAULT_CAMERA_CONFIG when robot not found
|
||||
instead of raising an error (default: False)
|
||||
|
||||
Returns:
|
||||
Camera configuration for the specified robot
|
||||
|
||||
Raises:
|
||||
KeyError: If robot_name not found and use_default_fallback is False
|
||||
"""
|
||||
if robot_name in ROBOT_CAMERA_CONFIGS:
|
||||
return ROBOT_CAMERA_CONFIGS[robot_name]
|
||||
|
||||
if use_default_fallback:
|
||||
return DEFAULT_CAMERA_CONFIG
|
||||
|
||||
available_robots = ", ".join(sorted(ROBOT_CAMERA_CONFIGS.keys()))
|
||||
raise KeyError(
|
||||
f"Robot '{robot_name}' not found in camera configurations. "
|
||||
f"Available robots: {available_robots}"
|
||||
)
|
||||
+361
@@ -0,0 +1,361 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Depth noise generation and encoding for navigation tasks."""
|
||||
|
||||
import os
|
||||
from collections import OrderedDict
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torchvision.models import regnet_x_400mf
|
||||
from torchvision.ops import Conv2dNormActivation, FeaturePyramidNetwork
|
||||
|
||||
from isaaclab_assets import ISAACLAB_ASSETS_DATA_DIR
|
||||
|
||||
from .camera_config import CameraConfig
|
||||
|
||||
# Default camera configuration (ZedX Camera for b2w and aow d)
|
||||
DEFAULT_CAMERA_CONFIG = CameraConfig(
|
||||
focal_length=25.0,
|
||||
baseline=0.12,
|
||||
min_depth=0.25,
|
||||
max_depth=10.0,
|
||||
depth_encoder_path=os.path.join(
|
||||
str(ISAACLAB_ASSETS_DATA_DIR),
|
||||
"Policies",
|
||||
"RSL-ETHZ/AoW_d/depth_encoder",
|
||||
"vae_pretrain_fuse.pth"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class VAESampler(nn.Module):
|
||||
def __init__(self, input_dim, latent_dim):
|
||||
super(VAESampler, self).__init__()
|
||||
self.input_dim = input_dim
|
||||
self.latent_dim = latent_dim
|
||||
self.conv = Conv2dNormActivation(
|
||||
input_dim, latent_dim, kernel_size=3, stride=1, padding=1, bias=False
|
||||
)
|
||||
|
||||
# Convolutional Layers for 2D mean and logvar
|
||||
self.mean_layers = nn.Sequential(
|
||||
Conv2dNormActivation(
|
||||
latent_dim, latent_dim, kernel_size=3, stride=1, padding=1, bias=False
|
||||
),
|
||||
nn.Conv2d(latent_dim, latent_dim, kernel_size=1, stride=1, padding=0)
|
||||
)
|
||||
|
||||
self.logvar_layers = nn.Sequential(
|
||||
Conv2dNormActivation(
|
||||
latent_dim, latent_dim, kernel_size=3, stride=1, padding=1, bias=False
|
||||
),
|
||||
nn.Conv2d(latent_dim, latent_dim, kernel_size=1, stride=1, padding=0)
|
||||
)
|
||||
|
||||
def reparameterize(self, mu, logvar):
|
||||
std = torch.exp(0.5 * logvar)
|
||||
eps = torch.randn_like(std)
|
||||
return mu + eps * std
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.mean_layers(x)
|
||||
return x
|
||||
|
||||
|
||||
class EncoderFPN(nn.Module):
|
||||
def __init__(self, in_channel, out_channel):
|
||||
super(EncoderFPN, self).__init__()
|
||||
encoder = regnet_x_400mf(weights=None)
|
||||
# Remove classification head from the encoder
|
||||
encoder = nn.Sequential(*list(encoder.children())[:-2])
|
||||
# Modify the first layer to accept the number of channels in the input image
|
||||
encoder[0][0] = nn.Conv2d(in_channel, 32, kernel_size=3, stride=2, padding=1, bias=False)
|
||||
self.enc = encoder[0]
|
||||
self.enc_1 = encoder[1][:2]
|
||||
self.enc_2 = encoder[1][2]
|
||||
self.enc_3 = encoder[1][3]
|
||||
|
||||
# Feature Pyramid Network
|
||||
self.fpn = FeaturePyramidNetwork([64, 160, 400], out_channel)
|
||||
|
||||
def forward(self, x):
|
||||
raise NotImplementedError("Forward method must be implemented in the child class.")
|
||||
|
||||
|
||||
# Depth image encoder
|
||||
class DepthEncoder(EncoderFPN):
|
||||
def __init__(self, out_channel):
|
||||
super(DepthEncoder, self).__init__(1, out_channel)
|
||||
|
||||
def forward(self, x):
|
||||
# check if depth has channel dimension
|
||||
if x.dim() == 3:
|
||||
x = x.unsqueeze(1)
|
||||
|
||||
out = OrderedDict()
|
||||
x = self.enc(x)
|
||||
out['feat1'] = self.enc_1(x)
|
||||
out['feat2'] = self.enc_2(out['feat1'])
|
||||
out['feat3'] = self.enc_3(out['feat2'])
|
||||
|
||||
out = self.fpn(out)
|
||||
|
||||
return out['feat1']
|
||||
|
||||
|
||||
class VAEDecoder(nn.Module):
|
||||
def __init__(self, input_dim, out_dim):
|
||||
super(VAEDecoder, self).__init__()
|
||||
self.input_dim = input_dim
|
||||
self.conv = Conv2dNormActivation(
|
||||
input_dim, input_dim, kernel_size=3, stride=1, padding=1, bias=False
|
||||
)
|
||||
|
||||
self.decoder = nn.Sequential(
|
||||
# Layer 0
|
||||
nn.ConvTranspose2d(input_dim, input_dim, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
# Layer 1
|
||||
nn.ConvTranspose2d(input_dim, input_dim, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
# Layer 2
|
||||
nn.ConvTranspose2d(input_dim, input_dim, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
# Final Output Layer
|
||||
nn.Conv2d(input_dim, out_dim, kernel_size=3, stride=1, padding=1)
|
||||
) # Expand the spatial dimensions by a factor of 2**4=16
|
||||
|
||||
def forward(self, z):
|
||||
z = self.conv(z)
|
||||
img = self.decoder(z)
|
||||
return img
|
||||
|
||||
|
||||
class DepthDecoder(VAEDecoder):
|
||||
def __init__(self, input_dim):
|
||||
super(DepthDecoder, self).__init__(input_dim, 1)
|
||||
|
||||
|
||||
class VAENet(nn.Module):
|
||||
def __init__(self, latent_dim):
|
||||
super(VAENet, self).__init__()
|
||||
self.depth_encoder = DepthEncoder(latent_dim)
|
||||
|
||||
self.vae_sampler = VAESampler(latent_dim, latent_dim)
|
||||
|
||||
self.depth_decoder = DepthDecoder(latent_dim)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.depth_encoder(x)
|
||||
x = self.vae_sampler(x)
|
||||
return x
|
||||
|
||||
def decode(self, z):
|
||||
return self.depth_decoder(z)
|
||||
|
||||
|
||||
class DepthNoise(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
focal_length,
|
||||
baseline,
|
||||
min_depth,
|
||||
max_depth,
|
||||
filter_size=3,
|
||||
inlier_thred_range=(0.01, 0.05),
|
||||
prob_range=(0.4, 0.6),
|
||||
invalid_disp=1e7
|
||||
):
|
||||
"""A Simply PyTorch module to add realistic noise to depth images.
|
||||
|
||||
Args:
|
||||
focal_length (float): Focal length of the camera (in pixels).
|
||||
baseline (float): Baseline distance between stereo cameras (in meters).
|
||||
min_depth (float): Minimum depth value after clamping.
|
||||
max_depth (float): Maximum depth value after clamping.
|
||||
filter_size (int): Kernel size for local mean disparity computation.
|
||||
inlier_thred_range (tuple): Threshold range for normalized disparity differences.
|
||||
prob_range (tuple): Probability range for matching pixels.
|
||||
invalid_disp (float): Invalid disparity
|
||||
"""
|
||||
super().__init__()
|
||||
self.focal_length = focal_length
|
||||
self.baseline = baseline
|
||||
self.min_depth = min_depth
|
||||
self.max_depth = max_depth
|
||||
self.invalid_disp = invalid_disp
|
||||
self.inlier_thred_range = inlier_thred_range
|
||||
self.prob_range = prob_range
|
||||
self.filter_size = filter_size
|
||||
|
||||
weights, substitutes = self._compute_weights(filter_size)
|
||||
self.register_buffer('weights', weights.view(1, 1, filter_size, filter_size))
|
||||
self.register_buffer('substitutes', substitutes.view(1, 1, filter_size, filter_size))
|
||||
|
||||
def _compute_weights(self, filter_size):
|
||||
"""Compute weights and substitutes for disparity filtering.
|
||||
|
||||
Args:
|
||||
filter_size (int): Kernel size for local mean disparity computation.
|
||||
"""
|
||||
center = filter_size // 2
|
||||
idx = torch.arange(filter_size) - center
|
||||
x_filter, y_filter = torch.meshgrid(idx, idx, indexing='ij')
|
||||
sqr_radius = x_filter ** 2 + y_filter ** 2
|
||||
sqrt_radius = torch.sqrt(sqr_radius)
|
||||
weights = 1 / torch.where(sqr_radius == 0, torch.ones_like(sqrt_radius), sqrt_radius)
|
||||
weights = weights / weights.sum()
|
||||
fill_weights = 1 / (1 + sqrt_radius)
|
||||
fill_weights = torch.where(sqr_radius > filter_size, -1.0, fill_weights)
|
||||
substitutes = (fill_weights > 0).float()
|
||||
|
||||
return weights, substitutes
|
||||
|
||||
def filter_disparity(self, disparity):
|
||||
"""Filter the disparity map using local mean disparity.
|
||||
|
||||
Args:
|
||||
disparity (torch.Tensor): Input disparity map tensor of shape (B, C, H, W).
|
||||
"""
|
||||
B, _, H, W = disparity.shape
|
||||
device = disparity.device
|
||||
center = self.filter_size // 2
|
||||
|
||||
output_disparity = torch.full_like(disparity, self.invalid_disp)
|
||||
|
||||
prob = (
|
||||
torch.rand(B, 1, 1, 1, device=device)
|
||||
* (self.prob_range[1] - self.prob_range[0])
|
||||
+ self.prob_range[0]
|
||||
)
|
||||
random_mask = (torch.rand(B, 1, H, W, device=device) < prob)
|
||||
|
||||
# Compute mean disparity
|
||||
weighted_disparity = F.conv2d(disparity, self.weights, padding=center)
|
||||
|
||||
# Compute differences
|
||||
differences = torch.abs(disparity - weighted_disparity)
|
||||
|
||||
# Normalize differences based on current image statistics
|
||||
differences_flat = differences.view(B, -1)
|
||||
mean_diff = torch.mean(differences_flat, dim=1, keepdim=True)
|
||||
std_diff = torch.std(differences_flat, dim=1, keepdim=True) + 1e-6
|
||||
|
||||
normalized_differences_flat = (differences_flat - mean_diff) / std_diff
|
||||
normalized_differences = normalized_differences_flat.view_as(differences)
|
||||
|
||||
# Use parameter-based threshold on normalized differences
|
||||
threshold = (
|
||||
torch.rand(B, 1, 1, 1, device=device)
|
||||
* (self.inlier_thred_range[1] - self.inlier_thred_range[0])
|
||||
+ self.inlier_thred_range[0]
|
||||
)
|
||||
update_mask = (normalized_differences < threshold) & random_mask
|
||||
|
||||
# Compute output value: round with 1/32 precision
|
||||
disparity = torch.round(disparity * 32.0) / 32.0
|
||||
|
||||
# Update output disparity
|
||||
output_disparity = torch.where(update_mask, disparity, output_disparity)
|
||||
|
||||
# Apply substitutes to fill neighboring pixels
|
||||
filled_values = F.conv2d(update_mask.float() * disparity, self.substitutes, padding=center)
|
||||
counts = F.conv2d(update_mask.float(), self.substitutes, padding=center) + 1e-9
|
||||
average_filled_values = filled_values / counts
|
||||
output_disparity = torch.where(counts >= 1, average_filled_values, output_disparity)
|
||||
|
||||
return output_disparity
|
||||
|
||||
def forward(self, depth, add_noise: bool) -> torch.Tensor:
|
||||
# correct input shape
|
||||
if len(depth.shape) == 3:
|
||||
depth = depth.unsqueeze(1) # add channel dimension
|
||||
|
||||
# check dimension (B, 1, H, W)
|
||||
assert depth.shape[1] == 1, "Input depth tensor must have shape (B, 1, H, W)."
|
||||
assert len(depth.shape) == 4, "Input depth tensor must have shape (B, 1, H, W)."
|
||||
|
||||
if add_noise:
|
||||
# Clamp the depth values
|
||||
depth = torch.clamp(depth, min=1. / self.invalid_disp)
|
||||
|
||||
# Step 1: Convert depth to disparity
|
||||
disparity = self.focal_length * self.baseline / depth
|
||||
|
||||
# Step 2: Filter the disparity map
|
||||
filtered_disparity = self.filter_disparity(disparity)
|
||||
|
||||
# Step 3: Recompute depth from disparity
|
||||
depth = self.focal_length * self.baseline / filtered_disparity
|
||||
|
||||
# Step 4: Clamp the depth values
|
||||
depth[depth < self.min_depth] = 0.0
|
||||
|
||||
# Step 5: Set invalid depth values to 0.0
|
||||
depth[depth > self.max_depth] = 0.0
|
||||
|
||||
return depth
|
||||
|
||||
|
||||
class DepthNoiseEncoder(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
feature_dim,
|
||||
camera_config: CameraConfig = None
|
||||
):
|
||||
"""A Simply PyTorch module to add realistic noise to depth images.
|
||||
|
||||
Args:
|
||||
feature_dim (int): Number of output channels from the encoder.
|
||||
camera_config (CameraConfig, optional): Camera configuration parameters.
|
||||
If None, uses DEFAULT_CAMERA_CONFIG.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Use provided config or default
|
||||
if camera_config is None:
|
||||
camera_config = DEFAULT_CAMERA_CONFIG
|
||||
|
||||
self.camera_config = camera_config
|
||||
self.depth_noise = DepthNoise(
|
||||
focal_length=camera_config.focal_length,
|
||||
baseline=camera_config.baseline,
|
||||
min_depth=camera_config.min_depth,
|
||||
max_depth=camera_config.max_depth
|
||||
)
|
||||
|
||||
self.encoder = VAENet(feature_dim)
|
||||
|
||||
try:
|
||||
self.encoder.load_state_dict(
|
||||
torch.load(camera_config.depth_encoder_path, weights_only=True),
|
||||
strict=True
|
||||
)
|
||||
print(
|
||||
'\033[92m'
|
||||
+ f'Successfully loaded pre-trained weights from {camera_config.depth_encoder_path}'
|
||||
+ '\033[0m'
|
||||
)
|
||||
except Exception as e:
|
||||
print('\033[91m' + f'Failed to load pre-trained weights: {e}' + '\033[0m')
|
||||
|
||||
def forward(
|
||||
self, depth: torch.Tensor, add_noise: bool = True
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
noisy_depth = self.depth_noise(depth, add_noise)
|
||||
encoded_depth = self.encoder(noisy_depth) # (B, 1, H, W) -> (B, C, H//8, W//8)
|
||||
return encoded_depth, noisy_depth
|
||||
|
||||
@torch.jit.export
|
||||
def decode(self, z: torch.Tensor) -> torch.Tensor:
|
||||
return self.encoder.decode(z)
|
||||
@@ -0,0 +1,229 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Custom event functions for navigation tasks.
|
||||
|
||||
These functions are ported from the internal Isaac Lab fork and provide
|
||||
domain randomization and environment events specific to navigation tasks.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from isaaclab.envs import ManagerBasedEnv
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
from isaaclab.utils import math as math_utils
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.sensors import RayCasterCamera
|
||||
from isaaclab_nav_task.navigation.mdp import PerceptiveNavigationSE2Action
|
||||
|
||||
|
||||
def randomize_camera_height(
|
||||
env: ManagerBasedEnv,
|
||||
env_ids: torch.Tensor,
|
||||
height_range: float,
|
||||
roll_angle_range: float,
|
||||
pitch_angle_range: float,
|
||||
yaw_angle_range: float,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("raycast_camera"),
|
||||
):
|
||||
"""Randomize the camera's height and orientation around the original values.
|
||||
|
||||
The function samples around the original camera offset and orientation:
|
||||
- height: original height +- height_range (meters)
|
||||
- roll: original roll +- roll_angle_range (degrees)
|
||||
- pitch: original pitch +- pitch_angle_range (degrees)
|
||||
- yaw: original yaw +- yaw_angle_range (degrees)
|
||||
|
||||
Args:
|
||||
env: The environment manager.
|
||||
env_ids: The environment ids to randomize.
|
||||
height_range: Range around original height (meters).
|
||||
roll_angle_range: Range around original roll angle (degrees).
|
||||
pitch_angle_range: Range around original pitch angle (degrees).
|
||||
yaw_angle_range: Range around original yaw angle (degrees).
|
||||
asset_cfg: Scene entity configuration for the camera.
|
||||
"""
|
||||
camera: RayCasterCamera = env.scene[asset_cfg.name]
|
||||
|
||||
# Get original camera position and orientation
|
||||
original_camera_pos = camera.cfg.offset.pos
|
||||
original_camera_quat = camera.cfg.offset.rot
|
||||
|
||||
# Sample position around original value
|
||||
num_envs = env_ids.numel()
|
||||
pos_delta = torch.zeros((num_envs, 3), device=env.device)
|
||||
pos_delta[:, 0] = torch.randn(num_envs, device=env.device) * (height_range / 2.0)
|
||||
pos_delta[:, 1] = torch.randn(num_envs, device=env.device) * (height_range / 2.0)
|
||||
pos_delta[:, 0] = torch.clamp(pos_delta[:, 0], -height_range, height_range)
|
||||
pos_delta[:, 1] = torch.clamp(pos_delta[:, 1], -height_range, height_range)
|
||||
pos_delta[:, 2] = torch.rand(num_envs, device=env.device) * 2.0 * height_range - height_range
|
||||
camera._offset_pos[env_ids] = torch.tensor(original_camera_pos, device=env.device) + pos_delta
|
||||
|
||||
# Convert original quaternion to Euler angles (XYZ convention)
|
||||
quat_tensor = torch.tensor(original_camera_quat, device=env.device).unsqueeze(0)
|
||||
original_euler = math_utils.euler_xyz_from_quat(quat_tensor)
|
||||
original_roll_deg = original_euler[0] * 180.0 / torch.pi
|
||||
original_pitch_deg = original_euler[1] * 180.0 / torch.pi
|
||||
original_yaw_deg = original_euler[2] * 180.0 / torch.pi
|
||||
|
||||
# Sample angles around original values
|
||||
roll_delta_deg = torch.randn(num_envs, device=env.device) * (roll_angle_range / 2.0)
|
||||
roll_delta_deg = torch.clamp(roll_delta_deg, -roll_angle_range, roll_angle_range)
|
||||
pitch_delta_deg = torch.rand(num_envs, device=env.device) * 2.0 * pitch_angle_range - pitch_angle_range
|
||||
yaw_delta_deg = torch.randn(num_envs, device=env.device) * (yaw_angle_range / 2.0)
|
||||
yaw_delta_deg = torch.clamp(yaw_delta_deg, -yaw_angle_range, yaw_angle_range)
|
||||
|
||||
# Add deltas to original angles
|
||||
sampled_roll_deg = original_roll_deg + roll_delta_deg
|
||||
sampled_pitch_deg = original_pitch_deg + pitch_delta_deg
|
||||
sampled_yaw_deg = original_yaw_deg + yaw_delta_deg
|
||||
|
||||
# Convert to radians and build quaternion (XYZ convention)
|
||||
sampled_roll_rad = sampled_roll_deg * torch.pi / 180.0
|
||||
sampled_pitch_rad = sampled_pitch_deg * torch.pi / 180.0
|
||||
sampled_yaw_rad = sampled_yaw_deg * torch.pi / 180.0
|
||||
|
||||
camera._offset_quat[env_ids] = math_utils.quat_from_euler_xyz(sampled_roll_rad, sampled_pitch_rad, sampled_yaw_rad)
|
||||
|
||||
|
||||
def randomize_action_scale(
|
||||
env: ManagerBasedEnv,
|
||||
env_ids: torch.Tensor,
|
||||
scale_range_x: tuple[float, float],
|
||||
scale_range_y: tuple[float, float],
|
||||
scale_range_theta: tuple[float, float],
|
||||
scale_range_xb: float,
|
||||
scale_range_yb: float,
|
||||
scale_range_thetab: float,
|
||||
action_term: str,
|
||||
):
|
||||
"""Change the action scale in the scene.
|
||||
|
||||
Args:
|
||||
env: The environment manager.
|
||||
env_ids: The environment ids to change the action scale in.
|
||||
scale_range_x: The scale range for x.
|
||||
scale_range_y: The scale range for y.
|
||||
scale_range_theta: The scale range for theta.
|
||||
scale_range_xb: The bias range for x.
|
||||
scale_range_yb: The bias range for y.
|
||||
scale_range_thetab: The bias range for theta.
|
||||
action_term: Name of the action term.
|
||||
"""
|
||||
action_term_obj: PerceptiveNavigationSE2Action = env.action_manager._terms[action_term]
|
||||
|
||||
action_term_obj._policy_scaling[env_ids] = torch.tensor(
|
||||
action_term_obj.cfg.policy_scaling, device=action_term_obj.device
|
||||
).expand(len(env_ids), -1)
|
||||
|
||||
random_scales_x = (
|
||||
torch.rand(len(env_ids), 1, device=env.device) * (scale_range_x[1] - scale_range_x[0]) + scale_range_x[0]
|
||||
)
|
||||
random_scales_y = (
|
||||
torch.rand(len(env_ids), 1, device=env.device) * (scale_range_y[1] - scale_range_y[0]) + scale_range_y[0]
|
||||
)
|
||||
random_scales_theta = (
|
||||
torch.rand(len(env_ids), 1, device=env.device) * (scale_range_theta[1] - scale_range_theta[0])
|
||||
+ scale_range_theta[0]
|
||||
)
|
||||
random_scales_xb = (torch.rand(len(env_ids), 1, device=env.device) * 2 - 1) * scale_range_xb
|
||||
random_scales_yb = (torch.rand(len(env_ids), 1, device=env.device) * 2 - 1) * scale_range_yb
|
||||
random_scales_thetab = (torch.rand(len(env_ids), 1, device=env.device) * 2 - 1) * scale_range_thetab
|
||||
|
||||
random_scales = torch.cat((random_scales_x, random_scales_y, random_scales_theta), dim=-1)
|
||||
random_scales_b = torch.cat((random_scales_xb, random_scales_yb, random_scales_thetab), dim=-1)
|
||||
|
||||
action_term_obj._policy_scaling[env_ids] *= random_scales
|
||||
action_term_obj._policy_bias[env_ids, :] = random_scales_b
|
||||
|
||||
|
||||
def reset_and_randomize_delay_buffer(
|
||||
env: ManagerBasedEnv,
|
||||
env_ids: torch.Tensor,
|
||||
):
|
||||
"""Reset and randomize delay buffers.
|
||||
|
||||
This function resets the delay buffers and randomizes the time lags for
|
||||
all observation types. Requires env.delay_manager to exist.
|
||||
|
||||
Args:
|
||||
env: The environment manager (must have delay_manager attribute).
|
||||
env_ids: The environment ids to reset.
|
||||
"""
|
||||
env.delay_manager.reset(env_ids)
|
||||
env.delay_manager.randomize_lags(env_ids)
|
||||
|
||||
|
||||
def randomize_low_pass_filter_alpha(
|
||||
env: ManagerBasedEnv,
|
||||
env_ids: torch.Tensor,
|
||||
alpha_range: tuple[float, float],
|
||||
action_term: str,
|
||||
per_dimension: bool = True,
|
||||
alpha_range_vx: tuple[float, float] | None = None,
|
||||
alpha_range_vy: tuple[float, float] | None = None,
|
||||
alpha_range_omega: tuple[float, float] | None = None,
|
||||
):
|
||||
"""Randomize the low-pass filter alpha parameter for velocity command smoothing.
|
||||
|
||||
Args:
|
||||
env: The environment manager.
|
||||
env_ids: The environment ids to randomize the alpha parameter for.
|
||||
alpha_range: Tuple of (min_alpha, max_alpha) values.
|
||||
action_term: Name of the action term to modify.
|
||||
per_dimension: Whether to use different alpha for each dimension.
|
||||
alpha_range_vx: Optional range for vx dimension.
|
||||
alpha_range_vy: Optional range for vy dimension.
|
||||
alpha_range_omega: Optional range for omega dimension.
|
||||
"""
|
||||
action_term_obj: PerceptiveNavigationSE2Action = env.action_manager._terms[action_term]
|
||||
|
||||
if per_dimension:
|
||||
# Use dimension-specific ranges or fall back to global range
|
||||
vx_range = alpha_range_vx if alpha_range_vx is not None else alpha_range
|
||||
vy_range = alpha_range_vy if alpha_range_vy is not None else alpha_range
|
||||
omega_range = alpha_range_omega if alpha_range_omega is not None else alpha_range
|
||||
|
||||
alpha_vx = torch.rand(len(env_ids), device=env.device) * (vx_range[1] - vx_range[0]) + vx_range[0]
|
||||
alpha_vy = torch.rand(len(env_ids), device=env.device) * (vy_range[1] - vy_range[0]) + vy_range[0]
|
||||
alpha_omega = torch.rand(len(env_ids), device=env.device) * (omega_range[1] - omega_range[0]) + omega_range[0]
|
||||
|
||||
if hasattr(action_term_obj, "_per_env_per_dim_low_pass_alpha"):
|
||||
action_term_obj._per_env_per_dim_low_pass_alpha[env_ids, 0] = alpha_vx
|
||||
action_term_obj._per_env_per_dim_low_pass_alpha[env_ids, 1] = alpha_vy
|
||||
action_term_obj._per_env_per_dim_low_pass_alpha[env_ids, 2] = alpha_omega
|
||||
else:
|
||||
random_alpha = (
|
||||
torch.rand(len(env_ids), device=env.device) * (alpha_range[1] - alpha_range[0]) + alpha_range[0]
|
||||
)
|
||||
if hasattr(action_term_obj, "_per_env_per_dim_low_pass_alpha"):
|
||||
action_term_obj._per_env_per_dim_low_pass_alpha[env_ids] = random_alpha.unsqueeze(-1).expand(-1, 3)
|
||||
|
||||
|
||||
def disable_backward_penalty_after_steps(
|
||||
env: ManagerBasedEnv,
|
||||
env_ids: torch.Tensor,
|
||||
disable_after_steps: int,
|
||||
action_term: str,
|
||||
):
|
||||
"""Disable backward movement penalty after a certain number of steps.
|
||||
|
||||
Args:
|
||||
env: The environment manager.
|
||||
env_ids: The environment ids.
|
||||
disable_after_steps: Number of steps after which to disable the penalty.
|
||||
action_term: Name of the action term.
|
||||
"""
|
||||
action_term_obj: PerceptiveNavigationSE2Action = env.action_manager._terms[action_term]
|
||||
|
||||
if hasattr(action_term_obj, "disable_backward_penalty"):
|
||||
# Check which environments have exceeded the step threshold
|
||||
if hasattr(env, "episode_length_buf"):
|
||||
exceeded_steps = env.episode_length_buf[env_ids] >= disable_after_steps
|
||||
action_term_obj.disable_backward_penalty[env_ids[exceeded_steps]] = True
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Height scan encoder utilities for navigation tasks."""
|
||||
|
||||
from .heightscan_encoder import HeightScanFeatEncoder
|
||||
|
||||
__all__ = ["HeightScanFeatEncoder"]
|
||||
+192
@@ -0,0 +1,192 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from typing import Tuple
|
||||
import torch.nn.functional as F
|
||||
from collections import OrderedDict
|
||||
from torchvision.ops import FeaturePyramidNetwork
|
||||
from torchvision.models import regnet_x_400mf
|
||||
from torchvision.ops import Conv2dNormActivation
|
||||
|
||||
# Local path to the encoder weights (relative to this file's directory)
|
||||
_ASSETS_DIR = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "assets", "data"))
|
||||
HEIGHTSCAN_ENCODER_PATH = os.path.join(_ASSETS_DIR, "Policies", "heightscan_encoder", "vae_heightscan3.pth")
|
||||
|
||||
class VAESampler(nn.Module):
|
||||
def __init__(self, input_dim, latent_dim):
|
||||
super(VAESampler, self).__init__()
|
||||
self.input_dim = input_dim
|
||||
self.latent_dim = latent_dim
|
||||
self.conv = Conv2dNormActivation(input_dim, latent_dim, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
|
||||
# Convolutional Layers for 2D mean and logvar
|
||||
self.mean_layers = nn.Sequential(
|
||||
Conv2dNormActivation(latent_dim, latent_dim, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
nn.Conv2d(latent_dim, latent_dim, kernel_size=1, stride=1, padding=0)
|
||||
)
|
||||
|
||||
self.logvar_layers = nn.Sequential(
|
||||
Conv2dNormActivation(latent_dim, latent_dim, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
nn.Conv2d(latent_dim, latent_dim, kernel_size=1, stride=1, padding=0)
|
||||
)
|
||||
|
||||
def reparameterize(self, mu, logvar):
|
||||
std = torch.exp(0.5 * logvar)
|
||||
eps = torch.randn_like(std)
|
||||
return mu + eps * std
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.mean_layers(x)
|
||||
return x
|
||||
|
||||
class EncoderFPN(nn.Module):
|
||||
def __init__(self, in_channel, out_channel):
|
||||
super(EncoderFPN, self).__init__()
|
||||
encoder = regnet_x_400mf(weights=None)
|
||||
# Remove classification head from the encoder
|
||||
encoder = nn.Sequential(*list(encoder.children())[:-2])
|
||||
# Modify the first layer to accept the number of channels in the input image
|
||||
encoder[0][0] = nn.Conv2d(in_channel, 32, kernel_size=3, stride=2, padding=1, bias=False)
|
||||
self.enc = encoder[0]
|
||||
self.enc_1 = encoder[1][:2]
|
||||
self.enc_2 = encoder[1][2]
|
||||
self.enc_3 = encoder[1][3]
|
||||
|
||||
# Feature Pyramid Network
|
||||
self.fpn = FeaturePyramidNetwork([64, 160, 400], out_channel)
|
||||
|
||||
def forward(self, x):
|
||||
raise NotImplementedError("Forward method must be implemented in the child class.")
|
||||
|
||||
# Depth image encoder
|
||||
class HeightScanEncoder(EncoderFPN):
|
||||
def __init__(self, out_channel):
|
||||
super(HeightScanEncoder, self).__init__(1, out_channel)
|
||||
|
||||
def forward(self, x):
|
||||
# check if depth has channel dimension
|
||||
if x.dim() == 3:
|
||||
x = x.unsqueeze(1)
|
||||
|
||||
out = OrderedDict()
|
||||
x = self.enc(x)
|
||||
out['feat1'] = self.enc_1(x)
|
||||
out['feat2'] = self.enc_2(out['feat1'])
|
||||
out['feat3'] = self.enc_3(out['feat2'])
|
||||
|
||||
out = self.fpn(out)
|
||||
|
||||
return out['feat1']
|
||||
|
||||
class VAEDecoder(nn.Module):
|
||||
def __init__(self, input_dim, out_dim):
|
||||
super(VAEDecoder, self).__init__()
|
||||
self.input_dim = input_dim
|
||||
self.conv = Conv2dNormActivation(input_dim, input_dim, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
|
||||
self.decoder = nn.Sequential(
|
||||
# Layer 0
|
||||
nn.ConvTranspose2d(input_dim, input_dim, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
# Layer 1
|
||||
nn.ConvTranspose2d(input_dim, input_dim, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
# Layer 2
|
||||
nn.ConvTranspose2d(input_dim, input_dim, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
# Final Output Layer
|
||||
nn.Conv2d(input_dim, out_dim, kernel_size=3, stride=1, padding=1)
|
||||
) # Expand the spatial dimensions by a factor of 2**4=16
|
||||
|
||||
def forward(self, z):
|
||||
z = self.conv(z)
|
||||
img = self.decoder(z)
|
||||
return img
|
||||
|
||||
class HeightScanDecoder(VAEDecoder):
|
||||
def __init__(self, input_dim):
|
||||
super(HeightScanDecoder, self).__init__(input_dim, 1)
|
||||
|
||||
class VAENet(nn.Module):
|
||||
def __init__(self, latent_dim):
|
||||
super(VAENet, self).__init__()
|
||||
self.depth_encoder = HeightScanEncoder(latent_dim)
|
||||
|
||||
self.vae_sampler = VAESampler(latent_dim, latent_dim)
|
||||
|
||||
self.depth_decoder = HeightScanDecoder(latent_dim)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.depth_encoder(x)
|
||||
x = self.vae_sampler(x)
|
||||
return x
|
||||
|
||||
def decode(self, z):
|
||||
return self.depth_decoder(z)
|
||||
|
||||
|
||||
class HeightScanFeatEncoder(torch.nn.Module):
|
||||
def __init__(self,
|
||||
feature_dim,
|
||||
):
|
||||
"""
|
||||
A Simply PyTorch module to add realistic noise to depth images.
|
||||
|
||||
Args:
|
||||
feature_dim (int): Number of output channels from the encoder.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.encoder = VAENet(feature_dim)
|
||||
|
||||
try:
|
||||
self.encoder.load_state_dict(torch.load(HEIGHTSCAN_ENCODER_PATH, weights_only=True), strict=True)
|
||||
print('\033[92m' + f'Successfully loaded pre-trained weights from {HEIGHTSCAN_ENCODER_PATH}' + '\033[0m')
|
||||
except Exception as e:
|
||||
print('\033[91m' + f'Failed to load pre-trained weights: {e}' + '\033[0m')
|
||||
|
||||
def forward(self, scan: torch.Tensor) -> torch.Tensor:
|
||||
encoded_scan = self.encoder(scan) # (B, 1, H, W) -> (B, C, H//8, W//8)
|
||||
return encoded_scan
|
||||
|
||||
@torch.jit.export
|
||||
def decode(self, z: torch.Tensor) -> torch.Tensor:
|
||||
return self.encoder.decode(z)
|
||||
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
# parameters for the depth noise
|
||||
feature_dim = 64
|
||||
|
||||
# Initialize the DepthNoise class
|
||||
heightscan_feat_encoder = HeightScanFeatEncoder(feature_dim=feature_dim)
|
||||
|
||||
# compile the depth_noise module
|
||||
compiled_scan_encoder = torch.jit.script(heightscan_feat_encoder)
|
||||
|
||||
# random height scan image
|
||||
scan = np.random.rand(51, 51) * 5
|
||||
|
||||
# convert depth to tensor
|
||||
scan = torch.tensor(scan).unsqueeze(0).unsqueeze(0).float()
|
||||
clipped_scan = torch.clamp(scan, min=-5, max=5)
|
||||
|
||||
# Add noise to the depth image
|
||||
noisy_scan = compiled_scan_encoder(clipped_scan)
|
||||
|
||||
noisy_scan = noisy_scan.numpy().squeeze().squeeze()
|
||||
# check the min and max values of the noisy depth
|
||||
print("Min depth: ", noisy_scan.min())
|
||||
print("Max depth: ", noisy_scan.max())
|
||||
@@ -0,0 +1,50 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Custom math utilities for navigation tasks."""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def vec_to_quat(v: torch.Tensor) -> torch.Tensor:
|
||||
"""Converts a unit vector to a quaternion representing the rotation from the X-axis to the vector v.
|
||||
|
||||
Args:
|
||||
v (torch.Tensor): A tensor of shape (..., 3) representing the unit vectors.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: A tensor of shape (..., 4) representing the quaternions (w, x, y, z).
|
||||
"""
|
||||
# Ensure v is a unit vector
|
||||
norm_v = v / torch.clamp(torch.norm(v, dim=-1, keepdim=True), min=1e-8)
|
||||
|
||||
# Reference vector (X-axis)
|
||||
a = torch.tensor([1.0, 0.0, 0.0], device=norm_v.device, dtype=norm_v.dtype).expand_as(norm_v)
|
||||
|
||||
# Compute dot and cross products
|
||||
dot = torch.sum(a * norm_v, dim=-1, keepdim=True)
|
||||
w = torch.cross(a, norm_v, dim=-1)
|
||||
|
||||
# Compute quaternion components
|
||||
s = torch.sqrt((1.0 + dot) * 0.5).clamp(min=1e-8) # Avoid division by zero
|
||||
q = torch.cat([s, w / (2.0 * s)], dim=-1)
|
||||
|
||||
# Handle edge case when dot == -1 (vectors are opposite)
|
||||
mask = (dot < -0.999999).squeeze(-1)
|
||||
if mask.any():
|
||||
# Choose an arbitrary orthogonal vector
|
||||
orthogonal = torch.cross(
|
||||
a[mask],
|
||||
torch.tensor([0.0, 1.0, 0.0], device=norm_v.device, dtype=norm_v.dtype).expand_as(a[mask]),
|
||||
dim=-1
|
||||
)
|
||||
orthogonal = orthogonal / torch.clamp(torch.norm(orthogonal, dim=-1, keepdim=True), min=1e-8)
|
||||
q[mask] = torch.cat([torch.zeros_like(s[mask]), orthogonal], dim=-1)
|
||||
|
||||
# Normalize the quaternion
|
||||
q = q / torch.clamp(torch.norm(q, dim=-1, keepdim=True), min=1e-8)
|
||||
return q
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Navigation-specific MDP components including goal commands and SE2 actions."""
|
||||
|
||||
from .goal_commands import * # noqa: F401, F403
|
||||
from .goal_commands_cfg import * # noqa: F401, F403
|
||||
from .actions import * # noqa: F401, F403
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
|
||||
from .navigation_se2_actions import PerceptiveNavigationSE2Action
|
||||
from .navigation_se2_actions_cfg import PerceptiveNavigationSE2ActionCfg
|
||||
|
||||
__all__ = ["PerceptiveNavigationSE2Action", "PerceptiveNavigationSE2ActionCfg"]
|
||||
+249
@@ -0,0 +1,249 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Navigation SE2 action term for hierarchical control."""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from isaaclab.envs import ManagerBasedRLEnv
|
||||
from isaaclab.managers.action_manager import ActionTerm
|
||||
from isaaclab.utils.assets import check_file_path, read_file
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .navigation_se2_actions_cfg import PerceptiveNavigationSE2ActionCfg
|
||||
|
||||
|
||||
class PerceptiveNavigationSE2Action(ActionTerm):
|
||||
"""Actions to navigate a robot using hierarchical control with a pre-trained locomotion policy."""
|
||||
|
||||
cfg: PerceptiveNavigationSE2ActionCfg
|
||||
_env: ManagerBasedRLEnv
|
||||
|
||||
def __init__(self, cfg: PerceptiveNavigationSE2ActionCfg, env: ManagerBasedRLEnv):
|
||||
super().__init__(cfg, env)
|
||||
# Check if policy file exists
|
||||
if not check_file_path(cfg.low_level_policy_file):
|
||||
raise FileNotFoundError(f"Policy file '{cfg.low_level_policy_file}' does not exist.")
|
||||
# Load pre-trained locomotion policy
|
||||
file_bytes = read_file(self.cfg.low_level_policy_file)
|
||||
self.low_level_policy = torch.jit.load(file_bytes, map_location=self.device)
|
||||
self.low_level_policy.eval()
|
||||
|
||||
# prepare joint position actions
|
||||
self.low_level_position_action_term: ActionTerm = self.cfg.low_level_position_action.class_type(cfg.low_level_position_action, env)
|
||||
# Velocity branch is optional (None for legs-only robots like Go2).
|
||||
if self.cfg.low_level_velocity_action is not None:
|
||||
self.low_level_velocity_action_term: ActionTerm | None = self.cfg.low_level_velocity_action.class_type(cfg.low_level_velocity_action, env)
|
||||
else:
|
||||
self.low_level_velocity_action_term = None
|
||||
|
||||
# prepare buffers
|
||||
self._action_dim = 3 # [vx, vy, omega]
|
||||
|
||||
# set up buffers
|
||||
self._init_buffers()
|
||||
|
||||
# Low-pass filter state for velocity commands
|
||||
self._prev_filtered_velocity_commands = torch.zeros((self.num_envs, self._action_dim), device=self.device)
|
||||
self._low_pass_alpha = self.cfg.low_pass_filter_alpha if hasattr(self.cfg, 'low_pass_filter_alpha') else 0.8
|
||||
self._enable_low_pass_filter = self.cfg.enable_low_pass_filter if hasattr(self.cfg, 'enable_low_pass_filter') else True
|
||||
# Per-environment per-dimension alpha values (initialized to default, can be randomized per episode)
|
||||
# Shape: [num_envs, action_dim] where action_dim = 3 (vx, vy, omega)
|
||||
self._per_env_per_dim_low_pass_alpha = torch.full((self.num_envs, self._action_dim), self._low_pass_alpha, device=self.device)
|
||||
|
||||
|
||||
"""
|
||||
Properties.
|
||||
"""
|
||||
|
||||
@property
|
||||
def action_dim(self) -> int:
|
||||
return self._action_dim
|
||||
|
||||
@property
|
||||
def raw_actions(self) -> torch.Tensor:
|
||||
return self._raw_navigation_velocity_actions
|
||||
|
||||
@property
|
||||
def processed_actions(self) -> torch.Tensor:
|
||||
return self._processed_navigation_velocity_actions
|
||||
|
||||
@property
|
||||
def filtered_velocity_commands(self) -> torch.Tensor:
|
||||
"""Get the current filtered (smoothed) velocity commands."""
|
||||
return self._prev_filtered_velocity_commands
|
||||
|
||||
@property
|
||||
def low_pass_alpha_values(self) -> torch.Tensor:
|
||||
"""Get the current per-environment per-dimension low-pass filter alpha values.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Alpha values with shape [num_envs, action_dim] where:
|
||||
- action_dim = 3 for [vx, vy, omega]
|
||||
- Each environment can have different alpha values for each command dimension
|
||||
"""
|
||||
return self._per_env_per_dim_low_pass_alpha
|
||||
|
||||
@property
|
||||
def low_level_actions(self) -> torch.Tensor:
|
||||
if self.low_level_velocity_action_term is None:
|
||||
return self._low_level_position_actions
|
||||
return torch.cat((self._low_level_position_actions, self._low_level_velocity_actions), dim=1)
|
||||
|
||||
@property
|
||||
def low_level_position_actions(self) -> torch.Tensor:
|
||||
return self._low_level_position_actions
|
||||
|
||||
@property
|
||||
def prev_low_level_position_actions(self) -> torch.Tensor:
|
||||
return self._prev_low_level_position_actions
|
||||
|
||||
@property
|
||||
def low_level_velocity_actions(self) -> torch.Tensor:
|
||||
return self._low_level_velocity_actions
|
||||
|
||||
@property
|
||||
def prev_low_level_velocity_actions(self) -> torch.Tensor:
|
||||
return self._prev_low_level_velocity_actions
|
||||
|
||||
"""
|
||||
Operations.
|
||||
"""
|
||||
|
||||
def apply_low_pass_filter(self, velocity_commands: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply low-pass filter to velocity commands for smoother locomotion.
|
||||
|
||||
The low-pass filter implements exponential smoothing:
|
||||
filtered_cmd(t) = alpha * filtered_cmd(t-1) + (1 - alpha) * new_cmd(t)
|
||||
|
||||
Where alpha is the smoothing factor:
|
||||
- alpha = 0.0: No smoothing (pass through)
|
||||
- alpha = 1.0: Maximum smoothing (no change)
|
||||
- alpha = 0.8: Good balance for locomotion (default)
|
||||
|
||||
This implementation supports per-environment per-dimension alpha values, allowing:
|
||||
- Different smoothing for vx, vy, and omega in each environment
|
||||
- Independent control over linear and angular velocity response
|
||||
|
||||
Args:
|
||||
velocity_commands (torch.Tensor): Raw velocity commands [num_envs, 3] (vx, vy, omega)
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Filtered velocity commands with same shape as input
|
||||
"""
|
||||
if not self._enable_low_pass_filter:
|
||||
return velocity_commands
|
||||
|
||||
# Use per-environment per-dimension alpha values for filtering
|
||||
# Shape: [num_envs, action_dim] - already matches velocity_commands shape
|
||||
alpha_values = self._per_env_per_dim_low_pass_alpha
|
||||
|
||||
# Apply exponential smoothing (low-pass filter) with per-environment per-dimension alpha
|
||||
filtered_commands = (
|
||||
alpha_values * self._prev_filtered_velocity_commands
|
||||
+ (1.0 - alpha_values) * velocity_commands
|
||||
)
|
||||
|
||||
# Update previous filtered commands for next iteration
|
||||
self._prev_filtered_velocity_commands.copy_(filtered_commands)
|
||||
|
||||
return filtered_commands
|
||||
|
||||
def process_actions(self, actions):
|
||||
"""Process low-level navigation actions. This function is called with a frequency of 10Hz.
|
||||
|
||||
Args:
|
||||
actions (torch.Tensor): The low-level navigation actions.
|
||||
"""
|
||||
# Store the raw low-level navigation actions
|
||||
self._raw_navigation_velocity_actions[:] = actions
|
||||
# Apply the affine transformations
|
||||
if not self.cfg.use_raw_actions:
|
||||
self._processed_navigation_velocity_actions = (
|
||||
self._raw_navigation_velocity_actions * self._scale + self._offset
|
||||
)
|
||||
else:
|
||||
self._processed_navigation_velocity_actions[:] = self._raw_navigation_velocity_actions
|
||||
|
||||
if self.cfg.policy_distr_type == "gaussian":
|
||||
# scale the actions to the range [-1, 1] for gaussian distribution
|
||||
self._processed_navigation_velocity_actions = torch.tanh(self._processed_navigation_velocity_actions)
|
||||
elif self.cfg.policy_distr_type == "beta":
|
||||
# scale the actions to the range [-1, 1] for beta distribution
|
||||
self._processed_navigation_velocity_actions = (self._processed_navigation_velocity_actions - 0.5) * 2.0
|
||||
else:
|
||||
raise ValueError(f"Unknown policy distribution type: {self.cfg.policy_distr_type}")
|
||||
|
||||
# compute the current speed of the robot to generate low-level actions based on the current speed
|
||||
observations = self._env.observation_manager.compute_group(group_name=self.cfg.observation_group)
|
||||
base_lin_vel = observations[:, 0:3]
|
||||
vel_xyz = base_lin_vel.norm(dim=1, keepdim=True)
|
||||
|
||||
# [vx, vy, omega]
|
||||
self._processed_navigation_velocity_actions = (self._processed_navigation_velocity_actions + vel_xyz * self._policy_bias) * self._policy_scaling
|
||||
|
||||
# Apply low-pass filter to smooth velocity commands and add delay effect
|
||||
self._processed_navigation_velocity_actions = self.apply_low_pass_filter(self._processed_navigation_velocity_actions)
|
||||
|
||||
@torch.inference_mode()
|
||||
def apply_actions(self):
|
||||
"""Apply low-level actions for the simulator to the physics engine. This functions is called with the
|
||||
simulation frequency of 200Hz. Since low-level locomotion runs at 50Hz, we need to decimate the actions."""
|
||||
|
||||
if self._counter % self.cfg.low_level_decimation == 0:
|
||||
self._counter = 0
|
||||
self._prev_low_level_position_actions[:] = self._low_level_position_actions.clone()
|
||||
if self.low_level_velocity_action_term is not None:
|
||||
self._prev_low_level_velocity_actions[:] = self._low_level_velocity_actions.clone()
|
||||
|
||||
# Get low level actions from low level policy
|
||||
actions_phase = self.low_level_policy(
|
||||
self._env.observation_manager.compute_group(group_name=self.cfg.observation_group)
|
||||
)
|
||||
|
||||
# Process actions and bring them in the right order
|
||||
pos_dim = self.low_level_position_action_term.action_dim
|
||||
self._low_level_position_actions[:] = actions_phase[:, :pos_dim]
|
||||
self.low_level_position_action_term.process_actions(self._low_level_position_actions)
|
||||
if self.low_level_velocity_action_term is not None:
|
||||
self._low_level_velocity_actions[:] = actions_phase[:, pos_dim:]
|
||||
self.low_level_velocity_action_term.process_actions(self._low_level_velocity_actions)
|
||||
|
||||
# Apply low level actions
|
||||
self.low_level_position_action_term.apply_actions()
|
||||
if self.low_level_velocity_action_term is not None:
|
||||
self.low_level_velocity_action_term.apply_actions()
|
||||
self._counter += 1
|
||||
|
||||
def reset_low_pass_filter(self, env_ids: torch.Tensor):
|
||||
"""Reset low-pass filter state for specified environments.
|
||||
|
||||
Args:
|
||||
env_ids: Environment indices to reset.
|
||||
"""
|
||||
self._prev_filtered_velocity_commands[env_ids] = 0.0
|
||||
|
||||
"""
|
||||
Helper functions
|
||||
"""
|
||||
|
||||
def _init_buffers(self):
|
||||
# Prepare buffers
|
||||
self._raw_navigation_velocity_actions = torch.zeros(self.num_envs, self._action_dim, device=self.device)
|
||||
self._processed_navigation_velocity_actions = torch.zeros((self.num_envs, self._action_dim), device=self.device)
|
||||
self._low_level_position_actions = torch.zeros(self.num_envs, self.low_level_position_action_term.action_dim, device=self.device)
|
||||
vel_dim = self.low_level_velocity_action_term.action_dim if self.low_level_velocity_action_term is not None else 0
|
||||
self._low_level_velocity_actions = torch.zeros(self.num_envs, vel_dim, device=self.device)
|
||||
self._prev_low_level_position_actions = torch.zeros_like(self._low_level_position_actions)
|
||||
self._prev_low_level_velocity_actions = torch.zeros_like(self._low_level_velocity_actions)
|
||||
self._low_level_step_dt = self.cfg.low_level_decimation * self._env.physics_dt
|
||||
self._counter = 0
|
||||
self._scale = torch.tensor(self.cfg.scale, device=self.device)
|
||||
self._offset = torch.tensor(self.cfg.offset, device=self.device)
|
||||
self._policy_scaling = torch.tensor(self.cfg.policy_scaling, device=self.device).repeat(self.num_envs, 1)
|
||||
self._policy_bias = torch.zeros(self.num_envs, self._action_dim, device=self.device)
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import MISSING
|
||||
|
||||
from isaaclab.managers.action_manager import ActionTerm, ActionTermCfg
|
||||
from isaaclab.utils import configclass
|
||||
|
||||
from .navigation_se2_actions import PerceptiveNavigationSE2Action
|
||||
|
||||
|
||||
@configclass
|
||||
class PerceptiveNavigationSE2ActionCfg(ActionTermCfg):
|
||||
class_type: type[ActionTerm] = PerceptiveNavigationSE2Action
|
||||
""" Class of the action term."""
|
||||
low_level_decimation: int = 4
|
||||
"""Decimation factor for the low level action term."""
|
||||
use_raw_actions: bool = False
|
||||
"""Whether to use raw actions or not."""
|
||||
scale: list[float] = [1.0, 1.0, 1.0]
|
||||
"""Scale for the actions [vx, vy, w]."""
|
||||
offset: list[float] = [0.0, 0.0, 0.0]
|
||||
"""Offset for the actions [vx, vy, w]."""
|
||||
low_level_velocity_action: ActionTermCfg | None = None
|
||||
"""Configuration of the low level velocity action term.
|
||||
|
||||
Set to ``None`` for legs-only quadrupeds (e.g. Go2) whose locomotion policy
|
||||
only outputs joint position targets. The position-action branch then handles
|
||||
all 12 leg joints, and the velocity-action branch is fully bypassed."""
|
||||
low_level_position_action: ActionTermCfg = MISSING
|
||||
"""Configuration of the low level position action term."""
|
||||
low_level_policy_file: str = MISSING
|
||||
"""Path to the low level policy file."""
|
||||
observation_group: str = "policy"
|
||||
"""Observation group to use for the low level policy."""
|
||||
policy_scaling: list[float] = [1.0, 1.0, 1.0]
|
||||
"""Policy dependent scaling for the actions [vx, vy, w]."""
|
||||
reorder_joint_list: list[str] | None = None
|
||||
"""Reorder the joint actions given from the low-level policy to match the Isaac Sim order if policy has been
|
||||
trained with a different order. Set to None to disable reordering."""
|
||||
policy_distr_type: str = "gaussian"
|
||||
"""Policy distribution type: 'gaussian', 'beta'."""
|
||||
# Low-pass filter parameters
|
||||
enable_low_pass_filter: bool = True
|
||||
"""Whether to enable low-pass filtering for velocity commands."""
|
||||
low_pass_filter_alpha: float = 0.5
|
||||
"""Low-pass filter smoothing factor (0.0 = no smoothing, 1.0 = maximum smoothing).
|
||||
Formula: filtered_cmd = alpha * prev_filtered_cmd + (1 - alpha) * new_cmd"""
|
||||
+786
@@ -0,0 +1,786 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Goal command generator for robot navigation tasks.
|
||||
|
||||
Simplified architecture:
|
||||
1. Terrain generation creates `valid_mask` (boolean mask of valid positions with safety padding)
|
||||
2. This module samples goal/spawn positions uniformly from valid positions
|
||||
3. Z-height is looked up from the visual height field
|
||||
|
||||
The terrain module handles:
|
||||
- Height field generation
|
||||
- Obstacle detection
|
||||
- Safety padding (dilation)
|
||||
- Border exclusion
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Tuple, Optional
|
||||
|
||||
import torch
|
||||
|
||||
import isaaclab.utils.math as math_utils
|
||||
from isaaclab.assets import Articulation
|
||||
from isaaclab.managers import CommandTerm
|
||||
from isaaclab.markers import VisualizationMarkers
|
||||
from isaaclab.markers.config import (
|
||||
CUBOID_MARKER_CFG,
|
||||
GREEN_ARROW_X_MARKER_CFG,
|
||||
RED_ARROW_X_MARKER_CFG,
|
||||
)
|
||||
from isaaclab.utils.math import subtract_frame_transforms, transform_points, yaw_quat
|
||||
|
||||
from isaaclab_nav_task.navigation.mdp.math_utils import vec_to_quat
|
||||
from isaaclab_nav_task.terrains.terrain_constants import VERTICAL_SCALE
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedRLEnv
|
||||
from .goal_commands_cfg import RobotNavigationGoalCommandCfg
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Position Sampler (Simplified)
|
||||
# =============================================================================
|
||||
|
||||
class PositionSampler:
|
||||
"""Samples positions uniformly from valid terrain cells.
|
||||
|
||||
Uses pre-computed valid_mask from terrain generation (already has safety padding).
|
||||
For spawn positions, uses spawn_mask with larger padding to account for robot body.
|
||||
|
||||
Coordinate System:
|
||||
- The height field has shape (num_cells_x, num_cells_y), e.g., (299, 299)
|
||||
- Each cell is horizontal_scale meters (e.g., 0.1m)
|
||||
- The mesh is generated with a border: border_pixels = int(border_width / horizontal_scale) + 1
|
||||
- Our valid_mask corresponds to the inner terrain (no border)
|
||||
- Local coordinates are centered: (-terrain_size/2, -terrain_size/2) to (+terrain_size/2, +terrain_size/2)
|
||||
- Border pixel offset is computed dynamically from terrain config (not hardcoded)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
heights: torch.Tensor,
|
||||
valid_mask: torch.Tensor,
|
||||
platform_mask: torch.Tensor,
|
||||
terrain_size: float,
|
||||
horizontal_scale: float,
|
||||
device: torch.device,
|
||||
platform_repeat_count: int = 10,
|
||||
spawn_mask: torch.Tensor = None,
|
||||
border_width: float = 0.0,
|
||||
):
|
||||
"""Initialize the sampler.
|
||||
|
||||
Args:
|
||||
heights: Height field (num_terrains, width, height) for Z-lookup.
|
||||
valid_mask: Boolean mask of valid goal positions (num_terrains, width, height).
|
||||
platform_mask: Boolean mask of platform positions for curriculum.
|
||||
terrain_size: Size of each terrain in meters (full mesh size).
|
||||
horizontal_scale: Resolution of height field in meters per cell.
|
||||
device: Torch device.
|
||||
platform_repeat_count: Repetition count for platform positions.
|
||||
spawn_mask: Boolean mask of valid spawn positions with larger padding.
|
||||
If None, defaults to valid_mask.
|
||||
border_width: Border width around terrain in meters (from terrain config).
|
||||
"""
|
||||
self.device = device
|
||||
self.terrain_size = terrain_size
|
||||
self.horizontal_scale = horizontal_scale
|
||||
self.heights = heights
|
||||
self.valid_mask = valid_mask
|
||||
self.platform_mask = platform_mask
|
||||
# Use spawn_mask if provided, otherwise fall back to valid_mask
|
||||
self.spawn_mask = spawn_mask if spawn_mask is not None else valid_mask
|
||||
|
||||
# Use horizontal_scale as cell size (correct resolution)
|
||||
self.cell_size = horizontal_scale
|
||||
|
||||
# Compute border pixel offset dynamically based on terrain configuration
|
||||
# Formula matches patches.py: border_pixels = int(border_width / horizontal_scale) + 1
|
||||
# This ensures valid_mask indices map correctly to mesh coordinates
|
||||
self.border_pixels = int(border_width / horizontal_scale) + 1
|
||||
|
||||
# Mesh center offset (mesh is centered at origin after transform)
|
||||
self.mesh_center = terrain_size / 2 # e.g., 30 / 2 = 15m
|
||||
|
||||
# Build position tables for both goal and spawn sampling
|
||||
self._build_position_tables(platform_repeat_count)
|
||||
|
||||
def _build_position_tables(self, platform_repeat_count: int):
|
||||
"""Build pre-computed position tensors for efficient sampling.
|
||||
|
||||
Creates two sets of position tables:
|
||||
- Goal positions: from valid_mask with platform repetition for curriculum
|
||||
- Spawn positions: from spawn_mask (larger padding for robot body clearance)
|
||||
"""
|
||||
num_terrains = self.valid_mask.shape[0]
|
||||
|
||||
# =========================
|
||||
# Build GOAL position table (from valid_mask with platform repetition)
|
||||
# =========================
|
||||
valid_indices = self.valid_mask.nonzero(as_tuple=False)
|
||||
|
||||
# Build enhanced indices with platform repetition
|
||||
enhanced_indices = []
|
||||
for terrain_idx in range(num_terrains):
|
||||
terrain_valid = valid_indices[valid_indices[:, 0] == terrain_idx]
|
||||
|
||||
if len(terrain_valid) == 0:
|
||||
enhanced_indices.append(terrain_valid)
|
||||
continue
|
||||
|
||||
# Find platform positions
|
||||
terrain_platform = self.platform_mask[terrain_idx]
|
||||
platform_positions = terrain_platform.nonzero(as_tuple=False)
|
||||
|
||||
if len(platform_positions) > 0:
|
||||
# Check which valid positions are platforms (vectorized)
|
||||
valid_xy = terrain_valid[:, 1:] # (num_valid, 2)
|
||||
plat_xy = platform_positions # (num_platforms, 2)
|
||||
|
||||
# Broadcast compare: (num_valid, 1, 2) vs (1, num_platforms, 2)
|
||||
matches = (valid_xy.unsqueeze(1) == plat_xy.unsqueeze(0)).all(dim=2)
|
||||
is_platform = matches.any(dim=1)
|
||||
|
||||
# Repeat platform positions
|
||||
platform_valid = terrain_valid[is_platform]
|
||||
if len(platform_valid) > 0:
|
||||
repeated = platform_valid.repeat(platform_repeat_count, 1)
|
||||
terrain_valid = torch.cat([terrain_valid, repeated], dim=0)
|
||||
|
||||
enhanced_indices.append(terrain_valid)
|
||||
|
||||
# Count positions per terrain for goals
|
||||
self.count_per_terrain = torch.zeros(num_terrains, dtype=torch.long, device=self.device)
|
||||
for terrain_idx in range(num_terrains):
|
||||
self.count_per_terrain[terrain_idx] = len(enhanced_indices[terrain_idx])
|
||||
|
||||
# Create padded tensor for goal positions
|
||||
max_count = max(1, self.count_per_terrain.max().item())
|
||||
self.positions = torch.full(
|
||||
(num_terrains, max_count, 3), -1, dtype=torch.long, device=self.device
|
||||
)
|
||||
|
||||
# Fill goal position tables
|
||||
for terrain_idx in range(num_terrains):
|
||||
terrain_positions = enhanced_indices[terrain_idx]
|
||||
num_pos = terrain_positions.shape[0]
|
||||
if num_pos > 0:
|
||||
self.positions[terrain_idx, :num_pos] = terrain_positions
|
||||
|
||||
# =========================
|
||||
# Build SPAWN position table (from spawn_mask, no platform repetition)
|
||||
# =========================
|
||||
spawn_indices = self.spawn_mask.nonzero(as_tuple=False)
|
||||
|
||||
# Count spawn positions per terrain
|
||||
self.spawn_count_per_terrain = torch.zeros(num_terrains, dtype=torch.long, device=self.device)
|
||||
spawn_positions_list = []
|
||||
for terrain_idx in range(num_terrains):
|
||||
terrain_spawn = spawn_indices[spawn_indices[:, 0] == terrain_idx]
|
||||
self.spawn_count_per_terrain[terrain_idx] = len(terrain_spawn)
|
||||
spawn_positions_list.append(terrain_spawn)
|
||||
|
||||
# Create padded tensor for spawn positions
|
||||
max_spawn_count = max(1, self.spawn_count_per_terrain.max().item())
|
||||
self.spawn_positions = torch.full(
|
||||
(num_terrains, max_spawn_count, 3), -1, dtype=torch.long, device=self.device
|
||||
)
|
||||
|
||||
# Fill spawn position tables
|
||||
for terrain_idx in range(num_terrains):
|
||||
terrain_positions = spawn_positions_list[terrain_idx]
|
||||
num_pos = terrain_positions.shape[0]
|
||||
if num_pos > 0:
|
||||
self.spawn_positions[terrain_idx, :num_pos] = terrain_positions
|
||||
|
||||
def sample(self, terrain_indices: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Sample GOAL positions for given terrain indices.
|
||||
|
||||
Uses valid_mask with platform repetition for curriculum learning.
|
||||
|
||||
Args:
|
||||
terrain_indices: Tensor of terrain indices to sample from.
|
||||
|
||||
Returns:
|
||||
Tuple of (x, y, z) local coordinates in meters.
|
||||
"""
|
||||
return self._sample_from_table(
|
||||
terrain_indices,
|
||||
self.positions,
|
||||
self.count_per_terrain
|
||||
)
|
||||
|
||||
def sample_spawn(self, terrain_indices: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Sample SPAWN positions for given terrain indices.
|
||||
|
||||
Uses spawn_mask with larger padding to account for:
|
||||
- Robot body dimensions with random orientation
|
||||
- Platform edge safety margins
|
||||
- Controller startup behavior
|
||||
|
||||
Args:
|
||||
terrain_indices: Tensor of terrain indices to sample from.
|
||||
|
||||
Returns:
|
||||
Tuple of (x, y, z) local coordinates in meters.
|
||||
"""
|
||||
return self._sample_from_table(
|
||||
terrain_indices,
|
||||
self.spawn_positions,
|
||||
self.spawn_count_per_terrain
|
||||
)
|
||||
|
||||
def _sample_from_table(
|
||||
self,
|
||||
terrain_indices: torch.Tensor,
|
||||
positions_table: torch.Tensor,
|
||||
count_per_terrain: torch.Tensor
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Internal method to sample from a position table.
|
||||
|
||||
Args:
|
||||
terrain_indices: Tensor of terrain indices to sample from.
|
||||
positions_table: Pre-computed position table (num_terrains, max_count, 3).
|
||||
count_per_terrain: Number of valid positions per terrain.
|
||||
|
||||
Returns:
|
||||
Tuple of (x, y, z) local coordinates in meters.
|
||||
"""
|
||||
num_samples = terrain_indices.shape[0]
|
||||
|
||||
# Random indices within valid range
|
||||
valid_counts = count_per_terrain[terrain_indices].float().clamp(min=1)
|
||||
random_indices = (torch.rand(num_samples, device=self.device) * valid_counts).long()
|
||||
|
||||
# Lookup positions
|
||||
selected = positions_table[terrain_indices, random_indices] # (n, 3)
|
||||
is_valid = selected[:, 0] >= 0
|
||||
|
||||
local_x = torch.zeros(num_samples, device=self.device)
|
||||
local_y = torch.zeros(num_samples, device=self.device)
|
||||
local_z = torch.zeros(num_samples, device=self.device)
|
||||
|
||||
if is_valid.any():
|
||||
valid_selected = selected[is_valid]
|
||||
x_idx = valid_selected[:, 1]
|
||||
y_idx = valid_selected[:, 2]
|
||||
|
||||
# Convert to meters (accounting for border pixel offset)
|
||||
# The mesh is generated with @height_field_to_mesh which adds a border
|
||||
# border_pixels = int(border_width / horizontal_scale) + 1 (computed dynamically)
|
||||
# Mesh vertex at (i, j) has position: (i * h_scale - terrain_size/2, j * h_scale - terrain_size/2)
|
||||
# Our valid_mask[i, j] corresponds to mesh heights[i + border_pixels, j + border_pixels]
|
||||
# So the world position for valid_mask[i, j] is:
|
||||
# x = (i + border_pixels) * h_scale - terrain_size/2
|
||||
# y = (j + border_pixels) * h_scale - terrain_size/2
|
||||
local_x[is_valid] = (x_idx.float() + self.border_pixels) * self.cell_size - self.mesh_center
|
||||
local_y[is_valid] = (y_idx.float() + self.border_pixels) * self.cell_size - self.mesh_center
|
||||
|
||||
# Lookup Z from heights (our heights tensor matches valid_mask dimensions, no border offset needed)
|
||||
height_values = self.heights[valid_selected[:, 0], x_idx, y_idx]
|
||||
local_z[is_valid] = height_values.float() * VERTICAL_SCALE
|
||||
|
||||
return local_x, local_y, local_z
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Success Rate Tracker
|
||||
# =============================================================================
|
||||
|
||||
class SuccessRateTracker:
|
||||
"""Tracks navigation success rates using a rolling buffer."""
|
||||
|
||||
def __init__(self, num_envs: int, device: torch.device, buffer_size: int = 10):
|
||||
self.device = device
|
||||
self.buffer_size = buffer_size
|
||||
self.buffer = torch.full((num_envs, buffer_size), -1.0, device=device)
|
||||
self.write_index = torch.zeros(num_envs, dtype=torch.long, device=device)
|
||||
|
||||
def record_result(self, success: torch.Tensor, env_ids: torch.Tensor):
|
||||
indices = self.write_index[env_ids] % self.buffer_size
|
||||
self.buffer[env_ids, indices] = success[env_ids].float()
|
||||
self.write_index[env_ids] += 1
|
||||
|
||||
def add(self, results: torch.Tensor, env_ids: torch.Tensor):
|
||||
"""Legacy alias."""
|
||||
self.record_result(results, env_ids)
|
||||
|
||||
def clear(self, env_ids: torch.Tensor):
|
||||
self.buffer[env_ids] = -1.0
|
||||
self.write_index[env_ids] = 0
|
||||
|
||||
def get_success_rate(self) -> torch.Tensor:
|
||||
filled_count = (self.buffer >= 0).sum(dim=1).clamp(min=1)
|
||||
success_count = (self.buffer > 0).sum(dim=1)
|
||||
return success_count.float() / filled_count.float()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Main Navigation Goal Command Generator
|
||||
# =============================================================================
|
||||
|
||||
class RobotNavigationGoalCommand(CommandTerm):
|
||||
"""Command generator for robot navigation goal positions.
|
||||
|
||||
Samples goal and spawn positions from terrain-provided valid_mask.
|
||||
"""
|
||||
|
||||
cfg: RobotNavigationGoalCommandCfg
|
||||
|
||||
def __init__(self, cfg: RobotNavigationGoalCommandCfg, env: ManagerBasedRLEnv):
|
||||
super().__init__(cfg, env)
|
||||
|
||||
self.env = env
|
||||
self.robot: Articulation = env.scene[cfg.asset_name]
|
||||
|
||||
# Terrain configuration
|
||||
terrain_cfg = env.scene.terrain.cfg.terrain_generator
|
||||
self.num_terrain_rows = terrain_cfg.num_rows
|
||||
self.num_terrain_cols = terrain_cfg.num_cols
|
||||
self.terrain_size = terrain_cfg.size[0]
|
||||
|
||||
# Initialize buffers
|
||||
self._init_command_buffers()
|
||||
self._init_tracking_buffers()
|
||||
self._init_metrics()
|
||||
|
||||
# Position sampling (lazy initialization)
|
||||
self._sampling_initialized = False
|
||||
self._position_sampler: Optional[PositionSampler] = None
|
||||
|
||||
def _init_command_buffers(self):
|
||||
"""Initialize command state buffers."""
|
||||
# Goal in body frame: [direction_x, direction_y, direction_z, log_distance]
|
||||
self.goal_command_body = torch.zeros(self.num_envs, 4, device=self.device)
|
||||
self.goal_command_body_unscaled = torch.ones(self.num_envs, 3, device=self.device)
|
||||
|
||||
# World frame positions
|
||||
self.goal_position_world = torch.zeros(self.num_envs, 3, device=self.device)
|
||||
self.goal_position_world[:, 2] = 0.5
|
||||
|
||||
self.spawn_position_world = torch.zeros(self.num_envs, 3, device=self.device)
|
||||
self.spawn_position_world[:, 2] = 0.5
|
||||
self.spawn_position_world[:, :2] = self.env.scene.env_origins[:, :2]
|
||||
|
||||
self.spawn_heading_world = torch.zeros(self.num_envs, device=self.device)
|
||||
|
||||
def _init_tracking_buffers(self):
|
||||
"""Initialize goal tracking buffers."""
|
||||
self.steps_at_goal = torch.zeros(self.num_envs, device=self.device)
|
||||
self.time_at_goal = torch.zeros(self.num_envs, device=self.device)
|
||||
self.required_steps_at_goal = 4.0 / self.env.step_dt
|
||||
|
||||
self.initial_distance_to_goal = torch.zeros(self.num_envs, device=self.device)
|
||||
self.distance_to_goal = torch.zeros(self.num_envs, device=self.device)
|
||||
self.closest_distance_to_goal = torch.zeros(self.num_envs, device=self.device)
|
||||
|
||||
self.total_distance_traveled = torch.zeros(self.num_envs, device=self.device)
|
||||
self.previous_position = torch.zeros(self.num_envs, 3, device=self.device)
|
||||
|
||||
self.goal_reach_count = torch.zeros(self.num_envs, device=self.device)
|
||||
self.success_tracker = SuccessRateTracker(self.num_envs, self.device, buffer_size=10)
|
||||
self.success_rate_buffer = torch.full((self.num_envs, 10), -1.0, device=self.device)
|
||||
|
||||
def _init_metrics(self):
|
||||
"""Initialize performance metrics."""
|
||||
self.metrics["velocity_toward_goal"] = torch.zeros(self.num_envs, device=self.device)
|
||||
self.metrics["velocity_magnitude"] = torch.zeros(self.num_envs, device=self.device)
|
||||
self.metrics["success_rate"] = torch.zeros(self.num_envs, device=self.device)
|
||||
|
||||
# =========================================================================
|
||||
# Command Interface
|
||||
# =========================================================================
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"NavigationGoalCommand:\n\tCommand dimension: {tuple(self.command.shape[1:])}\n"
|
||||
|
||||
@property
|
||||
def command(self) -> torch.Tensor:
|
||||
return self.goal_command_body
|
||||
|
||||
def _get_unscaled_command(self) -> torch.Tensor:
|
||||
return self.goal_command_body_unscaled
|
||||
|
||||
# =========================================================================
|
||||
# Position Sampling
|
||||
# =========================================================================
|
||||
|
||||
def _initialize_position_sampling(self):
|
||||
"""Initialize position sampling from terrain-provided masks."""
|
||||
if self._sampling_initialized:
|
||||
return
|
||||
|
||||
# Get terrain data from scene.terrain (stored by patches system)
|
||||
terrain = self.env.scene.terrain
|
||||
|
||||
# Check for height field data
|
||||
heights_raw = getattr(terrain, '_height_field_visual', None)
|
||||
valid_mask_raw = getattr(terrain, '_height_field_valid_mask', None)
|
||||
platform_mask_raw = getattr(terrain, '_height_field_platform_mask', None)
|
||||
spawn_mask_raw = getattr(terrain, '_height_field_spawn_mask', None)
|
||||
|
||||
if heights_raw is None or valid_mask_raw is None:
|
||||
raise ValueError(
|
||||
"No height field data found on terrain. "
|
||||
"Ensure add_goal=True is set in terrain configuration and patches are applied."
|
||||
)
|
||||
|
||||
# Move to device
|
||||
heights = heights_raw.to(self.device)
|
||||
valid_mask = valid_mask_raw.to(self.device)
|
||||
|
||||
# Platform mask defaults to empty if not provided
|
||||
if platform_mask_raw is not None:
|
||||
platform_mask = platform_mask_raw.to(self.device)
|
||||
else:
|
||||
platform_mask = torch.zeros_like(valid_mask)
|
||||
|
||||
# Spawn mask defaults to valid_mask if not provided
|
||||
if spawn_mask_raw is not None:
|
||||
spawn_mask = spawn_mask_raw.to(self.device)
|
||||
else:
|
||||
spawn_mask = valid_mask # Fall back to goal mask
|
||||
|
||||
# Get terrain configuration parameters
|
||||
terrain_cfg = self.env.scene.terrain.cfg.terrain_generator
|
||||
horizontal_scale = terrain_cfg.horizontal_scale
|
||||
# Note: border_width for height_field_to_mesh comes from sub-terrain config (HfTerrainBaseCfg),
|
||||
# NOT from TerrainGeneratorCfg. Sub-terrain configs default to border_width=0.0.
|
||||
# TerrainGeneratorCfg.border_width (e.g., 30.0) is for the outer grid border, not per-tile.
|
||||
sub_terrain_border_width = 0.0 # Default from HfTerrainBaseCfg
|
||||
|
||||
# Create sampler with both goal (valid_mask) and spawn (spawn_mask) masks
|
||||
self._position_sampler = PositionSampler(
|
||||
heights=heights,
|
||||
valid_mask=valid_mask,
|
||||
platform_mask=platform_mask,
|
||||
terrain_size=self.terrain_size,
|
||||
horizontal_scale=horizontal_scale,
|
||||
device=self.device,
|
||||
spawn_mask=spawn_mask,
|
||||
border_width=sub_terrain_border_width,
|
||||
)
|
||||
|
||||
self._sampling_initialized = True
|
||||
|
||||
def _get_terrain_indices(self, env_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""Get terrain indices for given environment IDs.
|
||||
|
||||
The terrain index formula depends on the generation order:
|
||||
- curriculum=True: column-major (row + col * num_rows)
|
||||
- curriculum=False: row-major (row * num_cols + col)
|
||||
|
||||
Note: terrain_levels corresponds to row, terrain_types to column.
|
||||
"""
|
||||
terrain = self.env.scene.terrain
|
||||
levels = terrain.terrain_levels[env_ids] # row
|
||||
types = terrain.terrain_types[env_ids] # col
|
||||
|
||||
# Check if curriculum mode
|
||||
terrain_cfg = self.env.scene.terrain.cfg.terrain_generator
|
||||
if terrain_cfg.curriculum:
|
||||
# Column-major order (curriculum mode iterates: for col: for row:)
|
||||
return levels + types * self.num_terrain_rows
|
||||
else:
|
||||
# Row-major order (random mode uses np.unravel_index with (num_rows, num_cols))
|
||||
return levels * self.num_terrain_cols + types
|
||||
|
||||
# =========================================================================
|
||||
# Command Sampling and Update
|
||||
# =========================================================================
|
||||
|
||||
def _resample_command(self, env_ids: Sequence[int]):
|
||||
"""Sample new goal and spawn positions.
|
||||
|
||||
Goal positions use valid_mask (smaller padding, robot just needs to reach).
|
||||
Spawn positions use spawn_mask (larger padding for robot body clearance).
|
||||
"""
|
||||
self._initialize_position_sampling()
|
||||
|
||||
if isinstance(env_ids, torch.Tensor):
|
||||
env_ids_tensor = env_ids.clone().to(device=self.device, dtype=torch.long)
|
||||
else:
|
||||
env_ids_tensor = torch.as_tensor(env_ids, device=self.device, dtype=torch.long)
|
||||
|
||||
# Reset tracking
|
||||
self._reset_tracking_state(env_ids)
|
||||
|
||||
# Get terrain indices
|
||||
terrain_indices = self._get_terrain_indices(env_ids_tensor)
|
||||
|
||||
# Sample goal positions (from valid_mask with smaller padding)
|
||||
goal_x, goal_y, goal_z = self._position_sampler.sample(terrain_indices)
|
||||
# Sample spawn positions (from spawn_mask with larger padding for robot body)
|
||||
spawn_x, spawn_y, spawn_z = self._position_sampler.sample_spawn(terrain_indices)
|
||||
|
||||
# Convert to world coordinates
|
||||
terrain = self.env.scene.terrain
|
||||
levels = terrain.terrain_levels[env_ids]
|
||||
types = terrain.terrain_types[env_ids]
|
||||
terrain_origins = terrain.terrain_origins[levels, types]
|
||||
|
||||
# Update goal position
|
||||
self.goal_position_world[env_ids, 0] = terrain_origins[:, 0] + goal_x
|
||||
self.goal_position_world[env_ids, 1] = terrain_origins[:, 1] + goal_y
|
||||
h_min, h_max = self.cfg.goal_height_offset_range
|
||||
height_offset = torch.rand(len(env_ids), device=self.device) * (h_max - h_min) + h_min
|
||||
self.goal_position_world[env_ids, 2] = goal_z + height_offset
|
||||
|
||||
# Small spawn height offset to prevent clipping into terrain
|
||||
# Note: The robot's default_root_state already includes standing height (~0.5m)
|
||||
spawn_offset = 0.05
|
||||
|
||||
# Update spawn/env origin
|
||||
terrain.env_origins[env_ids, 0] = terrain_origins[:, 0] + spawn_x
|
||||
terrain.env_origins[env_ids, 1] = terrain_origins[:, 1] + spawn_y
|
||||
terrain.env_origins[env_ids, 2] = spawn_z + spawn_offset
|
||||
|
||||
# Track spawn position
|
||||
self.spawn_position_world[env_ids, 0] = terrain_origins[:, 0] + spawn_x
|
||||
self.spawn_position_world[env_ids, 1] = terrain_origins[:, 1] + spawn_y
|
||||
self.spawn_position_world[env_ids, 2] = spawn_z + spawn_offset
|
||||
|
||||
# Initialize distance metrics
|
||||
self.initial_distance_to_goal[env_ids] = torch.norm(
|
||||
self.robot.data.root_pos_w[env_ids] - self.goal_position_world[env_ids], dim=1
|
||||
)
|
||||
self.closest_distance_to_goal[env_ids] = self.initial_distance_to_goal[env_ids]
|
||||
|
||||
def _reset_tracking_state(self, env_ids: Sequence[int]):
|
||||
"""Reset tracking state for specified environments."""
|
||||
self.steps_at_goal[env_ids] = 0
|
||||
self.time_at_goal[env_ids] = 0
|
||||
self.total_distance_traveled[env_ids] = 0.0
|
||||
self.previous_position[env_ids] = self.robot.data.root_pos_w[env_ids].clone()
|
||||
|
||||
def _update_command(self):
|
||||
"""Update command in body frame."""
|
||||
# Transform goal to body frame
|
||||
inverse_pos, inverse_rot = subtract_frame_transforms(
|
||||
self.robot.data.root_pos_w,
|
||||
self.robot.data.root_quat_w
|
||||
)
|
||||
goal_in_body = transform_points(
|
||||
self.goal_position_world.unsqueeze(1),
|
||||
inverse_pos,
|
||||
inverse_rot
|
||||
).squeeze(1)
|
||||
|
||||
self.goal_command_body_unscaled = goal_in_body.clone()
|
||||
|
||||
# Normalized direction and log distance
|
||||
distance = torch.norm(goal_in_body, dim=-1, keepdim=True)
|
||||
direction = goal_in_body / torch.clamp(distance, min=1e-6)
|
||||
log_distance = torch.log(distance + 1.0)
|
||||
|
||||
self.goal_command_body[:, :3] = direction
|
||||
self.goal_command_body[:, 3:] = log_distance
|
||||
|
||||
self._update_distance_tracking()
|
||||
|
||||
def _update_distance_tracking(self):
|
||||
"""Update distance metrics."""
|
||||
self.distance_to_goal = torch.norm(
|
||||
self.robot.data.root_pos_w - self.goal_position_world, dim=1
|
||||
)
|
||||
self.closest_distance_to_goal = torch.min(
|
||||
self.closest_distance_to_goal, self.distance_to_goal
|
||||
)
|
||||
|
||||
step_distance = torch.norm(
|
||||
self.robot.data.root_pos_w - self.previous_position, dim=1
|
||||
)
|
||||
self.total_distance_traveled += step_distance
|
||||
self.previous_position = self.robot.data.root_pos_w.clone()
|
||||
|
||||
def _resample_spawn_positions(self, env_ids: Sequence[int]):
|
||||
"""Update spawn position tracking."""
|
||||
self.spawn_position_world[env_ids, :2] = self.env.scene.env_origins[env_ids, :2]
|
||||
|
||||
# =========================================================================
|
||||
# Metrics and Reset
|
||||
# =========================================================================
|
||||
|
||||
def _update_metrics(self):
|
||||
"""Update performance metrics."""
|
||||
position_error = self.goal_position_world - self.robot.data.root_pos_w[:, :3]
|
||||
position_error_2d = position_error[:, :2]
|
||||
velocity_2d = self.robot.data.root_state_w[:, 7:9]
|
||||
|
||||
self.metrics["velocity_magnitude"] = torch.norm(velocity_2d, dim=1)
|
||||
|
||||
direction_to_goal = position_error_2d / torch.clamp(torch.norm(position_error_2d, dim=1, keepdim=True), min=1e-6)
|
||||
self.metrics["velocity_toward_goal"] = (velocity_2d * direction_to_goal).sum(dim=1)
|
||||
self.metrics["success_rate"] = self.success_tracker.get_success_rate()
|
||||
|
||||
def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]:
|
||||
"""Reset command generator and compute episode metrics."""
|
||||
metrics_obs = self.env.observation_manager.compute_group(group_name="metrics")
|
||||
success = metrics_obs["in_goal"][env_ids].squeeze(-1)
|
||||
failed = ~success
|
||||
|
||||
# Update legacy success rate buffer
|
||||
self.success_rate_buffer[env_ids] = torch.roll(
|
||||
self.success_rate_buffer[env_ids], 1, dims=1
|
||||
)
|
||||
rate = success.float() - failed.float()
|
||||
rate[rate == 0] = -1
|
||||
self.success_rate_buffer[env_ids, 0] = rate
|
||||
|
||||
# Reset command state
|
||||
if env_ids is None:
|
||||
env_ids = slice(None)
|
||||
self.command_counter[env_ids] = 0
|
||||
self._resample(env_ids)
|
||||
|
||||
# Return mean metrics
|
||||
extras = {}
|
||||
for name, value in self.metrics.items():
|
||||
extras[name] = torch.mean(value[env_ids]).item()
|
||||
value[env_ids] = 0.0
|
||||
|
||||
return extras
|
||||
|
||||
# =========================================================================
|
||||
# Success/Failure Tracking
|
||||
# =========================================================================
|
||||
|
||||
def update_success(self, at_goal: torch.Tensor):
|
||||
self.goal_reach_count += at_goal.int()
|
||||
|
||||
def update_failures(self, failed: torch.Tensor):
|
||||
self.goal_reach_count -= failed.int()
|
||||
|
||||
# =========================================================================
|
||||
# Visualization
|
||||
# =========================================================================
|
||||
|
||||
def _set_debug_vis_impl(self, debug_vis: bool):
|
||||
if debug_vis:
|
||||
self._setup_visualizers()
|
||||
else:
|
||||
self._hide_visualizers()
|
||||
|
||||
def _setup_visualizers(self):
|
||||
"""Create visualization markers."""
|
||||
if not hasattr(self, "goal_marker"):
|
||||
cfg = CUBOID_MARKER_CFG.copy()
|
||||
cfg.prim_path = "/Visuals/Command/goal_position"
|
||||
cfg.markers["cuboid"].size = (0.2, 0.2, 0.2)
|
||||
cfg.markers["cuboid"].visual_material.diffuse_color = (0.0, 0.0, 1.0)
|
||||
self.goal_marker = VisualizationMarkers(cfg)
|
||||
|
||||
if not hasattr(self, "spawn_marker"):
|
||||
cfg = CUBOID_MARKER_CFG.copy()
|
||||
cfg.prim_path = "/Visuals/Command/spawn_position"
|
||||
cfg.markers["cuboid"].size = (0.2, 0.2, 0.2)
|
||||
cfg.markers["cuboid"].visual_material.diffuse_color = (1.0, 0.5, 0.0)
|
||||
self.spawn_marker = VisualizationMarkers(cfg)
|
||||
|
||||
if not hasattr(self, "desired_velocity_marker"):
|
||||
cfg = GREEN_ARROW_X_MARKER_CFG.copy()
|
||||
cfg.prim_path = "/Visuals/Command/desired_velocity"
|
||||
cfg.markers["arrow"].scale = (0.5, 0.5, 0.5)
|
||||
self.desired_velocity_marker = VisualizationMarkers(cfg)
|
||||
|
||||
if not hasattr(self, "current_velocity_marker"):
|
||||
cfg = RED_ARROW_X_MARKER_CFG.copy()
|
||||
cfg.prim_path = "/Visuals/Command/current_velocity"
|
||||
cfg.markers["arrow"].scale = (0.5, 0.5, 0.5)
|
||||
self.current_velocity_marker = VisualizationMarkers(cfg)
|
||||
|
||||
self.goal_marker.set_visibility(True)
|
||||
self.spawn_marker.set_visibility(True)
|
||||
self.desired_velocity_marker.set_visibility(True)
|
||||
self.current_velocity_marker.set_visibility(True)
|
||||
|
||||
def _hide_visualizers(self):
|
||||
for name in ["goal_marker", "spawn_marker", "desired_velocity_marker", "current_velocity_marker"]:
|
||||
if hasattr(self, name):
|
||||
getattr(self, name).set_visibility(False)
|
||||
|
||||
def _debug_vis_callback(self, event):
|
||||
"""Update visualization markers."""
|
||||
self.goal_marker.visualize(self.goal_position_world)
|
||||
self.spawn_marker.visualize(self.spawn_position_world)
|
||||
|
||||
arrow_position = self.robot.data.root_pos_w.clone()
|
||||
arrow_position[:, 2] += 0.5
|
||||
|
||||
desired_scale, desired_quat = self._compute_velocity_arrow(
|
||||
self.command[:, :3], is_goal_direction=True
|
||||
)
|
||||
self.desired_velocity_marker.visualize(arrow_position, desired_quat, desired_scale)
|
||||
|
||||
current_scale, current_quat = self._compute_velocity_arrow(
|
||||
self.robot.data.root_lin_vel_b, is_goal_direction=False
|
||||
)
|
||||
self.current_velocity_marker.visualize(arrow_position, current_quat, current_scale)
|
||||
|
||||
def _compute_velocity_arrow(
|
||||
self,
|
||||
velocity: torch.Tensor,
|
||||
is_goal_direction: bool
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Compute arrow visualization parameters."""
|
||||
base_scale = torch.tensor(
|
||||
self.desired_velocity_marker.cfg.markers["arrow"].scale,
|
||||
device=self.device
|
||||
).repeat(velocity.shape[0], 1)
|
||||
|
||||
if not is_goal_direction:
|
||||
velocity = velocity.clone()
|
||||
velocity[:, 2] = 0.0
|
||||
|
||||
base_scale[:, 0] *= torch.norm(velocity, dim=1) * 3.0
|
||||
quat = vec_to_quat(velocity)
|
||||
|
||||
if is_goal_direction:
|
||||
quat = math_utils.quat_mul(self.robot.data.root_quat_w, quat)
|
||||
else:
|
||||
quat = math_utils.quat_mul(yaw_quat(self.robot.data.root_quat_w), quat)
|
||||
|
||||
return base_scale, quat
|
||||
|
||||
# =========================================================================
|
||||
# Utility Methods
|
||||
# =========================================================================
|
||||
|
||||
def _is_position_in_bounds(self, position: torch.Tensor) -> torch.Tensor:
|
||||
half_side = math.ceil(self.num_envs ** 0.5) * self.env.cfg.scene.env_spacing / 2
|
||||
return (position[:, :2].abs() < half_side).all(dim=1)
|
||||
|
||||
def _clamp_to_bounds(self, position: torch.Tensor) -> torch.Tensor:
|
||||
origins = self.env.scene.terrain.terrain_origins.view(-1, 3)
|
||||
bounds_max = origins.max(dim=0)[0][:2]
|
||||
bounds_min = origins.min(dim=0)[0][:2]
|
||||
position[:, 0] = position[:, 0].clamp(bounds_min[0], bounds_max[0])
|
||||
position[:, 1] = position[:, 1].clamp(bounds_min[1], bounds_max[1])
|
||||
return position
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Legacy Aliases
|
||||
# =============================================================================
|
||||
|
||||
RobotNavigationGoalCommand.pos_command_b = property(lambda self: self.goal_command_body)
|
||||
RobotNavigationGoalCommand.pos_command_w = property(lambda self: self.goal_position_world)
|
||||
RobotNavigationGoalCommand.pos_spawn_w = property(lambda self: self.spawn_position_world)
|
||||
RobotNavigationGoalCommand.closes_distance_to_goal = property(
|
||||
lambda self: self.closest_distance_to_goal
|
||||
)
|
||||
RobotNavigationGoalCommand.time_at_goal_in_steps = property(lambda self: self.steps_at_goal)
|
||||
RobotNavigationGoalCommand.required_time_at_goal_in_steps = property(
|
||||
lambda self: self.required_steps_at_goal
|
||||
)
|
||||
RobotNavigationGoalCommand.goal_reached_buffer = property(lambda self: self.success_tracker)
|
||||
RobotNavigationGoalCommand.goal_reached_counter = property(lambda self: self.goal_reach_count)
|
||||
RobotNavigationGoalCommand.distance_traveled = property(lambda self: self.total_distance_traveled)
|
||||
RobotNavigationGoalCommand.previous_pos_3d = property(lambda self: self.previous_position)
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import MISSING
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
from isaaclab.managers import CommandTermCfg
|
||||
from isaaclab.utils import configclass
|
||||
from isaaclab.utils.noise import AdditiveUniformNoiseCfg
|
||||
from .goal_commands import RobotNavigationGoalCommand
|
||||
|
||||
|
||||
"""
|
||||
Base command generator.
|
||||
"""
|
||||
|
||||
@configclass
|
||||
class RobotNavigationGoalCommandCfg(CommandTermCfg):
|
||||
"""Configuration for the robot goal command generator."""
|
||||
|
||||
class_type: type = RobotNavigationGoalCommand
|
||||
|
||||
asset_name: str = MISSING
|
||||
"""Name of the asset in the environment for which the commands are generated."""
|
||||
|
||||
robot_to_goal_line_vis: bool = True
|
||||
"""If true, visualize the line from the robot to the goal."""
|
||||
|
||||
goal_height_offset_range: tuple[float, float] = (0.2, 0.8)
|
||||
"""Range (min, max) in meters for the random goal height offset above the ground.
|
||||
|
||||
The goal cube is placed at ``terrain_ground_height + U(min, max)``. The default
|
||||
(0.2, 0.8) is tuned for the B2W body height (~0.5-0.6 m). For shorter robots
|
||||
such as the Go2 (standing height ~0.3 m), lower this range (e.g. (0.1, 0.4)) so
|
||||
the goal marker sits near the robot body. Note: success/termination is purely
|
||||
horizontal (xy), so this offset only affects visualization and the minor 3D
|
||||
``reach_goal_xy_soft`` reward term.
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,507 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Observation functions for navigation tasks.
|
||||
|
||||
These functions can be passed to :class:`isaaclab.managers.ObservationTermCfg`
|
||||
to specify observations for the policy.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Optional, cast
|
||||
|
||||
import torch
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from isaaclab.assets import Articulation, RigidObject
|
||||
from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
|
||||
import isaaclab_nav_task.navigation.mdp as mdp
|
||||
|
||||
from .depth_utils.camera_config import CameraConfig, DEFAULT_CAMERA_CONFIG, get_camera_config
|
||||
from .depth_utils.depth_noise_encoder import DepthNoiseEncoder
|
||||
from .heightscan_utils.heightscan_encoder import HeightScanFeatEncoder
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab_nav_task.navigation.mdp import PerceptiveNavigationSE2Action
|
||||
from isaaclab_nav_task.navigation.mdp.navigation.goal_commands import RobotNavigationGoalCommand
|
||||
|
||||
# Visualization state
|
||||
_DEPTH_VIZ_WINDOWS: dict[str, dict] = {}
|
||||
|
||||
# Global depth noise generator - will be initialized with camera config
|
||||
DEPTH_NOISE_GENERATOR: Optional[DepthNoiseEncoder] = None
|
||||
JIT_DEPTH_NOISE_GENERATOR: Optional[torch.jit.ScriptModule] = None
|
||||
|
||||
# Global height scan encoder - initialized once on first use
|
||||
HEIGHTSCAN_FEAT_ENCODER: Optional[HeightScanFeatEncoder] = None
|
||||
JIT_HEIGHTSCAN_FEAT_ENCODER: Optional[torch.jit.ScriptModule] = None
|
||||
|
||||
# Global flags/params for depth encoding
|
||||
use_jit = True
|
||||
min_depth = 0.0
|
||||
max_depth = 0.0
|
||||
|
||||
|
||||
def _update_depth_window(camera_name: str, depth_tensor: torch.Tensor, height: int, width: int, title: str) -> None:
|
||||
"""Update (or lazily create) a persistent matplotlib window for a camera.
|
||||
|
||||
This enables continuous, non-blocking visualization across steps for one or more cameras.
|
||||
"""
|
||||
try:
|
||||
# Convert to numpy and handle different tensor shapes
|
||||
if depth_tensor.dim() == 4: # [B, 1, H, W]
|
||||
depth_np = depth_tensor[0, 0].detach().cpu().numpy()
|
||||
elif depth_tensor.dim() == 2: # [B, H*W]
|
||||
depth_np = depth_tensor[0].detach().cpu().numpy().reshape(height, width)
|
||||
else:
|
||||
return
|
||||
|
||||
# Lazily create window if missing
|
||||
if camera_name not in _DEPTH_VIZ_WINDOWS:
|
||||
plt.ion()
|
||||
fig, ax = plt.subplots(1, 1, figsize=(8, 6))
|
||||
im = ax.imshow(depth_np, cmap="plasma", aspect="equal")
|
||||
cbar = plt.colorbar(im, ax=ax, label="Depth (meters)")
|
||||
ax.set_title(title)
|
||||
ax.set_xlabel("Width (pixels)")
|
||||
ax.set_ylabel("Height (pixels)")
|
||||
_DEPTH_VIZ_WINDOWS[camera_name] = {"fig": fig, "ax": ax, "im": im, "cbar": cbar}
|
||||
else:
|
||||
handle = _DEPTH_VIZ_WINDOWS[camera_name]
|
||||
im = handle["im"]
|
||||
ax = handle["ax"]
|
||||
ax.set_title(title)
|
||||
im.set_data(depth_np)
|
||||
# Optional: keep consistent color scaling per window
|
||||
im.set_clim(vmin=depth_np.min(), vmax=depth_np.max())
|
||||
|
||||
# Lightweight draw
|
||||
_DEPTH_VIZ_WINDOWS[camera_name]["fig"].canvas.draw_idle()
|
||||
plt.pause(0.001)
|
||||
except Exception as e:
|
||||
print(f"Error updating depth window for {camera_name}: {e}")
|
||||
|
||||
|
||||
def _ensure_depth_noise_generator_initialized(
|
||||
camera_config: Optional[CameraConfig] = None,
|
||||
use_jit_precompiled: bool = True,
|
||||
feature_dim: int = 64,
|
||||
):
|
||||
"""Ensure the depth noise generator is initialized with the correct configuration.
|
||||
|
||||
This is called automatically by observation functions that need depth encoding.
|
||||
If not explicitly initialized via initialize_depth_noise_generator(), it will use defaults.
|
||||
|
||||
Args:
|
||||
camera_config: The camera configuration to use. If None, uses DEFAULT_CAMERA_CONFIG.
|
||||
use_jit_precompiled: Whether to use JIT compilation for faster inference.
|
||||
feature_dim: Feature dimension for the encoder output.
|
||||
"""
|
||||
global DEPTH_NOISE_GENERATOR, JIT_DEPTH_NOISE_GENERATOR, use_jit, min_depth, max_depth
|
||||
|
||||
# Only initialize if not already done
|
||||
if DEPTH_NOISE_GENERATOR is not None:
|
||||
return
|
||||
|
||||
# Use provided camera config or default
|
||||
config = camera_config if camera_config is not None else DEFAULT_CAMERA_CONFIG
|
||||
|
||||
# Extract camera parameters
|
||||
min_depth = config.min_depth
|
||||
max_depth = config.max_depth
|
||||
resolution = config.resolution
|
||||
|
||||
print("=" * 80)
|
||||
print("Initializing depth noise generator for navigation observation:")
|
||||
print(f" Resolution: {resolution}")
|
||||
print(f" Depth range: [{min_depth}, {max_depth}]")
|
||||
print(f" Feature dim: {feature_dim}")
|
||||
print(f" Encoder path: {config.depth_encoder_path}")
|
||||
print(f" Use JIT: {use_jit_precompiled}")
|
||||
|
||||
# Initialize encoder with the new simplified API
|
||||
# The DepthNoiseEncoder now takes camera_config directly
|
||||
DEPTH_NOISE_GENERATOR = DepthNoiseEncoder(
|
||||
feature_dim=feature_dim,
|
||||
camera_config=config,
|
||||
).to(torch.device("cuda"))
|
||||
DEPTH_NOISE_GENERATOR.eval()
|
||||
|
||||
# Create JIT version for inference (optional optimization)
|
||||
use_jit = use_jit_precompiled
|
||||
if use_jit:
|
||||
# Try to find a JIT compiled version
|
||||
jit_path = config.depth_encoder_path.replace('.pth', '_jit.pt') if config.depth_encoder_path else None
|
||||
if jit_path and os.path.exists(jit_path):
|
||||
print(f" Loading precompiled JIT model from: {jit_path}")
|
||||
JIT_DEPTH_NOISE_GENERATOR = torch.jit.load(jit_path, map_location="cuda")
|
||||
JIT_DEPTH_NOISE_GENERATOR = torch.jit.optimize_for_inference(JIT_DEPTH_NOISE_GENERATOR)
|
||||
else:
|
||||
print(f" JIT compilation requested but no precompiled model found.")
|
||||
print(f" Creating JIT model from encoder...")
|
||||
example_input = torch.randn(1, 1, resolution[1], resolution[0]).cuda() # (B, C, H, W)
|
||||
JIT_DEPTH_NOISE_GENERATOR = torch.jit.trace(DEPTH_NOISE_GENERATOR, example_input)
|
||||
JIT_DEPTH_NOISE_GENERATOR = torch.jit.optimize_for_inference(JIT_DEPTH_NOISE_GENERATOR)
|
||||
else:
|
||||
JIT_DEPTH_NOISE_GENERATOR = DEPTH_NOISE_GENERATOR
|
||||
|
||||
print(" Depth noise generator initialized successfully")
|
||||
print("=" * 80)
|
||||
|
||||
|
||||
def initialize_depth_noise_generator(
|
||||
camera_config: Optional[CameraConfig] = None,
|
||||
robot_name: Optional[str] = None,
|
||||
use_jit_precompiled: bool = True,
|
||||
feature_dim: int = 64,
|
||||
):
|
||||
"""Initialize the depth noise generator with specific configuration.
|
||||
|
||||
This function should be called in the environment's __post_init__() to set up
|
||||
the depth encoder before any observations are computed.
|
||||
|
||||
Args:
|
||||
camera_config: The camera configuration to use. If None and robot_name is provided,
|
||||
uses the config for that robot. If both are None, uses DEFAULT_CAMERA_CONFIG.
|
||||
robot_name: Name of the robot (e.g., 'b2w', 'aow_d'). If provided and
|
||||
camera_config is None, automatically loads the appropriate camera config.
|
||||
use_jit_precompiled: Whether to use JIT compilation for faster inference. Defaults to True.
|
||||
feature_dim: Feature dimension for the encoder output. Defaults to 64.
|
||||
|
||||
Examples:
|
||||
# Using robot name (recommended for multi-robot training)
|
||||
initialize_depth_noise_generator(robot_name="b2w")
|
||||
|
||||
# Using explicit camera config
|
||||
initialize_depth_noise_generator(camera_config=ZEDX_CAMERA_CONFIG)
|
||||
"""
|
||||
# If camera_config not provided, try to get it from robot_name
|
||||
if camera_config is None and robot_name is not None:
|
||||
camera_config = get_camera_config(robot_name, use_default_fallback=False)
|
||||
|
||||
_ensure_depth_noise_generator_initialized(camera_config, use_jit_precompiled, feature_dim)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Observation Functions
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def generated_commands_reshaped(
|
||||
env: ManagerBasedRLEnv, command_name: str, unsqueeze_pos: int = 1, flatten: bool = False
|
||||
) -> torch.Tensor:
|
||||
"""The generated command from command term in the command manager with the given name."""
|
||||
if flatten:
|
||||
return env.command_manager.get_command(command_name)
|
||||
return env.command_manager.get_command(command_name).unsqueeze(unsqueeze_pos)
|
||||
|
||||
|
||||
def base_lin_vel_delayed(
|
||||
env: ManagerBasedEnv,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> torch.Tensor:
|
||||
"""Compute delayed root linear velocity.
|
||||
|
||||
Requires env.delay_manager to exist. Each environment has a random delay
|
||||
sampled from [0, max_delay] at episode reset.
|
||||
|
||||
Args:
|
||||
env: The environment object (must have delay_manager attribute).
|
||||
asset_cfg: The name of the asset.
|
||||
|
||||
Returns:
|
||||
The delayed linear velocity in the asset's root frame.
|
||||
"""
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
lin_vel = asset.data.root_lin_vel_b
|
||||
return env.delay_manager.compute_delayed_lin_vel(lin_vel)
|
||||
|
||||
|
||||
def base_ang_vel_delayed(
|
||||
env: ManagerBasedEnv,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> torch.Tensor:
|
||||
"""Compute delayed root angular velocity.
|
||||
|
||||
Requires env.delay_manager to exist. Each environment has a random delay
|
||||
sampled from [0, max_delay] at episode reset.
|
||||
|
||||
Args:
|
||||
env: The environment object (must have delay_manager attribute).
|
||||
asset_cfg: The name of the asset.
|
||||
|
||||
Returns:
|
||||
The delayed angular velocity in the asset's root frame.
|
||||
"""
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
ang_vel = asset.data.root_ang_vel_b
|
||||
return env.delay_manager.compute_delayed_ang_vel(ang_vel)
|
||||
|
||||
|
||||
def projected_gravity_delayed(
|
||||
env: ManagerBasedEnv,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> torch.Tensor:
|
||||
"""Compute delayed projected gravity.
|
||||
|
||||
Requires env.delay_manager to exist. Each environment has a random delay
|
||||
sampled from [0, max_delay] at episode reset.
|
||||
|
||||
Args:
|
||||
env: The environment object (must have delay_manager attribute).
|
||||
asset_cfg: The name of the asset.
|
||||
|
||||
Returns:
|
||||
The delayed projected gravity in the asset's root frame.
|
||||
"""
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
projected_gravity = asset.data.projected_gravity_b
|
||||
|
||||
return env.delay_manager.compute_delayed_projected_gravity(projected_gravity)
|
||||
|
||||
|
||||
def height_scan_feat(
|
||||
env: ManagerBasedEnv, sensor_cfg: SceneEntityCfg = SceneEntityCfg("height_scanner_critic"), offset: float = 0.5
|
||||
) -> torch.Tensor:
|
||||
"""Height scan feature from a ray caster sensor, encoded using a pre-trained VAE encoder.
|
||||
|
||||
The height scan is reshaped to a 51x51 grid and encoded using a VAE encoder that
|
||||
outputs a 64-channel, 7x7 spatial feature map (64*7*7 = 3136 features).
|
||||
|
||||
Args:
|
||||
env: The environment object.
|
||||
sensor_cfg: The configuration of the height scanner sensor.
|
||||
offset: Offset to subtract from the height values. Defaults to 0.5.
|
||||
|
||||
Returns:
|
||||
The encoded height scan features of shape (num_envs, 3136).
|
||||
"""
|
||||
global HEIGHTSCAN_FEAT_ENCODER, JIT_HEIGHTSCAN_FEAT_ENCODER
|
||||
|
||||
# Initialize the encoder on first call
|
||||
if HEIGHTSCAN_FEAT_ENCODER is None:
|
||||
print("Initializing height scan feature encoder...")
|
||||
HEIGHTSCAN_FEAT_ENCODER = HeightScanFeatEncoder(feature_dim=64).to(torch.device("cuda"))
|
||||
HEIGHTSCAN_FEAT_ENCODER.eval()
|
||||
JIT_HEIGHTSCAN_FEAT_ENCODER = torch.jit.script(HEIGHTSCAN_FEAT_ENCODER)
|
||||
JIT_HEIGHTSCAN_FEAT_ENCODER = torch.jit.optimize_for_inference(JIT_HEIGHTSCAN_FEAT_ENCODER)
|
||||
|
||||
# Get height scanner data
|
||||
height_scanner = env.scene.sensors[sensor_cfg.name]
|
||||
|
||||
# Compute height scan: sensor_height - hit_z - offset
|
||||
scan_data = height_scanner.data.pos_w[:, 2].unsqueeze(1) - height_scanner.data.ray_hits_w[..., 2] - offset
|
||||
|
||||
# Clamp the height scan data to the range [-5, 5]
|
||||
scan_data = torch.clamp(scan_data, min=-5.0, max=5.0)
|
||||
|
||||
# Expected grid size for the height scanner (51x51 = 2601 points)
|
||||
H = W = 51
|
||||
|
||||
# Reshape the height scan data to a 2D grid
|
||||
scan_data = scan_data.view(-1, H, W)
|
||||
|
||||
# Encode using the pre-trained encoder
|
||||
with torch.no_grad():
|
||||
encoded_scan = JIT_HEIGHTSCAN_FEAT_ENCODER(scan_data)
|
||||
|
||||
# Flatten and return: (batch, 64, 7, 7) -> (batch, 3136)
|
||||
return encoded_scan.view(env.num_envs, -1)
|
||||
|
||||
|
||||
def generated_actions(env: ManagerBasedRLEnv, action_name: str) -> torch.Tensor:
|
||||
"""The generated action from action term in the action manager with the given name.
|
||||
|
||||
Args:
|
||||
env: The environment object.
|
||||
action_name: The name of the action term.
|
||||
|
||||
Returns:
|
||||
The processed actions from the action term.
|
||||
"""
|
||||
return env.action_manager.get_term(action_name).processed_actions
|
||||
|
||||
|
||||
def generated_commands_reshaped_delayed(
|
||||
env: ManagerBasedRLEnv,
|
||||
command_name: str,
|
||||
unsqueeze_pos: int = 1,
|
||||
flatten: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""The generated command with delay applied.
|
||||
|
||||
Requires env.delay_manager to exist. Each environment has a random delay
|
||||
sampled from [0, max_delay] at episode reset.
|
||||
|
||||
Args:
|
||||
env: The environment object (must have delay_manager attribute).
|
||||
command_name: The name of the command term.
|
||||
unsqueeze_pos: Position to unsqueeze the command.
|
||||
flatten: Whether to flatten the command.
|
||||
|
||||
Returns:
|
||||
The delayed command.
|
||||
"""
|
||||
if flatten:
|
||||
command = env.command_manager.get_command(command_name)
|
||||
return env.delay_manager.compute_delayed_target_position(command)
|
||||
|
||||
command = env.command_manager.get_command(command_name).unsqueeze(unsqueeze_pos)
|
||||
return env.delay_manager.compute_delayed_target_position(command)
|
||||
|
||||
|
||||
def last_low_level_action(
|
||||
env: ManagerBasedEnv, action_term: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
|
||||
) -> torch.Tensor:
|
||||
"""The last low-level action.
|
||||
|
||||
Args:
|
||||
env: The environment object.
|
||||
action_term: The name of the action term.
|
||||
|
||||
Returns:
|
||||
The last low-level action.
|
||||
"""
|
||||
action_term: PerceptiveNavigationSE2Action = env.action_manager._terms[action_term]
|
||||
return action_term.low_level_actions[:, asset_cfg.joint_ids]
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def depth_image_prefect(env, sensor_cfg):
|
||||
"""Return the perfect (non-noisy) encoded depth image from the camera.
|
||||
|
||||
Args:
|
||||
env: The environment object.
|
||||
sensor_cfg: The sensor configuration.
|
||||
|
||||
Returns:
|
||||
The encoded depth image features.
|
||||
"""
|
||||
# Ensure encoder is initialized
|
||||
_ensure_depth_noise_generator_initialized()
|
||||
|
||||
depth_camera = env.scene.sensors[sensor_cfg.name]
|
||||
|
||||
# Get depth image tensor
|
||||
depth_tensor = depth_camera.data.output["distance_to_image_plane"].view(env.num_envs, -1)
|
||||
|
||||
# Handle NaN values by replacing them with 50.0
|
||||
depth_tensor = torch.nan_to_num(depth_tensor, nan=50.0, posinf=50.0, neginf=0.0)
|
||||
|
||||
# Reshape the tensor to [B, 1, H, W]
|
||||
H, W = depth_camera.image_shape
|
||||
depth_tensor = depth_tensor.view(-1, 1, H, W)
|
||||
|
||||
assert JIT_DEPTH_NOISE_GENERATOR is not None, (
|
||||
"Depth encoder JIT model is not initialized. Call initialize_depth_noise_generator first."
|
||||
)
|
||||
model = cast(torch.jit.ScriptModule, JIT_DEPTH_NOISE_GENERATOR)
|
||||
if use_jit:
|
||||
depth_tensor[depth_tensor > max_depth] = 0.0
|
||||
encoded_depth_tensor = model(depth_tensor)
|
||||
else:
|
||||
encoded_depth_tensor, _ = model(depth_tensor)
|
||||
|
||||
return encoded_depth_tensor.view(env.num_envs, -1)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def depth_image_noisy_delayed(
|
||||
env, sensor_cfg, visualize: bool = False
|
||||
):
|
||||
"""Return the noisy and delayed encoded depth image from the camera.
|
||||
|
||||
Requires env.delay_manager to exist. Each environment has a random delay
|
||||
sampled from [0, max_delay] at episode reset.
|
||||
|
||||
Args:
|
||||
env: The environment object (must have delay_manager attribute).
|
||||
sensor_cfg: The sensor configuration.
|
||||
visualize: Whether to visualize the depth image.
|
||||
|
||||
Returns:
|
||||
The delayed encoded depth image features.
|
||||
"""
|
||||
# Ensure encoder is initialized
|
||||
_ensure_depth_noise_generator_initialized()
|
||||
|
||||
depth_camera = env.scene.sensors[sensor_cfg.name]
|
||||
|
||||
# Get depth image tensor
|
||||
depth_tensor = depth_camera.data.output["distance_to_image_plane"].view(env.num_envs, -1)
|
||||
|
||||
# Handle NaN values by replacing them with 50.0
|
||||
depth_tensor = torch.nan_to_num(depth_tensor, nan=50.0, posinf=50.0, neginf=0.0)
|
||||
|
||||
# Reshape the tensor to [B, 1, H, W]
|
||||
H, W = depth_camera.image_shape
|
||||
depth_tensor = depth_tensor.view(-1, 1, H, W)
|
||||
|
||||
assert JIT_DEPTH_NOISE_GENERATOR is not None, (
|
||||
"Depth encoder JIT model is not initialized. Call initialize_depth_noise_generator first."
|
||||
)
|
||||
model = cast(torch.jit.ScriptModule, JIT_DEPTH_NOISE_GENERATOR)
|
||||
if use_jit:
|
||||
depth_tensor[depth_tensor > max_depth] = 0.0 # depth larger than depth max is invalid
|
||||
depth_tensor[depth_tensor < min_depth] = 0.0 # depth smaller than depth min is invalid
|
||||
encoded_depth_tensor = model(depth_tensor)
|
||||
noisy_depth_tensor = depth_tensor
|
||||
else:
|
||||
encoded_depth_tensor, noisy_depth_tensor = model(depth_tensor)
|
||||
|
||||
# Continuous visualization (single-env) with persistent windows
|
||||
if visualize and env.num_envs == 1 and not use_jit:
|
||||
camera_name = sensor_cfg.name if hasattr(sensor_cfg, "name") else "camera"
|
||||
_update_depth_window(camera_name, noisy_depth_tensor, H, W, title=f"Depth (Noisy Delayed) - {camera_name}")
|
||||
|
||||
encoded_depth_tensor_reshaped = encoded_depth_tensor.view(env.num_envs, -1)
|
||||
|
||||
# Apply delay using env's delay manager
|
||||
camera_name = sensor_cfg.name if hasattr(sensor_cfg, "name") else "depth"
|
||||
return env.delay_manager.compute_delayed_depth(encoded_depth_tensor_reshaped, camera_name)
|
||||
|
||||
|
||||
def in_goal(
|
||||
env: ManagerBasedRLEnv,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
distance_threshold: float = 0.5,
|
||||
goal_cmd_name: str = "robot_goal",
|
||||
) -> torch.Tensor:
|
||||
"""Check if the robot is within the goal distance threshold.
|
||||
|
||||
Args:
|
||||
env: The learning environment.
|
||||
asset_cfg: The name of the robot asset.
|
||||
distance_threshold: The distance threshold to the goal.
|
||||
goal_cmd_name: The name of the goal command.
|
||||
|
||||
Returns:
|
||||
Boolean tensor indicating whether the robot is within the goal.
|
||||
"""
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
goal_cmd_generator: RobotNavigationGoalCommand = env.command_manager._terms[goal_cmd_name]
|
||||
distance_goal = torch.norm(asset.data.root_pos_w[:, :2] - goal_cmd_generator.pos_command_w[:, :2], dim=1, p=2, keepdim=True)
|
||||
return distance_goal < distance_threshold
|
||||
|
||||
|
||||
def time_normalized(env: ManagerBasedRLEnv, command_name: str = "robot_goal") -> torch.Tensor:
|
||||
"""Time normalized to the maximum episode length.
|
||||
|
||||
Args:
|
||||
env: The learning environment.
|
||||
command_name: The name of the goal command.
|
||||
|
||||
Returns:
|
||||
The normalized time (current step / max steps).
|
||||
"""
|
||||
T_max = env.max_episode_length
|
||||
if hasattr(env, "episode_length_buf"):
|
||||
t = env.episode_length_buf.unsqueeze(-1)
|
||||
else:
|
||||
t = torch.tensor([0.0]).repeat(env.num_envs, 1).to(env.device)
|
||||
return t / T_max
|
||||
@@ -0,0 +1,134 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Reward functions for navigation tasks.
|
||||
|
||||
These functions can be passed to :class:`isaaclab.managers.RewardTermCfg`
|
||||
to specify the reward function and its parameters.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from isaaclab.assets import Articulation
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
|
||||
from isaaclab_nav_task.navigation.mdp.navigation.goal_commands import RobotNavigationGoalCommand
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedRLEnv
|
||||
|
||||
|
||||
def action_rate_l1(env: "ManagerBasedRLEnv") -> torch.Tensor:
|
||||
"""Penalize the rate of change of the actions using L1 kernel."""
|
||||
return torch.sum(torch.abs(env.action_manager.action - env.action_manager.prev_action), dim=1)
|
||||
|
||||
|
||||
def lateral_movement(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Reward the agent for moving laterally using L1-Kernel.
|
||||
|
||||
Args:
|
||||
env: The learning environment.
|
||||
asset_cfg: The name of the robot asset.
|
||||
|
||||
Returns:
|
||||
Dense reward [0, +1] based on the lateral velocity.
|
||||
"""
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
# compute the reward
|
||||
lateral_velocity = asset.data.root_lin_vel_b[:, 1]
|
||||
reward = torch.abs(lateral_velocity)
|
||||
return reward
|
||||
|
||||
|
||||
def rot_movement(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Reward the agent for rotating around the z-axis using an L2-Kernel.
|
||||
|
||||
Args:
|
||||
env: The learning environment.
|
||||
asset_cfg: The name of the robot asset.
|
||||
|
||||
Returns:
|
||||
Dense reward [0, +1] based on the rotational velocity.
|
||||
"""
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
rot_vel_norm = torch.norm(asset.data.root_ang_vel_b, dim=1)
|
||||
return rot_vel_norm
|
||||
|
||||
|
||||
def reach_goal_xyz(
|
||||
env: ManagerBasedRLEnv,
|
||||
command_name: str,
|
||||
sigmoid: float,
|
||||
T_r: float,
|
||||
probability: float,
|
||||
flat: bool,
|
||||
ratio: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Reward goal reaching with configurable sigmoid shaping.
|
||||
|
||||
Args:
|
||||
env: The learning environment.
|
||||
command_name: Name of the goal command.
|
||||
sigmoid: Sigmoid parameter for shaping.
|
||||
T_r: Time reward scaling factor.
|
||||
probability: Probability of random sampling.
|
||||
flat: Whether to only consider xy error (ignore z).
|
||||
ratio: Whether to scale by travel distance ratio.
|
||||
|
||||
Returns:
|
||||
Dense reward based on distance to goal.
|
||||
"""
|
||||
goal_cmd_generator: RobotNavigationGoalCommand = env.command_manager._terms[command_name]
|
||||
|
||||
t = env.episode_length_buf
|
||||
T = env.max_episode_length
|
||||
|
||||
if flat:
|
||||
xyz_error = torch.norm(goal_cmd_generator._get_unscaled_command()[:, :2], dim=1)
|
||||
else:
|
||||
xyz_error = torch.norm(goal_cmd_generator._get_unscaled_command(), dim=1)
|
||||
|
||||
reward = 1 / (1 + torch.square(xyz_error / sigmoid)) / T_r
|
||||
|
||||
timeup_mask = t > (T - goal_cmd_generator.required_time_at_goal_in_steps)
|
||||
random_mask = torch.rand_like(t.float()) < probability
|
||||
timeup_mask = torch.logical_or(timeup_mask, random_mask)
|
||||
|
||||
arrive_mask = goal_cmd_generator.time_at_goal > 0.0
|
||||
reward_mask = torch.logical_or(timeup_mask, arrive_mask)
|
||||
|
||||
if ratio:
|
||||
# Calculate the travel distance ratio relative to the initial goal distance
|
||||
travel_distance = torch.max(
|
||||
goal_cmd_generator.distance_traveled, goal_cmd_generator.initial_distance_to_goal
|
||||
)
|
||||
travel_distance_ratio = goal_cmd_generator.initial_distance_to_goal / (travel_distance + 1e-6)
|
||||
else:
|
||||
travel_distance_ratio = torch.ones_like(reward)
|
||||
|
||||
reward = reward * reward_mask.float() * travel_distance_ratio
|
||||
|
||||
return reward
|
||||
|
||||
|
||||
def backward_movement_penalty(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Small penalty for backward movement as a regularization term.
|
||||
|
||||
Args:
|
||||
env: The learning environment.
|
||||
asset_cfg: The name of the robot asset.
|
||||
|
||||
Returns:
|
||||
Penalty [0, +1] based on backward velocity (to be used with negative weight).
|
||||
"""
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
# compute the penalty
|
||||
forward_velocity = asset.data.root_lin_vel_b[:, 0]
|
||||
# Only penalize negative forward velocity (backward movement)
|
||||
backward_velocity = torch.clamp(-forward_velocity, min=0.0, max=1.0)
|
||||
return backward_velocity
|
||||
@@ -0,0 +1,211 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Common functions that can be used to activate certain terminations.
|
||||
|
||||
The functions can be passed to the :class:`isaaclab.managers.TerminationTermCfg` object to enable
|
||||
the termination introduced by the function.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from isaaclab.assets import Articulation
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
from isaaclab.sensors import ContactSensor
|
||||
from isaaclab.utils.math import quat_inv, yaw_quat, quat_mul, euler_xyz_from_quat
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedRLEnv
|
||||
from isaaclab_nav_task.navigation.mdp.navigation.goal_commands import RobotNavigationGoalCommand
|
||||
|
||||
|
||||
def euler_xyz_from_quat_wrapped(quat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Convert quaternion to Euler angles (XYZ convention) with wrapping to [-pi, pi].
|
||||
|
||||
Args:
|
||||
quat: Quaternion tensor of shape (..., 4) in (w, x, y, z) format.
|
||||
|
||||
Returns:
|
||||
Tuple of (roll, pitch, yaw) tensors.
|
||||
"""
|
||||
roll, pitch, yaw = euler_xyz_from_quat(quat)
|
||||
# Wrap to [-pi, pi]
|
||||
roll = torch.remainder(roll + torch.pi, 2 * torch.pi) - torch.pi
|
||||
pitch = torch.remainder(pitch + torch.pi, 2 * torch.pi) - torch.pi
|
||||
yaw = torch.remainder(yaw + torch.pi, 2 * torch.pi) - torch.pi
|
||||
return roll, pitch, yaw
|
||||
|
||||
|
||||
def time_out_navigation(
|
||||
env: "ManagerBasedRLEnv",
|
||||
goal_cmd_name: str = "robot_goal",
|
||||
distance_threshold: float = 0.5
|
||||
) -> torch.Tensor:
|
||||
"""Terminate the episode when the episode length exceeds the maximum episode length.
|
||||
|
||||
This also tracks success metrics by checking if the robot reached the goal before timeout.
|
||||
"""
|
||||
from isaaclab_nav_task.navigation.mdp.navigation.goal_commands import RobotNavigationGoalCommand
|
||||
|
||||
goal_cmd_generator: RobotNavigationGoalCommand = env.command_manager._terms[goal_cmd_name]
|
||||
|
||||
termination = env.episode_length_buf >= env.max_episode_length
|
||||
|
||||
env_ids = torch.where(termination)[0]
|
||||
|
||||
distance_goal = torch.norm(goal_cmd_generator._get_unscaled_command()[:, :2], dim=1)
|
||||
|
||||
# update time at goal
|
||||
goal_cmd_generator.time_at_goal[distance_goal < distance_threshold] += 1 * env.step_dt
|
||||
|
||||
if env_ids.numel() > 0: # Check if env_ids is not empty
|
||||
success_masks = goal_cmd_generator.time_at_goal > 0.0
|
||||
value_buffer = torch.zeros_like(distance_goal) # init with 0: Fail
|
||||
value_buffer[success_masks] = 1.0 # Success
|
||||
goal_cmd_generator.goal_reached_buffer.add(value_buffer, env_ids)
|
||||
|
||||
return termination
|
||||
|
||||
|
||||
def illegal_contact_navigation(
|
||||
env: "ManagerBasedRLEnv",
|
||||
threshold: float,
|
||||
sensor_cfg: SceneEntityCfg,
|
||||
goal_cmd_name: str = "robot_goal",
|
||||
) -> torch.Tensor:
|
||||
"""Terminate when the contact force on the sensor exceeds the force threshold."""
|
||||
from isaaclab_nav_task.navigation.mdp.navigation.goal_commands import RobotNavigationGoalCommand
|
||||
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
net_contact_forces = contact_sensor.data.net_forces_w_history
|
||||
goal_cmd_generator: RobotNavigationGoalCommand = env.command_manager._terms[goal_cmd_name]
|
||||
|
||||
termination = torch.any(
|
||||
torch.max(torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold, dim=1
|
||||
)
|
||||
|
||||
env_ids = torch.where(termination)[0]
|
||||
|
||||
if env_ids.numel() > 0: # Check if env_ids is not empty
|
||||
goal_cmd_generator.goal_reached_buffer.add(torch.zeros_like(termination, dtype=torch.float), env_ids)
|
||||
|
||||
return termination
|
||||
|
||||
|
||||
def large_angle_termination_navigation(
|
||||
env: "ManagerBasedRLEnv",
|
||||
threshold: float,
|
||||
goal_cmd_name: str = "robot_goal",
|
||||
) -> torch.Tensor:
|
||||
"""Terminate when the robot exceeds a pitch or roll angle threshold."""
|
||||
from isaaclab_nav_task.navigation.mdp.navigation.goal_commands import RobotNavigationGoalCommand
|
||||
|
||||
goal_cmd_generator: RobotNavigationGoalCommand = env.command_manager._terms[goal_cmd_name]
|
||||
|
||||
# degree to rad
|
||||
threshold_rad = threshold * torch.pi / 180.0
|
||||
|
||||
robot = env.scene["robot"]
|
||||
yaw_q = yaw_quat(robot.data.root_quat_w)
|
||||
base_quat_b = quat_mul(quat_inv(yaw_q), robot.data.root_quat_w)
|
||||
robot_roll, robot_pitch, _ = euler_xyz_from_quat_wrapped(base_quat_b)
|
||||
|
||||
termination = torch.logical_or(torch.abs(robot_pitch) > threshold_rad, torch.abs(robot_roll) > threshold_rad)
|
||||
|
||||
env_ids = torch.where(termination)[0]
|
||||
|
||||
if env_ids.numel() > 0: # Check if env_ids is not empty
|
||||
goal_cmd_generator.goal_reached_buffer.add(torch.zeros_like(termination, dtype=torch.float), env_ids)
|
||||
|
||||
return termination
|
||||
|
||||
|
||||
def at_goal_navigation(
|
||||
env: "ManagerBasedRLEnv",
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
distance_threshold: float = 0.5,
|
||||
goal_cmd_name: str = "robot_goal",
|
||||
) -> torch.Tensor:
|
||||
"""Terminate the episode when the goal is reached.
|
||||
|
||||
Args:
|
||||
env: The learning environment.
|
||||
asset_cfg: The name of the robot asset.
|
||||
distance_threshold: The distance threshold to the goal.
|
||||
goal_cmd_name: The name of the goal command term.
|
||||
|
||||
Returns:
|
||||
Boolean tensor indicating whether the goal is reached.
|
||||
"""
|
||||
from isaaclab_nav_task.navigation.mdp.navigation.goal_commands import RobotNavigationGoalCommand
|
||||
|
||||
# Extract the used quantities
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
goal_cmd_generator: RobotNavigationGoalCommand = env.command_manager._terms.get(goal_cmd_name)
|
||||
|
||||
# Calculate distance to goal
|
||||
xy_error = torch.norm(goal_cmd_generator._get_unscaled_command()[:, :2], dim=1)
|
||||
|
||||
# Check conditions for termination
|
||||
at_goal = xy_error < distance_threshold
|
||||
|
||||
# already at goal
|
||||
already_at_goal = goal_cmd_generator.time_at_goal > 0.0
|
||||
at_goal = torch.logical_or(at_goal, already_at_goal)
|
||||
|
||||
# Update the time at goal in steps
|
||||
goal_cmd_generator.time_at_goal_in_steps[at_goal] += 1 # Increment if at goal
|
||||
|
||||
# Determine if termination condition is met
|
||||
termination = goal_cmd_generator.time_at_goal_in_steps > goal_cmd_generator.required_time_at_goal_in_steps
|
||||
|
||||
# Update goal reached buffer if termination condition is met
|
||||
env_ids = torch.where(termination)[0]
|
||||
if env_ids.numel() > 0: # Check if any environments have met the termination condition
|
||||
goal_cmd_generator.goal_reached_buffer.add(torch.ones_like(termination, dtype=torch.float), env_ids)
|
||||
|
||||
return termination
|
||||
|
||||
|
||||
def terrain_fall(
|
||||
env: "ManagerBasedRLEnv",
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
fall_height_threshold: float = -1.0,
|
||||
goal_cmd_name: str = "robot_goal",
|
||||
) -> torch.Tensor:
|
||||
"""Terminate when the robot falls below a certain height threshold.
|
||||
|
||||
This termination is triggered when the robot's z-position falls below a
|
||||
specified threshold, indicating that the robot has fallen off the terrain
|
||||
or into a deep pit.
|
||||
|
||||
Args:
|
||||
env: The learning environment.
|
||||
asset_cfg: The configuration for the robot asset.
|
||||
fall_height_threshold: The z-height below which the robot is considered fallen (in meters).
|
||||
Default is -1.0m to account for pits which can be ~1.5m deep.
|
||||
goal_cmd_name: The name of the goal command term.
|
||||
|
||||
Returns:
|
||||
Boolean tensor indicating whether the robot has fallen.
|
||||
"""
|
||||
# Direct tensor access for z-coordinate (avoids intermediate variable allocation)
|
||||
termination = env.scene[asset_cfg.name].data.root_pos_w[:, 2] < fall_height_threshold
|
||||
|
||||
# Early exit if no terminations (common case - avoids torch.where overhead)
|
||||
if not termination.any():
|
||||
return termination
|
||||
|
||||
# Update goal reached buffer with failure
|
||||
goal_cmd = env.command_manager._terms.get(goal_cmd_name)
|
||||
if goal_cmd is not None:
|
||||
env_ids = termination.nonzero(as_tuple=False).squeeze(-1)
|
||||
goal_cmd.goal_reached_buffer.add(torch.zeros(env.num_envs, dtype=torch.float, device=env.device), env_ids)
|
||||
|
||||
return termination
|
||||
@@ -0,0 +1,44 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Custom navigation environment with observation delay support."""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from isaaclab.envs import ManagerBasedRLEnv
|
||||
|
||||
from isaaclab_nav_task.navigation.mdp.delay_manager import ObservationDelayManager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab_nav_task.navigation.navigation_env_cfg import NavigationEnvCfg
|
||||
|
||||
|
||||
class NavigationEnv(ManagerBasedRLEnv):
|
||||
"""Navigation environment with observation delay manager.
|
||||
|
||||
This environment extends the base ManagerBasedRLEnv to add support for
|
||||
simulating sensor delays. The delay_manager is created before managers
|
||||
are loaded, ensuring observations can access it during initialization.
|
||||
"""
|
||||
|
||||
cfg: NavigationEnvCfg
|
||||
|
||||
def load_managers(self):
|
||||
"""Load managers with delay manager initialization.
|
||||
|
||||
The delay_manager is created before other managers are loaded,
|
||||
ensuring observations can access it during initialization.
|
||||
"""
|
||||
# Create the delay manager before loading other managers
|
||||
# At this point, self.num_envs and self.device are available via scene
|
||||
self.delay_manager = ObservationDelayManager(
|
||||
cfg=self.cfg.delay_cfg,
|
||||
num_envs=self.num_envs,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
# Now load all other managers
|
||||
super().load_managers()
|
||||
@@ -0,0 +1,445 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Navigation environment configuration.
|
||||
|
||||
This module provides the base environment configuration for navigation tasks
|
||||
with visual perception using depth cameras.
|
||||
"""
|
||||
|
||||
import math
|
||||
import os
|
||||
from dataclasses import MISSING
|
||||
|
||||
import isaaclab.sim as sim_utils
|
||||
from isaaclab.assets import ArticulationCfg, AssetBaseCfg
|
||||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||||
from isaaclab.managers import CurriculumTermCfg as CurrTerm
|
||||
from isaaclab.managers import EventTermCfg as EventTerm
|
||||
from isaaclab.managers import ObservationGroupCfg as ObsGroup
|
||||
from isaaclab.managers import ObservationTermCfg as ObsTerm
|
||||
from isaaclab.managers import RewardTermCfg as RewTerm
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
from isaaclab.managers import TerminationTermCfg as DoneTerm
|
||||
from isaaclab.scene import InteractiveSceneCfg
|
||||
from isaaclab.sensors import ContactSensorCfg, RayCasterCfg, patterns, RayCasterCameraCfg
|
||||
from isaaclab.terrains import TerrainImporterCfg
|
||||
from isaaclab.utils import configclass
|
||||
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR
|
||||
from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise
|
||||
|
||||
import isaaclab_nav_task.navigation.mdp as mdp
|
||||
from isaaclab_nav_task.navigation.mdp.custom_noise import DeltaTransformationNoiseCfg
|
||||
from isaaclab_nav_task.navigation.mdp.delay_manager import ObservationDelayManagerCfg
|
||||
from isaaclab_nav_task.navigation.assets import ISAACLAB_NAV_TASKS_ASSETS_DIR
|
||||
|
||||
##
|
||||
# Pre-defined configs
|
||||
##
|
||||
from isaaclab_nav_task.terrains import MAZE_TERRAIN_CFG # isort: skip
|
||||
|
||||
# Constants
|
||||
LEG_JOINT_NAMES = [".*HAA", ".*HFE", ".*KFE"]
|
||||
LEG_BODY_NAMES = [".*HIP", ".*THIGH", ".*SHANK"]
|
||||
WHEEL_JOINT_NAMES = [".*WHEEL"]
|
||||
WHEEL_BODY_NAMES = [".*WHEEL_L"]
|
||||
PLANNING_FREQ = 5.0 # Hz
|
||||
|
||||
##
|
||||
# Scene definition
|
||||
##
|
||||
|
||||
|
||||
@configclass
|
||||
class MySceneCfg(InteractiveSceneCfg):
|
||||
"""Configuration for the terrain scene with a legged robot."""
|
||||
|
||||
# ground terrain
|
||||
terrain = TerrainImporterCfg(
|
||||
prim_path="/World/ground",
|
||||
terrain_type="generator",
|
||||
terrain_generator=MAZE_TERRAIN_CFG,
|
||||
max_init_terrain_level=0,
|
||||
collision_group=-1,
|
||||
physics_material=sim_utils.RigidBodyMaterialCfg(
|
||||
friction_combine_mode="multiply",
|
||||
restitution_combine_mode="multiply",
|
||||
restitution=0.1,
|
||||
static_friction=1.0,
|
||||
dynamic_friction=0.8,
|
||||
compliant_contact_stiffness=5e5,
|
||||
compliant_contact_damping=300.0,
|
||||
),
|
||||
visual_material=sim_utils.MdlFileCfg(
|
||||
mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/TilesMarbleSpiderWhiteBrickBondHoned.mdl",
|
||||
project_uvw=True,
|
||||
texture_scale=(0.25, 0.25),
|
||||
),
|
||||
debug_vis=False,
|
||||
)
|
||||
# robots
|
||||
robot: ArticulationCfg = MISSING
|
||||
|
||||
raycast_camera = RayCasterCameraCfg(
|
||||
prim_path="{ENV_REGEX_NS}/Robot/base",
|
||||
mesh_prim_paths=["/World/ground"],
|
||||
update_period=0,
|
||||
offset=RayCasterCameraCfg.OffsetCfg(
|
||||
pos=(0.4761, 0.0035, 0.1055), rot=(0.9848078, 0.0, 0.1736482, 0.0), convention="world" # 20 degrees
|
||||
),
|
||||
data_types=["distance_to_image_plane"],
|
||||
debug_vis=False,
|
||||
max_distance=11.0,
|
||||
pattern_cfg=patterns.PinholeCameraPatternCfg.from_ros_camera_info(
|
||||
# ZED camera parameters from ROS camera_info topic
|
||||
fx=72.7025,
|
||||
fy=72.7025,
|
||||
cx=94.4457,
|
||||
cy=62.5424,
|
||||
width=192,
|
||||
height=120,
|
||||
downsample_factor=3, # Downsample from 192x120 to 64x40
|
||||
),
|
||||
)
|
||||
|
||||
height_scanner_critic = RayCasterCfg(
|
||||
prim_path="{ENV_REGEX_NS}/Robot/base",
|
||||
offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
|
||||
# attach_yaw_only=True,
|
||||
ray_alignment='yaw',
|
||||
pattern_cfg=patterns.GridPatternCfg(resolution=0.2, size=[10.0, 10.0]),
|
||||
debug_vis=False,
|
||||
mesh_prim_paths=["/World/ground"],
|
||||
)
|
||||
|
||||
contact_forces = ContactSensorCfg(prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True)
|
||||
# lights
|
||||
sky_light = AssetBaseCfg(
|
||||
prim_path="/World/skyLight",
|
||||
spawn=sim_utils.DomeLightCfg(
|
||||
intensity=750.0,
|
||||
texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
##
|
||||
# MDP settings
|
||||
##
|
||||
|
||||
|
||||
@configclass
|
||||
class CommandsCfg:
|
||||
"""Command specifications for the MDP."""
|
||||
|
||||
robot_goal = mdp.RobotNavigationGoalCommandCfg(
|
||||
asset_name="robot",
|
||||
# Large value to disable automatic resampling - goals only change on episode reset
|
||||
# Note: math.inf doesn't work with PyTorch's uniform_(), so we use 1e9 (~31 years)
|
||||
resampling_time_range=(1e9, 1e9),
|
||||
debug_vis=True,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class ActionsCfg:
|
||||
"""Action specifications for the MDP."""
|
||||
|
||||
velocity_command = mdp.PerceptiveNavigationSE2ActionCfg(
|
||||
asset_name="robot",
|
||||
low_level_position_action=mdp.JointPositionActionCfg(
|
||||
asset_name="robot", joint_names=[".*HAA", ".*HFE", ".*KFE"], scale=0.5, use_default_offset=True
|
||||
),
|
||||
low_level_velocity_action=mdp.JointVelocityActionCfg(
|
||||
asset_name="robot", joint_names=[".*WHEEL"], scale=5.0, use_default_offset=True
|
||||
),
|
||||
low_level_decimation=4,
|
||||
low_level_policy_file=os.path.join(
|
||||
ISAACLAB_NAV_TASKS_ASSETS_DIR, "Policies", "locomotion", "aow_d", "policy_blind_3_1.pt"
|
||||
),
|
||||
observation_group="low_level_policy",
|
||||
policy_scaling=[1.5, 1.0, 1.0],
|
||||
use_raw_actions=True,
|
||||
policy_distr_type="gaussian",
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class ObservationsCfg:
|
||||
"""Observation specifications for the MDP."""
|
||||
|
||||
@configclass
|
||||
class PolicyCfg(ObsGroup):
|
||||
"""Observations for policy group."""
|
||||
|
||||
base_lin_vel = ObsTerm(
|
||||
func=mdp.base_lin_vel_delayed, noise=Unoise(n_min=-0.2, n_max=0.2)
|
||||
)
|
||||
base_ang_vel = ObsTerm(
|
||||
func=mdp.base_ang_vel_delayed, noise=Unoise(n_min=-0.1, n_max=0.1)
|
||||
)
|
||||
projected_gravity = ObsTerm(
|
||||
func=mdp.projected_gravity_delayed, noise=Unoise(n_min=-0.1, n_max=0.1)
|
||||
)
|
||||
last_action = ObsTerm(func=mdp.last_action)
|
||||
target_position = ObsTerm(
|
||||
func=mdp.generated_commands_reshaped_delayed,
|
||||
params={"command_name": "robot_goal", "flatten": True},
|
||||
noise=DeltaTransformationNoiseCfg(rotation=0.1, translation=0.5, noise_prob=0.1, remove_dist=False),
|
||||
)
|
||||
depth_image = ObsTerm(
|
||||
func=mdp.depth_image_noisy_delayed, params={"sensor_cfg": SceneEntityCfg("raycast_camera")}
|
||||
)
|
||||
|
||||
def __post_init__(self):
|
||||
self.enable_corruption = True
|
||||
self.concatenate_terms = True
|
||||
|
||||
@configclass
|
||||
class CriticCfg(ObsGroup):
|
||||
"""Observations for critic group."""
|
||||
|
||||
base_lin_vel = ObsTerm(func=mdp.base_lin_vel)
|
||||
base_ang_vel = ObsTerm(func=mdp.base_ang_vel)
|
||||
projected_gravity = ObsTerm(func=mdp.projected_gravity)
|
||||
last_action = ObsTerm(func=mdp.last_action)
|
||||
target_position = ObsTerm(
|
||||
func=mdp.generated_commands_reshaped, params={"command_name": "robot_goal", "flatten": True}
|
||||
)
|
||||
time_normalized = ObsTerm(func=mdp.time_normalized, params={"command_name": "robot_goal"})
|
||||
height_scan_critic = ObsTerm(
|
||||
func=mdp.height_scan_feat, params={"sensor_cfg": SceneEntityCfg("height_scanner_critic")}
|
||||
)
|
||||
depth_image = ObsTerm(func=mdp.depth_image_prefect, params={"sensor_cfg": SceneEntityCfg("raycast_camera")})
|
||||
|
||||
def __post_init__(self):
|
||||
self.enable_corruption = True
|
||||
self.concatenate_terms = True
|
||||
|
||||
@configclass
|
||||
class LowLevelPolicyCfg(ObsGroup):
|
||||
"""Observations for low-level policy."""
|
||||
|
||||
base_lin_vel = ObsTerm(func=mdp.base_lin_vel, noise=Unoise(n_min=-0.2, n_max=0.2))
|
||||
base_ang_vel = ObsTerm(func=mdp.base_ang_vel, noise=Unoise(n_min=-0.1, n_max=0.1))
|
||||
projected_gravity = ObsTerm(func=mdp.projected_gravity, noise=Unoise(n_min=-0.1, n_max=0.1))
|
||||
velocity_commands = ObsTerm(func=mdp.generated_actions, params={"action_name": "velocity_command"})
|
||||
joint_pos = ObsTerm(func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.2, n_max=0.2))
|
||||
joint_vel = ObsTerm(func=mdp.joint_vel_rel, noise=Unoise(n_min=-1.5, n_max=1.5))
|
||||
actions = ObsTerm(func=mdp.last_low_level_action, params={"action_term": "velocity_command"})
|
||||
|
||||
def __post_init__(self):
|
||||
self.enable_corruption = True
|
||||
self.concatenate_terms = True
|
||||
|
||||
@configclass
|
||||
class MetricsCfg(ObsGroup):
|
||||
"""Observations for metrics tracking."""
|
||||
|
||||
in_goal = ObsTerm(func=mdp.in_goal)
|
||||
|
||||
def __post_init__(self):
|
||||
self.enable_corruption = False
|
||||
self.concatenate_terms = False
|
||||
|
||||
# Observation groups
|
||||
metrics: MetricsCfg = MetricsCfg()
|
||||
policy: PolicyCfg = PolicyCfg()
|
||||
critic: CriticCfg = CriticCfg()
|
||||
low_level_policy: LowLevelPolicyCfg = LowLevelPolicyCfg()
|
||||
|
||||
@configclass
|
||||
class EventCfg:
|
||||
"""Configuration for events."""
|
||||
|
||||
# Startup events
|
||||
physics_material = EventTerm(
|
||||
func=mdp.randomize_rigid_body_material,
|
||||
mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=".*"),
|
||||
"static_friction_range": (0.8, 1.2),
|
||||
"dynamic_friction_range": (0.7, 1.0),
|
||||
"restitution_range": (0.0, 0.1),
|
||||
"num_buckets": 64,
|
||||
},
|
||||
)
|
||||
|
||||
# Reset events
|
||||
reset_base = EventTerm(
|
||||
func=mdp.reset_root_state_uniform,
|
||||
mode="reset",
|
||||
params={
|
||||
"pose_range": {"x": (-0.0, 0.0), "y": (-0.0, 0.0), "yaw": (-math.pi, math.pi)},
|
||||
"velocity_range": {"x": (-0.0, 0.0), "y": (-0.0, 0.0), "z": (-0.0, 0.0)},
|
||||
},
|
||||
)
|
||||
|
||||
reset_robot_joints = EventTerm(
|
||||
func=mdp.reset_joints_by_scale,
|
||||
mode="reset",
|
||||
params={"position_range": (0.5, 1.5), "velocity_range": (0.0, 0.0)},
|
||||
)
|
||||
|
||||
randomize_camera_height = EventTerm(
|
||||
func=mdp.randomize_camera_height,
|
||||
mode="reset",
|
||||
params={
|
||||
"height_range": 0.1,
|
||||
"roll_angle_range": 5.0,
|
||||
"pitch_angle_range": 5.0,
|
||||
"yaw_angle_range": 5.0,
|
||||
"asset_cfg": SceneEntityCfg("raycast_camera"),
|
||||
},
|
||||
)
|
||||
|
||||
randomize_action_scale = EventTerm(
|
||||
func=mdp.randomize_action_scale,
|
||||
mode="reset",
|
||||
params={
|
||||
"scale_range_x": (0.8, 1.2),
|
||||
"scale_range_y": (0.6, 1.0),
|
||||
"scale_range_theta": (0.8, 1.2),
|
||||
"scale_range_xb": 0.1,
|
||||
"scale_range_yb": 0.2,
|
||||
"scale_range_thetab": 0.1,
|
||||
"action_term": "velocity_command",
|
||||
},
|
||||
)
|
||||
|
||||
reset_delay_buffer = EventTerm(
|
||||
func=mdp.reset_and_randomize_delay_buffer,
|
||||
mode="reset",
|
||||
)
|
||||
|
||||
randomize_low_pass_filter_alpha = EventTerm(
|
||||
func=mdp.randomize_low_pass_filter_alpha,
|
||||
mode="reset",
|
||||
params={
|
||||
"alpha_range": (0.4, 0.9),
|
||||
"action_term": "velocity_command",
|
||||
"per_dimension": True,
|
||||
"alpha_range_vx": (0.4, 0.9),
|
||||
"alpha_range_vy": (0.4, 0.9),
|
||||
"alpha_range_omega": (0.4, 0.9),
|
||||
},
|
||||
)
|
||||
|
||||
# Interval events
|
||||
push_robot = EventTerm(
|
||||
func=mdp.push_by_setting_velocity,
|
||||
mode="interval",
|
||||
interval_range_s=(0.2, 5.0),
|
||||
params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}},
|
||||
)
|
||||
|
||||
@configclass
|
||||
class RewardsCfg:
|
||||
"""Reward terms for the MDP."""
|
||||
|
||||
# Penalties
|
||||
joint_acc_l2_joint = RewTerm(
|
||||
func=mdp.joint_acc_l2,
|
||||
weight=-1e-7,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=LEG_JOINT_NAMES + WHEEL_JOINT_NAMES)},
|
||||
)
|
||||
|
||||
lateral_movement = RewTerm(func=mdp.lateral_movement, weight=-0.1)
|
||||
rot_movement = RewTerm(func=mdp.rot_movement, weight=-1e-5)
|
||||
action_rate_l1 = RewTerm(func=mdp.action_rate_l1, weight=-0.1)
|
||||
episode_termination = RewTerm(func=mdp.is_terminated, weight=-50.0)
|
||||
|
||||
# Goal rewards
|
||||
reach_goal_xy_soft = RewTerm(
|
||||
func=mdp.reach_goal_xyz,
|
||||
weight=0.25,
|
||||
params={"command_name": "robot_goal", "sigmoid": 2.5, "T_r": 1.0, "probability": 0.01, "flat": False, "ratio": False},
|
||||
)
|
||||
reach_goal_xy_tight = RewTerm(
|
||||
func=mdp.reach_goal_xyz,
|
||||
weight=1.5,
|
||||
params={"command_name": "robot_goal", "sigmoid": 0.25, "T_r": 0.1, "probability": 0.01, "flat": True, "ratio": False},
|
||||
)
|
||||
|
||||
# Backward movement penalty (disabled by default, can be enabled via curriculum)
|
||||
backward_movement_penalty = RewTerm(func=mdp.backward_movement_penalty, weight=-0.0)
|
||||
|
||||
|
||||
@configclass
|
||||
class TerminationsCfg:
|
||||
"""Termination terms for the MDP."""
|
||||
|
||||
time_out = DoneTerm(func=mdp.time_out_navigation, time_out=True, params={"distance_threshold": 0.5})
|
||||
base_contact = DoneTerm(
|
||||
func=mdp.illegal_contact_navigation,
|
||||
params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=["base", ".*HIP", ".*THIGH"]), "threshold": 1.0},
|
||||
)
|
||||
large_pitch_angle = DoneTerm(func=mdp.large_angle_termination_navigation, params={"threshold": 40})
|
||||
early_termination = DoneTerm(func=mdp.at_goal_navigation, time_out=True, params={"distance_threshold": 0.5})
|
||||
# Terrain fall termination (robot fell off terrain or into deep pit)
|
||||
terrain_fall = DoneTerm(
|
||||
func=mdp.terrain_fall,
|
||||
time_out=True,
|
||||
params={"fall_height_threshold": -2.0},
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class CurriculumCfg:
|
||||
"""Curriculum terms for the MDP."""
|
||||
|
||||
disable_backward_penalty = CurrTerm(
|
||||
func=mdp.disable_backward_penalty_after_steps,
|
||||
params={"disable_after_steps": 500, "action_term": "velocity_command"},
|
||||
)
|
||||
|
||||
##
|
||||
# Environment configuration
|
||||
##
|
||||
|
||||
|
||||
@configclass
|
||||
class NavigationEnvCfg(ManagerBasedRLEnvCfg):
|
||||
"""Configuration for the navigation environment with velocity-tracking."""
|
||||
|
||||
# Scene settings
|
||||
scene: MySceneCfg = MySceneCfg(num_envs=2048, env_spacing=2.5, replicate_physics=False)
|
||||
# Basic settings
|
||||
observations: ObservationsCfg = ObservationsCfg()
|
||||
actions: ActionsCfg = ActionsCfg()
|
||||
commands: CommandsCfg = CommandsCfg()
|
||||
# MDP settings
|
||||
rewards: RewardsCfg = RewardsCfg()
|
||||
terminations: TerminationsCfg = TerminationsCfg()
|
||||
events: EventCfg = EventCfg()
|
||||
curriculum: CurriculumCfg = CurriculumCfg()
|
||||
# Observation delay configuration
|
||||
delay_cfg: ObservationDelayManagerCfg = ObservationDelayManagerCfg()
|
||||
|
||||
def __post_init__(self):
|
||||
"""Post initialization."""
|
||||
# Simulation settings: low_level_policy runs at 50Hz
|
||||
self.sim.dt = 0.005
|
||||
self.is_finite_horizon = True
|
||||
self.low_level_decimation = 4
|
||||
self.decimation = int((1 / self.sim.dt) / PLANNING_FREQ)
|
||||
self.episode_length_s = 60.0
|
||||
self.sim.render_interval = self.low_level_decimation
|
||||
self.sim.disable_contact_processing = True
|
||||
self.sim.physics_material = self.scene.terrain.physics_material
|
||||
|
||||
# Update sensor periods
|
||||
if self.scene.height_scanner_critic is not None:
|
||||
self.scene.height_scanner_critic.update_period = self.decimation * self.sim.dt
|
||||
if self.scene.raycast_camera is not None:
|
||||
self.scene.raycast_camera.update_period = self.decimation * self.sim.dt
|
||||
if self.scene.contact_forces is not None:
|
||||
self.scene.contact_forces.update_period = self.sim.dt
|
||||
self.scene.contact_forces.history_length = self.decimation
|
||||
|
||||
# Terrain curriculum settings
|
||||
self.scene.terrain.max_init_terrain_level = 10
|
||||
self.scene.terrain.terrain_generator.difficulty_range = [0.5, 1.0]
|
||||
self.scene.terrain.terrain_generator.curriculum = False
|
||||
@@ -0,0 +1,63 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Maze terrain generation for navigation tasks.
|
||||
|
||||
This module provides:
|
||||
- Maze terrain generation with various obstacle types
|
||||
- Stair/platform terrain for elevation training
|
||||
- Pit terrain with negative obstacles
|
||||
- Monkey-patches for Isaac Lab terrain system integration
|
||||
- Optimized mesh generation for reduced GPU memory
|
||||
|
||||
Key files:
|
||||
- terrain_constants.py: Centralized height values and thresholds
|
||||
- hf_terrains_maze.py: Terrain generation functions
|
||||
- hf_terrains_maze_cfg.py: Terrain configurations
|
||||
- patches.py: Isaac Lab integration patches
|
||||
"""
|
||||
|
||||
# Patches first (no Isaac Sim required)
|
||||
from .patches import apply_terrain_patches
|
||||
|
||||
# Constants (no Isaac Sim required)
|
||||
from .terrain_constants import (
|
||||
HEIGHTS,
|
||||
THRESHOLDS,
|
||||
PADDING,
|
||||
STAIRS,
|
||||
HeightValues,
|
||||
SamplingThresholds,
|
||||
VERTICAL_SCALE,
|
||||
HORIZONTAL_SCALE,
|
||||
CELL_SIZE,
|
||||
CELL_PIXELS,
|
||||
)
|
||||
|
||||
# Terrain generation (requires Isaac Sim)
|
||||
from .hf_terrains_maze_cfg import HfMazeTerrainCfg
|
||||
from .hf_terrains_maze import maze_terrain
|
||||
from .maze_config import MAZE_TERRAIN_CFG
|
||||
|
||||
__all__ = [
|
||||
# Patches
|
||||
"apply_terrain_patches",
|
||||
# Constants
|
||||
"HEIGHTS",
|
||||
"THRESHOLDS",
|
||||
"PADDING",
|
||||
"STAIRS",
|
||||
"HeightValues",
|
||||
"SamplingThresholds",
|
||||
"VERTICAL_SCALE",
|
||||
"HORIZONTAL_SCALE",
|
||||
"CELL_SIZE",
|
||||
"CELL_PIXELS",
|
||||
# Terrain configs
|
||||
"HfMazeTerrainCfg",
|
||||
"MAZE_TERRAIN_CFG",
|
||||
# Terrain functions
|
||||
"maze_terrain",
|
||||
]
|
||||
@@ -0,0 +1,714 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Maze terrain generation for navigation tasks.
|
||||
|
||||
This module generates terrain height fields with explicit valid position masks.
|
||||
The key simplification is that terrain generation directly outputs:
|
||||
- `heights`: Actual terrain heights for rendering/physics
|
||||
- `valid_mask`: Boolean mask of valid goal/spawn positions
|
||||
|
||||
This eliminates the need for complex height-based classification in goal sampling.
|
||||
|
||||
The terrain data is stored on the config during generation, then picked up by
|
||||
the patches system and stored on TerrainImporter for access via:
|
||||
- self.env.scene.terrain._height_field_visual
|
||||
- self.env.scene.terrain._height_field_valid_mask
|
||||
- self.env.scene.terrain._height_field_platform_mask
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
from dataclasses import dataclass, field
|
||||
from scipy.ndimage import binary_dilation, rotate, shift
|
||||
from typing import TYPE_CHECKING, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from isaaclab.terrains.height_field.utils import height_field_to_mesh
|
||||
|
||||
from .terrain_constants import HEIGHTS, PADDING, STAIRS, OBSTACLES, ObstacleType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from . import hf_terrains_maze_cfg
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Terrain Data Container
|
||||
# =============================================================================
|
||||
|
||||
@dataclass
|
||||
class TerrainData:
|
||||
"""Container for terrain height field and valid position mask.
|
||||
|
||||
Attributes:
|
||||
heights: Height field for rendering/physics (actual terrain heights).
|
||||
valid_mask: Boolean mask where True = valid for goals/spawns.
|
||||
platform_mask: Boolean mask where True = elevated platform (for curriculum).
|
||||
"""
|
||||
heights: np.ndarray
|
||||
valid_mask: np.ndarray
|
||||
platform_mask: np.ndarray = field(default_factory=lambda: np.array([]))
|
||||
|
||||
@classmethod
|
||||
def create(cls, width: int, height: int) -> "TerrainData":
|
||||
"""Create empty terrain data with ground-level heights."""
|
||||
return cls(
|
||||
heights=np.zeros((width, height), dtype=np.int16),
|
||||
valid_mask=np.ones((width, height), dtype=bool), # Start all valid
|
||||
platform_mask=np.zeros((width, height), dtype=bool),
|
||||
)
|
||||
|
||||
def set_obstacle(
|
||||
self,
|
||||
x_start: int, x_end: int,
|
||||
y_start: int, y_end: int,
|
||||
height_value: int
|
||||
):
|
||||
"""Set a region as an obstacle (invalid for goals)."""
|
||||
self.heights[x_start:x_end, y_start:y_end] = height_value
|
||||
self.valid_mask[x_start:x_end, y_start:y_end] = False
|
||||
|
||||
def set_platform(
|
||||
self,
|
||||
x_start: int, x_end: int,
|
||||
y_start: int, y_end: int,
|
||||
height_value: int
|
||||
):
|
||||
"""Set a region as a platform (valid for goals, elevated)."""
|
||||
self.heights[x_start:x_end, y_start:y_end] = height_value
|
||||
self.valid_mask[x_start:x_end, y_start:y_end] = True
|
||||
self.platform_mask[x_start:x_end, y_start:y_end] = True
|
||||
|
||||
def set_ground(self, x_start: int, x_end: int, y_start: int, y_end: int):
|
||||
"""Set a region as flat ground (valid for goals)."""
|
||||
self.heights[x_start:x_end, y_start:y_end] = HEIGHTS.GROUND
|
||||
self.valid_mask[x_start:x_end, y_start:y_end] = True
|
||||
self.platform_mask[x_start:x_end, y_start:y_end] = False
|
||||
|
||||
def apply_padding(self, padding_cells: int):
|
||||
"""Dilate invalid regions by padding cells for safety margin."""
|
||||
obstacles = ~self.valid_mask
|
||||
kernel = np.ones((2 * padding_cells + 1, 2 * padding_cells + 1), dtype=bool)
|
||||
dilated = binary_dilation(obstacles, structure=kernel)
|
||||
self.valid_mask = ~dilated
|
||||
|
||||
def create_spawn_mask(self, spawn_padding_cells: int) -> np.ndarray:
|
||||
"""Create a mask for spawn positions with larger padding than goals."""
|
||||
extra_padding = spawn_padding_cells - PADDING.GOAL_PADDING
|
||||
if extra_padding > 0:
|
||||
obstacles = ~self.valid_mask
|
||||
kernel = np.ones((2 * extra_padding + 1, 2 * extra_padding + 1), dtype=bool)
|
||||
dilated = binary_dilation(obstacles, structure=kernel)
|
||||
return ~dilated
|
||||
return self.valid_mask.copy()
|
||||
|
||||
def exclude_borders(self, border_cells: int = 2):
|
||||
"""Mark terrain borders as invalid."""
|
||||
self.valid_mask[:border_cells, :] = False
|
||||
self.valid_mask[-border_cells:, :] = False
|
||||
self.valid_mask[:, :border_cells] = False
|
||||
self.valid_mask[:, -border_cells:] = False
|
||||
|
||||
def apply_height_transition_padding(self, height_threshold: int, padding_cells: int):
|
||||
"""Mark cells near height transitions as invalid."""
|
||||
grad_x = np.abs(np.diff(self.heights, axis=0, prepend=self.heights[:1, :]))
|
||||
grad_y = np.abs(np.diff(self.heights, axis=1, prepend=self.heights[:, :1]))
|
||||
grad_x_back = np.abs(np.diff(self.heights, axis=0, append=self.heights[-1:, :]))
|
||||
grad_y_back = np.abs(np.diff(self.heights, axis=1, append=self.heights[:, -1:]))
|
||||
|
||||
max_grad = np.maximum.reduce([grad_x, grad_y, grad_x_back, grad_y_back])
|
||||
transition_mask = (max_grad >= height_threshold).astype(bool)
|
||||
|
||||
if padding_cells > 0:
|
||||
kernel = np.ones((2 * padding_cells + 1, 2 * padding_cells + 1), dtype=bool)
|
||||
transition_mask = binary_dilation(transition_mask, structure=kernel).astype(bool)
|
||||
|
||||
self.valid_mask = self.valid_mask & ~transition_mask
|
||||
|
||||
|
||||
def get_cell_bounds(
|
||||
cell_x: int, cell_y: int, cell_pixels: int, max_x: int, max_y: int
|
||||
) -> Tuple[int, int, int, int]:
|
||||
"""Get pixel bounds for a maze cell with clamping.
|
||||
|
||||
Returns:
|
||||
Tuple of (x_start, x_end, y_start, y_end).
|
||||
"""
|
||||
return (
|
||||
max(0, cell_x * cell_pixels),
|
||||
min(max_x, (cell_x + 1) * cell_pixels),
|
||||
max(0, cell_y * cell_pixels),
|
||||
min(max_y, (cell_y + 1) * cell_pixels),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Maze Generation
|
||||
# =============================================================================
|
||||
|
||||
def generate_maze(
|
||||
rng: np.random.Generator,
|
||||
width: int,
|
||||
height: int,
|
||||
open_prob: float
|
||||
) -> np.ndarray:
|
||||
"""Generate maze using DFS with random openings.
|
||||
|
||||
Args:
|
||||
rng: Random number generator for reproducibility.
|
||||
width: Maze width in cells.
|
||||
height: Maze height in cells.
|
||||
open_prob: Probability of random wall removal.
|
||||
|
||||
Returns:
|
||||
2D array where 1=wall, 0=path.
|
||||
"""
|
||||
maze = np.ones((width, height), dtype=np.uint8)
|
||||
stack = [(0, 0)]
|
||||
maze[0, 0] = 0
|
||||
|
||||
while stack:
|
||||
x, y = stack[-1]
|
||||
neighbors = []
|
||||
for dx, dy in [(-2, 0), (2, 0), (0, -2), (0, 2)]:
|
||||
nx, ny = x + dx, y + dy
|
||||
if 0 <= nx < width and 0 <= ny < height and maze[nx, ny] == 1:
|
||||
neighbors.append((nx, ny))
|
||||
|
||||
if neighbors:
|
||||
idx = rng.integers(len(neighbors))
|
||||
nx, ny = neighbors[idx]
|
||||
maze[(x + nx) // 2, (y + ny) // 2] = 0
|
||||
maze[nx, ny] = 0
|
||||
stack.append((nx, ny))
|
||||
else:
|
||||
stack.pop()
|
||||
|
||||
# Random openings
|
||||
maze[rng.random((width, height)) < open_prob] = 0
|
||||
return maze
|
||||
|
||||
|
||||
def clear_center(maze: np.ndarray, terrain: TerrainData, cell_pixels: int):
|
||||
"""Clear the center area for spawning."""
|
||||
cx, cy = maze.shape[0] // 2, maze.shape[1] // 2
|
||||
|
||||
for dx in [-1, 0, 1]:
|
||||
for dy in [-1, 0, 1]:
|
||||
if abs(dx) + abs(dy) <= 1: # Plus shape
|
||||
nx, ny = cx + dx, cy + dy
|
||||
if 0 <= nx < maze.shape[0] and 0 <= ny < maze.shape[1]:
|
||||
maze[nx, ny] = 0
|
||||
|
||||
x_start = (cx - 1) * cell_pixels
|
||||
x_end = (cx + 2) * cell_pixels
|
||||
y_start = (cy - 1) * cell_pixels
|
||||
y_end = (cy + 2) * cell_pixels
|
||||
terrain.set_ground(x_start, x_end, y_start, y_end)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Obstacle Generators
|
||||
# =============================================================================
|
||||
|
||||
def make_pillar(
|
||||
_rng: np.random.Generator,
|
||||
size: int,
|
||||
wall_height: int,
|
||||
scale: float,
|
||||
is_pit: bool,
|
||||
thickness: int
|
||||
) -> np.ndarray:
|
||||
"""Generate a centered pillar obstacle."""
|
||||
grid = np.zeros((size, size), dtype=np.int16)
|
||||
h = int(wall_height * scale) * (-1 if is_pit else 1)
|
||||
grid[thickness:size-thickness, thickness:size-thickness] = h
|
||||
return grid
|
||||
|
||||
|
||||
def make_bar(
|
||||
rng: np.random.Generator,
|
||||
size: int,
|
||||
wall_height: int,
|
||||
scale: float,
|
||||
is_pit: bool,
|
||||
thickness: int
|
||||
) -> np.ndarray:
|
||||
"""Generate a rotated bar obstacle."""
|
||||
grid = np.zeros((size, size), dtype=np.int16)
|
||||
center = size // 2
|
||||
h = int(wall_height * scale)
|
||||
grid[center - thickness//2:center + thickness//2, :] = h
|
||||
|
||||
angle = rng.uniform(-180, 180)
|
||||
grid = rotate(grid, angle, reshape=False, order=1).astype(np.int16)
|
||||
|
||||
if is_pit:
|
||||
grid = -grid
|
||||
return grid
|
||||
|
||||
|
||||
def make_cross(
|
||||
rng: np.random.Generator,
|
||||
size: int,
|
||||
wall_height: int,
|
||||
scale: float,
|
||||
is_pit: bool,
|
||||
thickness: int
|
||||
) -> np.ndarray:
|
||||
"""Generate a cross-shaped obstacle."""
|
||||
grid = np.zeros((size, size), dtype=np.int16)
|
||||
center = size // 2
|
||||
h = int(wall_height * scale)
|
||||
grid[center - thickness//2:center + thickness//2, :] = h
|
||||
grid[:, center - thickness//2:center + thickness//2] = h
|
||||
|
||||
angle = rng.uniform(-180, 180)
|
||||
grid = rotate(grid, angle, reshape=False, order=1).astype(np.int16)
|
||||
|
||||
if is_pit:
|
||||
grid = -grid
|
||||
return grid
|
||||
|
||||
|
||||
def make_shifted_block(
|
||||
rng: np.random.Generator,
|
||||
size: int,
|
||||
wall_height: int,
|
||||
scale: float,
|
||||
is_pit: bool,
|
||||
thickness: int
|
||||
) -> np.ndarray:
|
||||
"""Generate a randomly shifted block."""
|
||||
grid = np.zeros((size, size), dtype=np.int16)
|
||||
h = int(wall_height * scale)
|
||||
grid[thickness:size-thickness, thickness:size-thickness] = h
|
||||
|
||||
room = max(0, size // 2 - thickness)
|
||||
shift_amt = (
|
||||
rng.integers(-room, room + 1),
|
||||
rng.integers(-room, room + 1)
|
||||
)
|
||||
grid = shift(grid, shift=shift_amt, cval=0).astype(np.int16)
|
||||
|
||||
if is_pit:
|
||||
grid = -grid
|
||||
return grid
|
||||
|
||||
|
||||
# Obstacle generator lookup table
|
||||
_OBSTACLE_GENERATORS = {
|
||||
ObstacleType.PILLAR: make_pillar,
|
||||
ObstacleType.BAR: make_bar,
|
||||
ObstacleType.CROSS: make_cross,
|
||||
ObstacleType.SHIFTED_BLOCK: make_shifted_block,
|
||||
}
|
||||
|
||||
|
||||
def make_random_obstacle(
|
||||
rng: np.random.Generator,
|
||||
size: int,
|
||||
wall_height: int,
|
||||
is_pit: bool | None = None,
|
||||
pillar_weight: float | None = None
|
||||
) -> np.ndarray:
|
||||
"""Generate a random obstacle type.
|
||||
|
||||
Args:
|
||||
rng: Random number generator.
|
||||
size: Size of the obstacle grid in pixels.
|
||||
wall_height: Height of walls in terrain units.
|
||||
is_pit: Force pit (True) or wall (False). None = random.
|
||||
pillar_weight: Weight for pillars (0-1). None = uniform distribution.
|
||||
"""
|
||||
scale = rng.uniform(OBSTACLES.SCALE_MIN, OBSTACLES.SCALE_MAX)
|
||||
if is_pit is None:
|
||||
is_pit = rng.random() < OBSTACLES.DEFAULT_PIT_PROB
|
||||
thickness = rng.integers(OBSTACLES.THICKNESS_MIN, OBSTACLES.THICKNESS_MAX)
|
||||
|
||||
# Select obstacle type (with optional pillar weighting)
|
||||
if pillar_weight is not None and pillar_weight > 0:
|
||||
# Weighted selection: pillar_weight for pillars, rest split evenly
|
||||
other_weight = (1.0 - pillar_weight) / (ObstacleType.NUM_TYPES - 1)
|
||||
weights = [other_weight] * ObstacleType.NUM_TYPES
|
||||
weights[ObstacleType.PILLAR] = pillar_weight
|
||||
obstacle_type = rng.choice(ObstacleType.NUM_TYPES, p=weights)
|
||||
else:
|
||||
# Uniform selection
|
||||
obstacle_type = rng.integers(ObstacleType.NUM_TYPES)
|
||||
|
||||
generator = _OBSTACLE_GENERATORS[obstacle_type]
|
||||
return generator(rng, size, wall_height, scale, is_pit, thickness)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Stair/Platform Generator
|
||||
# =============================================================================
|
||||
|
||||
class StairGenerator:
|
||||
"""Generates stair structures with platforms."""
|
||||
|
||||
LAYOUTS = [
|
||||
{"platforms": [(1, 0), (1, 1), (1, 2), (0, 1), (2, 1)],
|
||||
"stairs": [(0, 0, "n"), (2, 2, "s"), (0, 2, "s"), (2, 0, "w")]},
|
||||
{"platforms": [(0, 1), (1, 1), (2, 1)],
|
||||
"stairs": [(0, 0, "n"), (2, 2, "s"), (0, 2, "s"), (2, 0, "n")]},
|
||||
{"platforms": [(0, 1), (1, 1), (2, 1), (1, 0), (1, 2)],
|
||||
"stairs": [(0, 0, "n"), (2, 2, "s"), (0, 2, "s"), (2, 0, "n")]},
|
||||
]
|
||||
|
||||
def __init__(self, wall_height: float, vertical_scale: float):
|
||||
self.wall_height = wall_height
|
||||
self.platform_height = int(wall_height - 0.5 / vertical_scale)
|
||||
self.vertical_scale = vertical_scale
|
||||
self._make_stair_templates()
|
||||
|
||||
def _make_stair_templates(self):
|
||||
"""Create stair templates for each direction."""
|
||||
cell_px = STAIRS.SINGLE_CELL_PIXELS
|
||||
step_res = cell_px // STAIRS.NUM_STEPS
|
||||
|
||||
east = np.zeros((cell_px, cell_px), dtype=np.float32)
|
||||
for i in range(STAIRS.NUM_STEPS):
|
||||
h = STAIRS.STEP_HEIGHT_METERS * (i + 1) / self.vertical_scale
|
||||
east[i * step_res:(i + 1) * step_res, :] = h
|
||||
|
||||
self.templates = {
|
||||
"e": east,
|
||||
"n": rotate(east, 90),
|
||||
"w": rotate(east, 180),
|
||||
"s": rotate(east, 270),
|
||||
}
|
||||
|
||||
def generate(self, rng: np.random.Generator) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Generate a 3x3 stair/platform structure.
|
||||
|
||||
Args:
|
||||
rng: Random number generator.
|
||||
|
||||
Returns:
|
||||
Tuple of (heights, valid_mask, platform_mask).
|
||||
"""
|
||||
layout = self.LAYOUTS[rng.integers(len(self.LAYOUTS))]
|
||||
size = STAIRS.STAIR_GRID_SIZE * STAIRS.SINGLE_CELL_PIXELS
|
||||
cell_px = STAIRS.SINGLE_CELL_PIXELS
|
||||
|
||||
heights = np.zeros((size, size), dtype=np.float32)
|
||||
valid_mask = np.zeros((size, size), dtype=bool)
|
||||
platform_mask = np.zeros((size, size), dtype=bool)
|
||||
|
||||
for gx, gy in layout["platforms"]:
|
||||
xs, xe = gx * cell_px, (gx + 1) * cell_px
|
||||
ys, ye = gy * cell_px, (gy + 1) * cell_px
|
||||
heights[xs:xe, ys:ye] = self.platform_height
|
||||
valid_mask[xs:xe, ys:ye] = True
|
||||
platform_mask[xs:xe, ys:ye] = True
|
||||
|
||||
for gx, gy, direction in layout["stairs"]:
|
||||
xs, xe = gx * cell_px, (gx + 1) * cell_px
|
||||
ys, ye = gy * cell_px, (gy + 1) * cell_px
|
||||
heights[xs:xe, ys:ye] = self.templates[direction]
|
||||
|
||||
return heights, valid_mask, platform_mask
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Main Terrain Generation
|
||||
# =============================================================================
|
||||
|
||||
def _get_rng(cfg: "hf_terrains_maze_cfg.HfMazeTerrainCfg") -> np.random.Generator:
|
||||
"""Get RNG from config or create a new one."""
|
||||
if cfg.rng is not None:
|
||||
return cfg.rng
|
||||
# Fallback: create unseeded RNG (non-reproducible)
|
||||
return np.random.default_rng()
|
||||
|
||||
|
||||
@height_field_to_mesh
|
||||
def maze_terrain(difficulty: float, cfg: "hf_terrains_maze_cfg.HfMazeTerrainCfg") -> np.ndarray:
|
||||
"""Generate maze terrain with obstacles and valid position mask.
|
||||
|
||||
Args:
|
||||
difficulty: Terrain difficulty (0-1).
|
||||
cfg: Terrain configuration.
|
||||
|
||||
Returns:
|
||||
Height field for mesh generation.
|
||||
"""
|
||||
rng = _get_rng(cfg)
|
||||
|
||||
# Setup dimensions
|
||||
cell_pixels = int(cfg.cell_size / cfg.horizontal_scale)
|
||||
wall_height = int(cfg.wall_height / cfg.vertical_scale)
|
||||
terrain_w = int(cfg.size[0] / cfg.horizontal_scale)
|
||||
terrain_h = int(cfg.size[1] / cfg.horizontal_scale)
|
||||
|
||||
terrain = TerrainData.create(terrain_w, terrain_h)
|
||||
stair_gen = StairGenerator(wall_height, cfg.vertical_scale)
|
||||
|
||||
# Generate base pattern
|
||||
if cfg.non_maze_terrain:
|
||||
maze = np.zeros(cfg.grid_size, dtype=np.uint8)
|
||||
obstacle_prob = difficulty * OBSTACLES.NON_MAZE_DENSITY
|
||||
maze[rng.random(cfg.grid_size) < obstacle_prob] = 1
|
||||
else:
|
||||
maze = generate_maze(rng, cfg.grid_size[0], cfg.grid_size[1], 1 - difficulty)
|
||||
|
||||
clear_center(maze, terrain, cell_pixels)
|
||||
|
||||
# Generate terrain features based on type
|
||||
if cfg.dynamic_obstacles:
|
||||
_add_pits(rng, terrain, cfg, difficulty, wall_height, cell_pixels)
|
||||
elif cfg.stairs:
|
||||
_add_stairs(rng, terrain, cfg, difficulty, wall_height, cell_pixels, stair_gen)
|
||||
else:
|
||||
_add_walls(rng, maze, terrain, cfg, wall_height, cell_pixels)
|
||||
|
||||
clear_center(maze, terrain, cell_pixels)
|
||||
|
||||
# Apply height transition padding for stair terrain
|
||||
if cfg.stairs:
|
||||
terrain.apply_height_transition_padding(
|
||||
height_threshold=PADDING.HEIGHT_TRANSITION_THRESHOLD,
|
||||
padding_cells=PADDING.HEIGHT_TRANSITION_PADDING
|
||||
)
|
||||
|
||||
# Apply safety padding and border exclusion
|
||||
terrain.apply_padding(PADDING.GOAL_PADDING)
|
||||
terrain.exclude_borders(PADDING.BORDER_CELLS)
|
||||
|
||||
# Create spawn mask with larger padding
|
||||
spawn_mask = terrain.create_spawn_mask(PADDING.SPAWN_PADDING)
|
||||
spawn_mask[:PADDING.BORDER_CELLS, :] = False
|
||||
spawn_mask[-PADDING.BORDER_CELLS:, :] = False
|
||||
spawn_mask[:, :PADDING.BORDER_CELLS] = False
|
||||
spawn_mask[:, -PADDING.BORDER_CELLS:] = False
|
||||
|
||||
# Store data on cfg for patches to pick up
|
||||
if cfg.add_goal:
|
||||
cfg.height_field_visual = torch.from_numpy(terrain.heights.copy()).unsqueeze(0)
|
||||
cfg.height_field_valid_mask = torch.from_numpy(terrain.valid_mask.copy()).unsqueeze(0)
|
||||
cfg.height_field_platform_mask = torch.from_numpy(terrain.platform_mask.copy()).unsqueeze(0)
|
||||
cfg.height_field_spawn_mask = torch.from_numpy(spawn_mask.copy()).unsqueeze(0)
|
||||
|
||||
return terrain.heights
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Terrain Type Generators
|
||||
# =============================================================================
|
||||
|
||||
def _add_walls(
|
||||
rng: np.random.Generator,
|
||||
maze: np.ndarray,
|
||||
terrain: TerrainData,
|
||||
cfg,
|
||||
wall_height: int,
|
||||
cell_pixels: int
|
||||
):
|
||||
"""Add wall obstacles to terrain based on maze pattern."""
|
||||
# Use pillar weighting for non-maze terrain (more thin pillars)
|
||||
pillar_weight = OBSTACLES.NON_MAZE_PILLAR_WEIGHT if cfg.non_maze_terrain else None
|
||||
|
||||
for x in range(cfg.grid_size[0]):
|
||||
for y in range(cfg.grid_size[1]):
|
||||
if maze[x, y] != 1:
|
||||
continue
|
||||
|
||||
xs, xe, ys, ye = get_cell_bounds(
|
||||
x, y, cell_pixels, terrain.heights.shape[0], terrain.heights.shape[1]
|
||||
)
|
||||
|
||||
if cfg.randomize_wall and rng.random() < cfg.random_wall_ratio:
|
||||
obs = make_random_obstacle(rng, cell_pixels, wall_height, pillar_weight=pillar_weight)
|
||||
terrain.heights[xs:xe, ys:ye] = obs[:xe - xs, :ye - ys]
|
||||
terrain.valid_mask[xs:xe, ys:ye] = False
|
||||
else:
|
||||
h = int(wall_height * rng.uniform(OBSTACLES.SCALE_MIN, OBSTACLES.SCALE_MAX))
|
||||
terrain.set_obstacle(xs, xe, ys, ye, h)
|
||||
|
||||
|
||||
def _add_stairs(
|
||||
rng: np.random.Generator,
|
||||
terrain: TerrainData,
|
||||
cfg,
|
||||
difficulty: float,
|
||||
wall_height: int,
|
||||
cell_pixels: int,
|
||||
stair_gen: StairGenerator
|
||||
):
|
||||
"""Add stair/platform structures to terrain."""
|
||||
grid_w, grid_h = cfg.grid_size
|
||||
grid_middle = grid_w // 2
|
||||
excluded = set(range(grid_middle - 1, grid_middle + 1))
|
||||
|
||||
# Compute stair placement locations (avoid center and edges)
|
||||
# Stairs are 3x3, so max position is grid_size - 4 to fit with margin
|
||||
stair_margin = 1
|
||||
max_x = grid_w - STAIRS.STAIR_GRID_SIZE - stair_margin
|
||||
max_y = grid_h - STAIRS.STAIR_GRID_SIZE - stair_margin
|
||||
num_locations = 6
|
||||
x_locs = set(np.round(np.linspace(stair_margin, max_x, num_locations)).astype(int)) - excluded
|
||||
y_locs = set(np.round(np.linspace(stair_margin, max_y, num_locations)).astype(int)) - excluded
|
||||
|
||||
processed = set()
|
||||
stair_size = STAIRS.STAIR_GRID_SIZE * STAIRS.SINGLE_CELL_PIXELS
|
||||
stair_prob = difficulty * OBSTACLES.STAIRS_PLACEMENT_PROB
|
||||
obstacle_prob = difficulty * OBSTACLES.STAIRS_OBSTACLE_DENSITY
|
||||
|
||||
for x in range(grid_w):
|
||||
for y in range(grid_h):
|
||||
if (x, y) in processed:
|
||||
continue
|
||||
|
||||
# Try placing stair structure at valid locations
|
||||
if x in x_locs and y in y_locs and rng.random() < stair_prob:
|
||||
heights, valid, platform = stair_gen.generate(rng)
|
||||
|
||||
xs = x * cell_pixels
|
||||
xe = min(terrain.heights.shape[0], xs + stair_size)
|
||||
ys = y * cell_pixels
|
||||
ye = min(terrain.heights.shape[1], ys + stair_size)
|
||||
|
||||
sx, sy = xe - xs, ye - ys
|
||||
terrain.heights[xs:xe, ys:ye] = heights[:sx, :sy]
|
||||
terrain.valid_mask[xs:xe, ys:ye] = valid[:sx, :sy]
|
||||
terrain.platform_mask[xs:xe, ys:ye] = platform[:sx, :sy]
|
||||
|
||||
# Mark 3x3 area as processed
|
||||
for dx in range(3):
|
||||
for dy in range(3):
|
||||
processed.add((x + dx, y + dy))
|
||||
|
||||
elif rng.random() < obstacle_prob:
|
||||
xs, xe, ys, ye = get_cell_bounds(
|
||||
x, y, cell_pixels, terrain.heights.shape[0], terrain.heights.shape[1]
|
||||
)
|
||||
# Check if area is clear before placing
|
||||
if terrain.valid_mask[xs + 1:xe - 1, ys + 1:ye - 1].all():
|
||||
obs = make_random_obstacle(rng, cell_pixels, wall_height)
|
||||
terrain.heights[xs:xe, ys:ye] = obs[:xe - xs, :ye - ys]
|
||||
terrain.valid_mask[xs:xe, ys:ye] = False
|
||||
|
||||
|
||||
def _add_pits(
|
||||
rng: np.random.Generator,
|
||||
terrain: TerrainData,
|
||||
cfg,
|
||||
difficulty: float,
|
||||
wall_height: int,
|
||||
cell_pixels: int
|
||||
):
|
||||
"""Add pit/trough obstacles to terrain.
|
||||
|
||||
Layout:
|
||||
- Two horizontal pit trenches with random bridges for crossing
|
||||
- Random obstacles (mostly pits) scattered in the middle area
|
||||
"""
|
||||
grid_w, grid_h = cfg.grid_size
|
||||
|
||||
# Pit trench rows (near top and bottom)
|
||||
trench_offset = OBSTACLES.PITS_TRENCH_ROW_OFFSET
|
||||
pit_rows = {trench_offset, grid_h - trench_offset - 1}
|
||||
|
||||
# Generate bridge positions for crossing pit trenches
|
||||
bridges = _generate_bridges(rng, grid_w)
|
||||
|
||||
# Add pit trenches (negative height = troughs)
|
||||
for pit_y in pit_rows:
|
||||
for x in range(grid_w):
|
||||
if x in bridges:
|
||||
continue
|
||||
xs, xe, ys, ye = get_cell_bounds(
|
||||
x, pit_y, cell_pixels, terrain.heights.shape[0], terrain.heights.shape[1]
|
||||
)
|
||||
terrain.set_obstacle(xs, xe, ys, ye, -wall_height)
|
||||
|
||||
# Add random obstacles in middle area (between pit trenches)
|
||||
_add_middle_obstacles(rng, terrain, cfg, difficulty, wall_height, cell_pixels, pit_rows)
|
||||
|
||||
|
||||
def _generate_bridges(rng: np.random.Generator, grid_width: int) -> set:
|
||||
"""Generate bridge positions across pit rows.
|
||||
|
||||
Returns set of x-coordinates where bridges (gaps in pits) are placed.
|
||||
Bridges are 2 cells wide for easier robot crossing.
|
||||
"""
|
||||
num_bridges = rng.integers(OBSTACLES.BRIDGE_COUNT_MIN, OBSTACLES.BRIDGE_COUNT_MAX)
|
||||
margin = OBSTACLES.PITS_EDGE_MARGIN
|
||||
available = list(range(margin, grid_width - margin))
|
||||
rng.shuffle(available)
|
||||
|
||||
bridges = set()
|
||||
for i in range(min(num_bridges, len(available))):
|
||||
pos = available[i]
|
||||
bridges.add(pos)
|
||||
# Make bridges 2 cells wide
|
||||
if pos + 1 < grid_width - margin:
|
||||
bridges.add(pos + 1)
|
||||
|
||||
return bridges
|
||||
|
||||
|
||||
def _add_middle_obstacles(
|
||||
rng: np.random.Generator,
|
||||
terrain: TerrainData,
|
||||
cfg,
|
||||
difficulty: float,
|
||||
wall_height: int,
|
||||
cell_pixels: int,
|
||||
pit_rows: set
|
||||
):
|
||||
"""Add random obstacles in the middle area between pit rows."""
|
||||
grid_w, grid_h = cfg.grid_size
|
||||
obstacle_prob = difficulty * OBSTACLES.PITS_DENSITY
|
||||
|
||||
# Compute valid placement bounds (avoid edges and pit rows)
|
||||
margin = OBSTACLES.PITS_EDGE_MARGIN
|
||||
trench_offset = OBSTACLES.PITS_TRENCH_ROW_OFFSET
|
||||
|
||||
x_range = range(margin, grid_w - margin)
|
||||
# Middle area: between the two pit trenches, with 1 cell buffer
|
||||
y_range = range(trench_offset + 1, grid_h - trench_offset - 1)
|
||||
|
||||
# Iterate only over valid cells (more efficient)
|
||||
for x in x_range:
|
||||
for y in y_range:
|
||||
if y in pit_rows:
|
||||
continue
|
||||
|
||||
if rng.random() < obstacle_prob:
|
||||
xs, xe, ys, ye = get_cell_bounds(
|
||||
x, y, cell_pixels, terrain.heights.shape[0], terrain.heights.shape[1]
|
||||
)
|
||||
obs = _generate_pit_obstacle(rng, cell_pixels, wall_height)
|
||||
terrain.heights[xs:xe, ys:ye] = obs[:xe - xs, :ye - ys]
|
||||
terrain.valid_mask[xs:xe, ys:ye] = False
|
||||
|
||||
|
||||
def _generate_pit_obstacle(
|
||||
rng: np.random.Generator,
|
||||
cell_pixels: int,
|
||||
wall_height: int
|
||||
) -> np.ndarray:
|
||||
"""Generate an obstacle for pit terrain with high pit probability.
|
||||
|
||||
Distribution:
|
||||
- 60% bars (75% negative/pits) -> 45% pit bars
|
||||
- 40% random shapes (50% negative/pits) -> 20% pit shapes
|
||||
- Total: ~65% negative obstacles
|
||||
"""
|
||||
if rng.random() < OBSTACLES.PITS_BAR_RATIO:
|
||||
# Bar obstacle with high pit probability
|
||||
is_pit = rng.random() < OBSTACLES.PITS_BAR_PIT_PROB
|
||||
scale = rng.uniform(OBSTACLES.SCALE_MIN, OBSTACLES.SCALE_MAX)
|
||||
thickness = rng.integers(OBSTACLES.THICKNESS_MIN, OBSTACLES.THICKNESS_MAX)
|
||||
return make_bar(rng, cell_pixels, wall_height, scale, is_pit, thickness)
|
||||
else:
|
||||
# Random obstacle type (pillar, cross, block) with moderate pit probability
|
||||
is_pit = rng.random() < OBSTACLES.PITS_RANDOM_PIT_PROB
|
||||
return make_random_obstacle(rng, cell_pixels, wall_height, is_pit=is_pit)
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Configuration for maze height field terrains."""
|
||||
|
||||
from dataclasses import MISSING
|
||||
from typing import Any, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from isaaclab.utils import configclass
|
||||
from isaaclab.terrains.height_field.hf_terrains_cfg import HfTerrainBaseCfg
|
||||
|
||||
from . import hf_terrains_maze
|
||||
|
||||
|
||||
@configclass
|
||||
class HfMazeTerrainCfg(HfTerrainBaseCfg):
|
||||
"""Configuration for a maze height field terrain.
|
||||
|
||||
This terrain generates a procedural maze with configurable wall structures,
|
||||
obstacles, and optional stairs. The maze can be used for navigation tasks
|
||||
with various difficulty levels.
|
||||
|
||||
Height Field Data (set during terrain generation):
|
||||
- height_field_visual: Heights for Z-lookup (num_terrains, W, H)
|
||||
- height_field_valid_mask: Valid goal positions with safety padding
|
||||
- height_field_platform_mask: Platform positions for curriculum
|
||||
- height_field_spawn_mask: Valid spawn positions with larger padding
|
||||
"""
|
||||
|
||||
function = hf_terrains_maze.maze_terrain
|
||||
|
||||
# =========================================================================
|
||||
# Height Field Storage (populated during terrain generation)
|
||||
# =========================================================================
|
||||
|
||||
height_field_visual: torch.Tensor = None
|
||||
"""Height field for Z-lookup (actual terrain heights)."""
|
||||
|
||||
height_field_valid_mask: torch.Tensor = None
|
||||
"""Boolean mask of valid goal positions (padded with GOAL_PADDING)."""
|
||||
|
||||
height_field_platform_mask: torch.Tensor = None
|
||||
"""Boolean mask of platform positions for curriculum learning."""
|
||||
|
||||
height_field_spawn_mask: torch.Tensor = None
|
||||
"""Boolean mask of valid spawn positions (larger padding for robot body)."""
|
||||
|
||||
# =========================================================================
|
||||
# Maze Generation Parameters
|
||||
# =========================================================================
|
||||
|
||||
maze: bool = True
|
||||
"""Flag indicating this is a maze terrain."""
|
||||
|
||||
open_probability: float = None
|
||||
"""Probability of a cell being open in the maze."""
|
||||
|
||||
grid_size: tuple[int, int] = (15, 15)
|
||||
"""Size of the maze grid (number of cells in width and height)."""
|
||||
|
||||
cell_size: float = 2.0
|
||||
"""Size of each cell in the maze grid (in meters)."""
|
||||
|
||||
wall_height: float = 1.5
|
||||
"""Height of the walls (in meters). Defaults to 1.5."""
|
||||
|
||||
# =========================================================================
|
||||
# Terrain Features
|
||||
# =========================================================================
|
||||
|
||||
add_goal: Any = MISSING
|
||||
"""Enable goal sampling data generation."""
|
||||
|
||||
add_noise_to_flat: Any = MISSING
|
||||
"""Add noise to flat areas of the maze."""
|
||||
|
||||
randomize_wall: Any = MISSING
|
||||
"""Use randomized obstacle shapes instead of full walls."""
|
||||
|
||||
random_wall_ratio: float = 0.5
|
||||
"""Mix ratio between randomized and standard walls. Defaults to 0.5."""
|
||||
|
||||
non_maze_terrain: bool = False
|
||||
"""Use non-maze terrain with random obstacles. Defaults to False."""
|
||||
|
||||
stairs: bool = False
|
||||
"""Add stairs to empty map. Defaults to False."""
|
||||
|
||||
add_stairs_to_maze: bool = False
|
||||
"""Add stairs to the maze. Defaults to False."""
|
||||
|
||||
dynamic_obstacles: bool = False
|
||||
"""Enable pit/trough obstacles. Defaults to False."""
|
||||
|
||||
# =========================================================================
|
||||
# Random Number Generator
|
||||
# =========================================================================
|
||||
|
||||
rng: Optional[np.random.Generator] = None
|
||||
"""Random number generator for reproducible terrain generation.
|
||||
|
||||
Set by the terrain generator (patches.py) before calling the terrain function.
|
||||
If None, will create a new unseeded generator (non-reproducible).
|
||||
"""
|
||||
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Configuration for maze terrains."""
|
||||
|
||||
from isaaclab.terrains.terrain_generator_cfg import TerrainGeneratorCfg
|
||||
|
||||
from .hf_terrains_maze_cfg import HfMazeTerrainCfg
|
||||
|
||||
MAZE_TERRAIN_CFG = TerrainGeneratorCfg(
|
||||
size=(30.0, 30.0),
|
||||
border_width=30.0, # Border around the entire terrain grid (not per-tile)
|
||||
num_rows=6,
|
||||
num_cols=30,
|
||||
horizontal_scale=0.1,
|
||||
vertical_scale=0.005,
|
||||
slope_threshold=0.75,
|
||||
use_cache=False,
|
||||
curriculum=False,
|
||||
difficulty_range=(0.5, 1.0),
|
||||
sub_terrains={
|
||||
"maze": HfMazeTerrainCfg(
|
||||
proportion=0.3,
|
||||
open_probability=0.9,
|
||||
grid_size=(15, 15),
|
||||
cell_size=2.0,
|
||||
add_noise_to_flat=False,
|
||||
add_goal=True,
|
||||
randomize_wall=True,
|
||||
random_wall_ratio=0.5,
|
||||
add_stairs_to_maze=True,
|
||||
),
|
||||
"non_maze": HfMazeTerrainCfg(
|
||||
proportion=0.2,
|
||||
open_probability=0.9,
|
||||
grid_size=(15, 15),
|
||||
cell_size=2.0,
|
||||
add_noise_to_flat=False,
|
||||
add_goal=True,
|
||||
randomize_wall=True,
|
||||
random_wall_ratio=1.0,
|
||||
non_maze_terrain=True,
|
||||
),
|
||||
"stairs": HfMazeTerrainCfg(
|
||||
proportion=0.3,
|
||||
open_probability=0.9,
|
||||
grid_size=(15, 15),
|
||||
cell_size=2.0,
|
||||
add_noise_to_flat=False,
|
||||
add_goal=True,
|
||||
randomize_wall=False,
|
||||
random_wall_ratio=1.0,
|
||||
non_maze_terrain=False,
|
||||
stairs=True,
|
||||
),
|
||||
"pits": HfMazeTerrainCfg(
|
||||
proportion=0.2,
|
||||
open_probability=0.9,
|
||||
grid_size=(15, 15),
|
||||
cell_size=2.0,
|
||||
add_noise_to_flat=False,
|
||||
add_goal=True,
|
||||
randomize_wall=True,
|
||||
random_wall_ratio=1.0,
|
||||
non_maze_terrain=True,
|
||||
dynamic_obstacles=True, # Enables pit/trough generation
|
||||
),
|
||||
},
|
||||
)
|
||||
"""Maze terrain configuration for navigation tasks."""
|
||||
@@ -0,0 +1,550 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Monkey-patches for Isaac Lab terrain system to support navigation tasks.
|
||||
|
||||
This module patches TerrainGenerator and TerrainImporter to:
|
||||
1. Store height field data for goal/spawn sampling
|
||||
2. Use optimized mesh generation to reduce GPU memory
|
||||
|
||||
The patches add these attributes to TerrainImporter (accessed via env.scene.terrain):
|
||||
- _height_field_visual: Heights for Z-lookup (num_terrains, W, H)
|
||||
- _height_field_valid_mask: Valid goal positions with safety padding
|
||||
- _height_field_platform_mask: Platform positions for curriculum
|
||||
- _height_field_spawn_mask: Valid spawn positions with larger padding
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import trimesh
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.terrains.terrain_generator import TerrainGenerator
|
||||
from isaaclab.terrains.terrain_importer import TerrainImporter
|
||||
|
||||
# Flag to track if patches have been applied
|
||||
_patches_applied = False
|
||||
|
||||
# Shared storage for passing height field data from TerrainGenerator to TerrainImporter
|
||||
# Works because TerrainImporter creates TerrainGenerator synchronously in __init__
|
||||
_height_field_storage = {
|
||||
"height_field_visual": None,
|
||||
"height_field_valid_mask": None,
|
||||
"height_field_platform_mask": None,
|
||||
"height_field_spawn_mask": None,
|
||||
}
|
||||
|
||||
|
||||
def apply_terrain_patches():
|
||||
"""Apply monkey-patches to Isaac Lab terrain classes.
|
||||
|
||||
Patches applied:
|
||||
1. height_field_to_mesh - optimized mesh generation (reduces GPU memory)
|
||||
2. TerrainGenerator - collects height field data for goal sampling
|
||||
3. TerrainImporter - stores height field data as attributes
|
||||
4. PinholeCameraPatternCfg - adds from_ros_camera_info convenience method
|
||||
|
||||
Safe to call multiple times - patches are only applied once.
|
||||
"""
|
||||
global _patches_applied
|
||||
if _patches_applied:
|
||||
return
|
||||
|
||||
_patch_height_field_to_mesh()
|
||||
_patch_terrain_generator()
|
||||
_patch_terrain_importer()
|
||||
_patch_pinhole_camera_pattern_cfg()
|
||||
|
||||
_patches_applied = True
|
||||
|
||||
|
||||
def _convert_height_field_to_mesh_with_optimization_dynamic(
|
||||
height_field: np.ndarray, horizontal_scale: float, vertical_scale: float, block_size: int = 20
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Convert a height-field array to a triangle mesh, optimizing large flat ground blocks (20x20, 10x10, 5x5).
|
||||
|
||||
This function optimizes mesh generation by detecting flat regions and simplifying them
|
||||
to just 2 triangles instead of generating full detail. It uses a hierarchical approach:
|
||||
- 20x20 blocks are checked first
|
||||
- Non-flat blocks are subdivided into 10x10, then 5x5
|
||||
- Only 5x5 or smaller blocks with height variation get full detail
|
||||
|
||||
This can dramatically reduce vertex count for terrains with large flat areas,
|
||||
significantly reducing GPU memory usage during simulation.
|
||||
|
||||
Args:
|
||||
height_field: The input height-field array.
|
||||
horizontal_scale: The discretization of the terrain along the x and y axis.
|
||||
vertical_scale: The discretization of the terrain along the z axis.
|
||||
block_size: Initial block size for optimization (default 20).
|
||||
|
||||
Returns:
|
||||
The vertices and triangles of the mesh:
|
||||
- **vertices** (np.ndarray(float)): Array of shape (num_vertices, 3).
|
||||
- **triangles** (np.ndarray(int)): Array of shape (num_triangles, 3).
|
||||
"""
|
||||
num_rows, num_cols = height_field.shape
|
||||
vertices = []
|
||||
triangles = []
|
||||
vertex_count = 0
|
||||
|
||||
def process_block(i, j, block_size):
|
||||
"""Process a block and decide whether to simplify or subdivide into smaller blocks."""
|
||||
nonlocal vertex_count
|
||||
|
||||
# Skip if starting position is at or beyond the terrain bounds
|
||||
if i >= num_rows or j >= num_cols:
|
||||
return
|
||||
|
||||
# Determine block dimensions within bounds of the height field
|
||||
block_end_row = min(i + block_size + 1, num_rows) # +1 to handle the right and bottom edges
|
||||
block_end_col = min(j + block_size + 1, num_cols) # +1 to handle the right and bottom edges
|
||||
|
||||
# Extract the block of height data
|
||||
block = height_field[i:block_end_row, j:block_end_col]
|
||||
|
||||
# Skip empty or single-element blocks (can't form triangles)
|
||||
if block.size == 0 or block.shape[0] < 2 or block.shape[1] < 2:
|
||||
return
|
||||
|
||||
# Check if the entire block is flat (all heights are the same)
|
||||
if np.all(block == block[0, 0]):
|
||||
# Simplify the block by using two large triangles for the whole flat region
|
||||
v0 = [i * horizontal_scale, j * horizontal_scale, block[0, 0] * vertical_scale]
|
||||
v1 = [block_end_row * horizontal_scale, j * horizontal_scale, block[0, 0] * vertical_scale]
|
||||
v2 = [i * horizontal_scale, block_end_col * horizontal_scale, block[0, 0] * vertical_scale]
|
||||
v3 = [block_end_row * horizontal_scale, block_end_col * horizontal_scale, block[0, 0] * vertical_scale]
|
||||
|
||||
# Add the vertices for the large triangles
|
||||
vertices.extend([v0, v1, v2, v3])
|
||||
|
||||
# Add the two triangles for the block
|
||||
triangles.append([vertex_count, vertex_count + 1, vertex_count + 2])
|
||||
triangles.append([vertex_count + 1, vertex_count + 3, vertex_count + 2])
|
||||
|
||||
vertex_count += 4
|
||||
else:
|
||||
# If block is not flat and the block size is larger than 5x5, subdivide into smaller blocks
|
||||
if block_size > 5:
|
||||
half_size = block_size // 2
|
||||
# Process each of the four quadrants of the block
|
||||
process_block(i, j, half_size)
|
||||
process_block(i + half_size, j, half_size)
|
||||
process_block(i, j + half_size, half_size)
|
||||
process_block(i + half_size, j + half_size, half_size)
|
||||
else:
|
||||
# If block size is 5x5 or smaller, generate detailed triangles for each grid point
|
||||
for x in range(block_end_row - i):
|
||||
for y in range(block_end_col - j):
|
||||
# Get the height at this point
|
||||
z = block[x, y] * vertical_scale
|
||||
v = [(i + x) * horizontal_scale, (j + y) * horizontal_scale, z]
|
||||
vertices.append(v)
|
||||
|
||||
# Now create triangles for all internal points including the last row and last column
|
||||
for x in range(block_end_row - i - 1): # Handle rows, including the last row
|
||||
for y in range(block_end_col - j - 1): # Handle columns, including the last column
|
||||
ind0 = vertex_count + x * (block_end_col - j) + y
|
||||
ind1 = ind0 + 1
|
||||
ind2 = ind0 + (block_end_col - j)
|
||||
ind3 = ind2 + 1
|
||||
|
||||
# Create two triangles for this grid cell
|
||||
triangles.append([ind0, ind3, ind1])
|
||||
triangles.append([ind0, ind2, ind3])
|
||||
|
||||
vertex_count += (block_end_row - i) * (block_end_col - j)
|
||||
|
||||
# Start by processing blocks with the initial block size (20x20)
|
||||
for i in range(0, num_rows, block_size):
|
||||
for j in range(0, num_cols, block_size):
|
||||
process_block(i, j, block_size)
|
||||
|
||||
vertices = np.array(vertices)
|
||||
triangles = np.array(triangles)
|
||||
|
||||
# Return vertices and triangles arrays
|
||||
return vertices, triangles
|
||||
|
||||
|
||||
def _patch_height_field_to_mesh():
|
||||
"""Patch height_field_to_mesh decorator to use optimized mesh generation.
|
||||
|
||||
This patch replaces the mesh generation in the height_field_to_mesh decorator
|
||||
with an optimized version that reduces vertex count for flat terrain regions,
|
||||
significantly reducing GPU memory usage during simulation with many environments.
|
||||
"""
|
||||
import copy
|
||||
import functools
|
||||
from isaaclab.terrains.height_field import utils as hf_utils
|
||||
|
||||
def _patched_height_field_to_mesh(func):
|
||||
"""Patched decorator that uses optimized mesh generation."""
|
||||
@functools.wraps(func)
|
||||
def wrapper(difficulty: float, cfg):
|
||||
# check valid border width
|
||||
if cfg.border_width > 0 and cfg.border_width < cfg.horizontal_scale:
|
||||
raise ValueError(
|
||||
f"The border width ({cfg.border_width}) must be greater than or equal to the"
|
||||
f" horizontal scale ({cfg.horizontal_scale})."
|
||||
)
|
||||
# allocate buffer for height field (with border)
|
||||
width_pixels = int(cfg.size[0] / cfg.horizontal_scale) + 1
|
||||
length_pixels = int(cfg.size[1] / cfg.horizontal_scale) + 1
|
||||
border_pixels = int(cfg.border_width / cfg.horizontal_scale) + 1
|
||||
heights = np.zeros((width_pixels, length_pixels), dtype=np.int16)
|
||||
# override size of the terrain to account for the border
|
||||
sub_terrain_size = [width_pixels - 2 * border_pixels, length_pixels - 2 * border_pixels]
|
||||
sub_terrain_size = [dim * cfg.horizontal_scale for dim in sub_terrain_size]
|
||||
# update the config
|
||||
terrain_size = copy.deepcopy(cfg.size)
|
||||
cfg.size = tuple(sub_terrain_size)
|
||||
# generate the height field
|
||||
z_gen = func(difficulty, cfg)
|
||||
# handle the border for the terrain
|
||||
heights[border_pixels:-border_pixels, border_pixels:-border_pixels] = z_gen
|
||||
# set terrain size back to config
|
||||
cfg.size = terrain_size
|
||||
|
||||
# PATCH: Use optimized mesh generation to reduce GPU memory usage
|
||||
vertices, triangles = _convert_height_field_to_mesh_with_optimization_dynamic(
|
||||
heights, cfg.horizontal_scale, cfg.vertical_scale, 20
|
||||
)
|
||||
mesh = trimesh.Trimesh(vertices=vertices, faces=triangles)
|
||||
|
||||
# compute origin
|
||||
x1 = int((cfg.size[0] * 0.5 - 1) / cfg.horizontal_scale)
|
||||
x2 = int((cfg.size[0] * 0.5 + 1) / cfg.horizontal_scale)
|
||||
y1 = int((cfg.size[1] * 0.5 - 1) / cfg.horizontal_scale)
|
||||
y2 = int((cfg.size[1] * 0.5 + 1) / cfg.horizontal_scale)
|
||||
origin_z = np.max(heights[x1:x2, y1:y2]) * cfg.vertical_scale
|
||||
origin = np.array([0.5 * cfg.size[0], 0.5 * cfg.size[1], origin_z])
|
||||
# return mesh and origin
|
||||
return [mesh], origin
|
||||
|
||||
return wrapper
|
||||
|
||||
# Apply patch to the height_field_to_mesh decorator
|
||||
hf_utils.height_field_to_mesh = _patched_height_field_to_mesh
|
||||
|
||||
|
||||
def _patch_terrain_generator():
|
||||
"""Patch TerrainGenerator to collect height field data for goal sampling."""
|
||||
from isaaclab.terrains.terrain_generator import TerrainGenerator
|
||||
from isaaclab.terrains.utils import color_meshes_by_height
|
||||
from isaaclab.utils.timer import Timer
|
||||
from isaaclab.utils.io import dump_yaml
|
||||
from isaaclab.utils.dict import dict_to_md5_hash
|
||||
|
||||
# Height field attribute names we collect
|
||||
_HEIGHT_FIELD_ATTRS = [
|
||||
"height_field_visual",
|
||||
"height_field_valid_mask",
|
||||
"height_field_platform_mask",
|
||||
"height_field_spawn_mask",
|
||||
]
|
||||
|
||||
def _patched_init(self, cfg, device: str = "cpu"):
|
||||
"""Patched __init__ that collects height field data."""
|
||||
# Validate inputs
|
||||
if len(cfg.sub_terrains) == 0:
|
||||
raise ValueError("No sub-terrains specified in the configuration!")
|
||||
if cfg.curriculum and cfg.difficulty_range is None:
|
||||
raise ValueError("Curriculum learning requires 'difficulty_range' to be specified!")
|
||||
|
||||
# Store inputs
|
||||
self.cfg = cfg
|
||||
self.device = device
|
||||
|
||||
# Set common values for all sub-terrains
|
||||
for sub_cfg in self.cfg.sub_terrains.values():
|
||||
sub_cfg.size = self.cfg.size
|
||||
sub_cfg.horizontal_scale = self.cfg.horizontal_scale
|
||||
sub_cfg.vertical_scale = self.cfg.vertical_scale
|
||||
sub_cfg.slope_threshold = self.cfg.slope_threshold
|
||||
|
||||
# Set seed for reproducibility
|
||||
# - seed=None (default): random terrain each time (for training variety)
|
||||
# - seed=<int>: reproducible terrain (for debugging/evaluation)
|
||||
# Note: Always create RNG because Isaac Lab's base code uses self.np_rng
|
||||
self.np_rng = np.random.default_rng(self.cfg.seed)
|
||||
self._reproducible = self.cfg.seed is not None
|
||||
|
||||
# Initialize buffers
|
||||
self.flat_patches = {}
|
||||
self.terrain_meshes = []
|
||||
self.terrain_origins = np.zeros((self.cfg.num_rows, self.cfg.num_cols, 3))
|
||||
|
||||
# PATCH: Temporary lists to collect height field data during generation
|
||||
self._height_field_lists = {attr: [] for attr in _HEIGHT_FIELD_ATTRS}
|
||||
|
||||
# Generate terrains
|
||||
if self.cfg.curriculum:
|
||||
with Timer("[INFO] Generating terrains based on curriculum took"):
|
||||
self._generate_curriculum_terrains()
|
||||
else:
|
||||
with Timer("[INFO] Generating terrains randomly took"):
|
||||
self._generate_random_terrains()
|
||||
|
||||
# Add border and combine meshes
|
||||
self._add_terrain_border()
|
||||
self.terrain_mesh = trimesh.util.concatenate(self.terrain_meshes)
|
||||
|
||||
# PATCH: Consolidate collected height fields into tensors and store in shared storage
|
||||
for attr in _HEIGHT_FIELD_ATTRS:
|
||||
data_list = self._height_field_lists[attr]
|
||||
if data_list:
|
||||
setattr(self, attr, torch.cat(data_list, dim=0))
|
||||
_height_field_storage[attr] = getattr(self, attr)
|
||||
else:
|
||||
setattr(self, attr, None)
|
||||
_height_field_storage[attr] = None
|
||||
del self._height_field_lists # Free memory
|
||||
|
||||
# Color the terrain mesh
|
||||
if self.cfg.color_scheme == "height":
|
||||
self.terrain_mesh = color_meshes_by_height(self.terrain_mesh)
|
||||
elif self.cfg.color_scheme == "random":
|
||||
self.terrain_mesh.visual.vertex_colors = self.np_rng.choice(
|
||||
range(256), size=(len(self.terrain_mesh.vertices), 4)
|
||||
)
|
||||
elif self.cfg.color_scheme != "none":
|
||||
raise ValueError(f"Unknown color scheme: {self.cfg.color_scheme}")
|
||||
|
||||
# Move flat patches to device
|
||||
for name in self.flat_patches:
|
||||
self.flat_patches[name] = self.flat_patches[name].to(self.device)
|
||||
|
||||
def _patched_get_terrain_mesh(self, difficulty: float, cfg) -> tuple:
|
||||
"""Patched _get_terrain_mesh that collects height field data from each terrain."""
|
||||
# Copy configuration and set parameters
|
||||
cfg = cfg.copy()
|
||||
cfg.difficulty = float(difficulty)
|
||||
cfg.seed = self.cfg.seed
|
||||
|
||||
# Clear non-serializable fields for hashing
|
||||
for attr in _HEIGHT_FIELD_ATTRS:
|
||||
if hasattr(cfg, attr):
|
||||
setattr(cfg, attr, None)
|
||||
# Clear RNG for hashing (seed already included, RNG is derived from it)
|
||||
if hasattr(cfg, 'rng'):
|
||||
cfg.rng = None
|
||||
|
||||
# Generate hash and cache paths
|
||||
sub_terrain_hash = dict_to_md5_hash(cfg.to_dict())
|
||||
sub_terrain_cache_dir = os.path.join(self.cfg.cache_dir, sub_terrain_hash)
|
||||
sub_terrain_obj_filename = os.path.join(sub_terrain_cache_dir, "mesh.obj")
|
||||
sub_terrain_csv_filename = os.path.join(sub_terrain_cache_dir, "origin.csv")
|
||||
sub_terrain_meta_filename = os.path.join(sub_terrain_cache_dir, "cfg.yaml")
|
||||
|
||||
# Check cache
|
||||
if self.cfg.use_cache and os.path.exists(sub_terrain_obj_filename):
|
||||
mesh = trimesh.load_mesh(sub_terrain_obj_filename, process=False)
|
||||
origin = np.loadtxt(sub_terrain_csv_filename, delimiter=",")
|
||||
return mesh, origin
|
||||
|
||||
# Set RNG for terrain generation (after hashing, before generation)
|
||||
# Only pass RNG if seed was explicitly set (for reproducibility)
|
||||
# Otherwise cfg.rng stays None and terrain uses fresh random each time
|
||||
if hasattr(cfg, 'rng') and self._reproducible:
|
||||
cfg.rng = self.np_rng.spawn(1)[0]
|
||||
|
||||
# Generate mesh
|
||||
meshes, origin = cfg.function(difficulty, cfg)
|
||||
if not isinstance(meshes, list):
|
||||
meshes = [meshes]
|
||||
mesh = trimesh.util.concatenate(meshes)
|
||||
|
||||
# Center the mesh
|
||||
transform = np.eye(4)
|
||||
transform[0:2, -1] = -cfg.size[0] * 0.5, -cfg.size[1] * 0.5
|
||||
mesh.apply_transform(transform)
|
||||
origin += transform[0:3, -1]
|
||||
|
||||
# PATCH: Collect height field data from cfg (set by terrain generation function)
|
||||
for attr in _HEIGHT_FIELD_ATTRS:
|
||||
if hasattr(cfg, attr):
|
||||
data = getattr(cfg, attr)
|
||||
if data is not None:
|
||||
self._height_field_lists[attr].append(data)
|
||||
setattr(cfg, attr, None) # Clear after collecting
|
||||
|
||||
# Cache if enabled
|
||||
if self.cfg.use_cache:
|
||||
os.makedirs(sub_terrain_cache_dir, exist_ok=True)
|
||||
mesh.export(sub_terrain_obj_filename)
|
||||
np.savetxt(sub_terrain_csv_filename, origin, delimiter=",", header="x,y,z")
|
||||
dump_yaml(sub_terrain_meta_filename, cfg)
|
||||
|
||||
return mesh, origin
|
||||
|
||||
# Apply patches
|
||||
TerrainGenerator.__init__ = _patched_init
|
||||
TerrainGenerator._get_terrain_mesh = _patched_get_terrain_mesh
|
||||
|
||||
|
||||
def _patch_terrain_importer():
|
||||
"""Patch TerrainImporter to store height field data as attributes.
|
||||
|
||||
After TerrainGenerator populates shared storage, we capture the data
|
||||
and store it as attributes on TerrainImporter for access via env.scene.terrain.
|
||||
"""
|
||||
from isaaclab.terrains.terrain_importer import TerrainImporter
|
||||
|
||||
_original_importer_init = TerrainImporter.__init__
|
||||
|
||||
# Attribute names (with underscore prefix for storage on importer)
|
||||
_HEIGHT_FIELD_ATTRS = [
|
||||
"height_field_visual",
|
||||
"height_field_valid_mask",
|
||||
"height_field_platform_mask",
|
||||
"height_field_spawn_mask",
|
||||
]
|
||||
|
||||
def _patched_importer_init(self, cfg):
|
||||
"""Patched __init__ that captures height field data from TerrainGenerator."""
|
||||
# Clear previous storage
|
||||
for attr in _HEIGHT_FIELD_ATTRS:
|
||||
_height_field_storage[attr] = None
|
||||
|
||||
# Call original __init__ - creates TerrainGenerator which populates storage
|
||||
_original_importer_init(self, cfg)
|
||||
|
||||
# Capture height field data from storage and store as attributes
|
||||
for attr in _HEIGHT_FIELD_ATTRS:
|
||||
setattr(self, f"_{attr}", _height_field_storage.get(attr))
|
||||
|
||||
TerrainImporter.__init__ = _patched_importer_init
|
||||
|
||||
|
||||
def _patch_pinhole_camera_pattern_cfg():
|
||||
"""Patch PinholeCameraPatternCfg to add from_ros_camera_info method."""
|
||||
from isaaclab.sensors.ray_caster.patterns.patterns_cfg import PinholeCameraPatternCfg
|
||||
|
||||
# Check if already patched
|
||||
if hasattr(PinholeCameraPatternCfg, 'from_ros_camera_info'):
|
||||
return
|
||||
|
||||
# Add downsample_factor attribute if not present
|
||||
if not hasattr(PinholeCameraPatternCfg, 'downsample_factor'):
|
||||
PinholeCameraPatternCfg.downsample_factor = 1
|
||||
|
||||
# Store original from_intrinsic_matrix
|
||||
_original_from_intrinsic_matrix = PinholeCameraPatternCfg.from_intrinsic_matrix
|
||||
|
||||
@classmethod
|
||||
def _patched_from_intrinsic_matrix(
|
||||
cls,
|
||||
intrinsic_matrix: list[float],
|
||||
width: int,
|
||||
height: int,
|
||||
focal_length: float = 1.0,
|
||||
downsample_factor: int = 1,
|
||||
) -> PinholeCameraPatternCfg:
|
||||
"""Create a PinholeCameraPatternCfg from an intrinsic matrix with downsampling support.
|
||||
|
||||
Args:
|
||||
intrinsic_matrix: The intrinsic matrix as a 9-element list [f_x, 0, c_x, 0, f_y, c_y, 0, 0, 1].
|
||||
width: Width of the image (in pixels).
|
||||
height: Height of the image (in pixels).
|
||||
focal_length: Focal length of the camera (in cm). Defaults to 1.0 cm.
|
||||
downsample_factor: Downsampling factor for RL training. Defaults to 1 (no downsampling).
|
||||
|
||||
Returns:
|
||||
An instance of the PinholeCameraPatternCfg class.
|
||||
"""
|
||||
# Extract standard intrinsic parameters (in pixels)
|
||||
f_x = intrinsic_matrix[0]
|
||||
c_x = intrinsic_matrix[2]
|
||||
f_y = intrinsic_matrix[4]
|
||||
c_y = intrinsic_matrix[5]
|
||||
|
||||
# Apply downsampling adjustments
|
||||
if downsample_factor > 1:
|
||||
f_x = f_x / downsample_factor
|
||||
f_y = f_y / downsample_factor
|
||||
c_x = c_x / downsample_factor
|
||||
c_y = c_y / downsample_factor
|
||||
width = width // downsample_factor
|
||||
height = height // downsample_factor
|
||||
|
||||
# Convert to USD camera parameters
|
||||
horizontal_aperture = width * focal_length / f_x
|
||||
vertical_aperture = height * focal_length / f_y
|
||||
|
||||
# Convert principal point offset from pixels to physical units (cm)
|
||||
horizontal_aperture_offset = (c_x - width / 2) * horizontal_aperture / width
|
||||
vertical_aperture_offset = (c_y - height / 2) * vertical_aperture / height
|
||||
|
||||
return cls(
|
||||
focal_length=focal_length,
|
||||
horizontal_aperture=horizontal_aperture,
|
||||
vertical_aperture=vertical_aperture,
|
||||
horizontal_aperture_offset=horizontal_aperture_offset,
|
||||
vertical_aperture_offset=vertical_aperture_offset,
|
||||
width=width,
|
||||
height=height,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_ros_camera_info(
|
||||
cls,
|
||||
fx: float,
|
||||
fy: float,
|
||||
cx: float,
|
||||
cy: float,
|
||||
width: int,
|
||||
height: int,
|
||||
focal_length: float = 1.0,
|
||||
downsample_factor: int = 1,
|
||||
) -> PinholeCameraPatternCfg:
|
||||
"""Create a PinholeCameraPatternCfg from ROS CameraInfo parameters.
|
||||
|
||||
This is a convenience method for creating camera configurations directly from
|
||||
ROS camera_info topic parameters, avoiding the need to construct the full
|
||||
intrinsic matrix.
|
||||
|
||||
Args:
|
||||
fx: Focal length in x-direction (pixels).
|
||||
fy: Focal length in y-direction (pixels).
|
||||
cx: Principal point x-coordinate (pixels).
|
||||
cy: Principal point y-coordinate (pixels).
|
||||
width: Image width (pixels).
|
||||
height: Image height (pixels).
|
||||
focal_length: USD focal length scaling factor (cm). Defaults to 1.0 cm.
|
||||
downsample_factor: Downsampling factor for RL training. Defaults to 1 (no downsampling).
|
||||
|
||||
Returns:
|
||||
A PinholeCameraPatternCfg instance.
|
||||
|
||||
Example:
|
||||
# From your camera info topic:
|
||||
# K: [419.663, 0, 425.728, 0, 419.663, 238.272, 0, 0, 1]
|
||||
# size: 848 x 480
|
||||
|
||||
# Original resolution
|
||||
cfg = PinholeCameraPatternCfg.from_ros_camera_info(
|
||||
fx=419.663, fy=419.663, cx=425.728, cy=238.272,
|
||||
width=848, height=480
|
||||
)
|
||||
|
||||
# 4x downsampled for RL training (212x120)
|
||||
cfg_rl = PinholeCameraPatternCfg.from_ros_camera_info(
|
||||
fx=419.663, fy=419.663, cx=425.728, cy=238.272,
|
||||
width=848, height=480, downsample_factor=4
|
||||
)
|
||||
"""
|
||||
# Create intrinsic matrix in row-major format
|
||||
intrinsic_matrix = [fx, 0, cx, 0, fy, cy, 0, 0, 1]
|
||||
return cls.from_intrinsic_matrix(intrinsic_matrix, width, height, focal_length, downsample_factor)
|
||||
|
||||
# Apply patches
|
||||
PinholeCameraPatternCfg.from_intrinsic_matrix = _patched_from_intrinsic_matrix
|
||||
PinholeCameraPatternCfg.from_ros_camera_info = from_ros_camera_info
|
||||
@@ -0,0 +1,341 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Constants for terrain generation and goal sampling.
|
||||
|
||||
This module centralizes all height field values and terrain parameters used across
|
||||
terrain generation and goal sampling code. Using these constants ensures consistency
|
||||
and makes the height semantics explicit.
|
||||
|
||||
Height Field Value Semantics (in discretized units):
|
||||
====================================================
|
||||
The height field stores heights as integers where:
|
||||
actual_height_meters = height_value * VERTICAL_SCALE
|
||||
|
||||
Key height values:
|
||||
- GROUND: 0 (flat walkable ground at z=0)
|
||||
- PLATFORM: ~200 (flat raised platforms at ~1.0m, valid for goals)
|
||||
- WALL: ~300 (obstacles/walls at ~1.5m, always excluded from goals)
|
||||
- PIT: ~-300 (negative obstacles/troughs at ~-1.5m, always excluded)
|
||||
|
||||
Goal Sampling Valid Ranges:
|
||||
- Ground: -10 to 50 (allows small noise/variation)
|
||||
- Platform: 150 to 250 (captures platform height with margin)
|
||||
- Excluded: < -10 (pits) or > 250 (walls)
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Scale Factors
|
||||
# =============================================================================
|
||||
|
||||
HORIZONTAL_SCALE: float = 0.1
|
||||
"""Horizontal resolution of height field in meters per cell."""
|
||||
|
||||
VERTICAL_SCALE: float = 0.005
|
||||
"""Vertical resolution of height field (height_meters = height_value * VERTICAL_SCALE)."""
|
||||
|
||||
CELL_SIZE: float = 2.0
|
||||
"""Size of each maze cell in meters (default: 2m x 2m)."""
|
||||
|
||||
CELL_PIXELS: int = int(CELL_SIZE / HORIZONTAL_SCALE)
|
||||
"""Number of height field pixels per maze cell (20 at default scale)."""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Height Values (in discretized units, multiply by VERTICAL_SCALE for meters)
|
||||
# =============================================================================
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HeightValues:
|
||||
"""Standard height values for terrain features."""
|
||||
|
||||
# Ground level (valid for goals/spawn)
|
||||
GROUND: int = 0
|
||||
|
||||
# Platform height: wall_height - 0.5m in vertical scale units
|
||||
# = 1.5/0.005 - 0.5/0.005 = 300 - 100 = 200
|
||||
PLATFORM: int = 200
|
||||
|
||||
# Wall/obstacle height: 1.5m in vertical scale units
|
||||
# = 1.5/0.005 = 300
|
||||
WALL: int = 300
|
||||
|
||||
# Pit/trough height (negative obstacle)
|
||||
PIT: int = -300
|
||||
|
||||
@property
|
||||
def platform_meters(self) -> float:
|
||||
"""Platform height in meters."""
|
||||
return self.PLATFORM * VERTICAL_SCALE
|
||||
|
||||
@property
|
||||
def wall_meters(self) -> float:
|
||||
"""Wall height in meters."""
|
||||
return self.WALL * VERTICAL_SCALE
|
||||
|
||||
|
||||
HEIGHTS = HeightValues()
|
||||
"""Singleton instance of standard height values."""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Goal Sampling Thresholds
|
||||
# =============================================================================
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SamplingThresholds:
|
||||
"""Thresholds for classifying height values during goal/spawn sampling.
|
||||
|
||||
These values match the old implementation for strict height classification.
|
||||
"""
|
||||
|
||||
# Ground range (strict: 0-10, allows minimal noise)
|
||||
GROUND_MIN: int = 0
|
||||
GROUND_MAX: int = 10
|
||||
|
||||
# Platform range (strict: 195-205 for stair platforms at ~1.0m)
|
||||
PLATFORM_MIN: int = 195
|
||||
PLATFORM_MAX: int = 205
|
||||
|
||||
# Wall threshold (anything above 10 is obstacle, but also check platform edges)
|
||||
WALL_THRESHOLD: int = 10
|
||||
|
||||
# Pit threshold (anything below 0 is pit/excluded)
|
||||
PIT_THRESHOLD: int = 0
|
||||
|
||||
# Edge detection threshold for Sobel filter (old: any edge > 0.0)
|
||||
EDGE_THRESHOLD: float = 0.0
|
||||
|
||||
# Obstacle marker value (used to mark excluded cells)
|
||||
OBSTACLE_MARKER: int = 110
|
||||
|
||||
# Extended platform range for edge detection (180-220)
|
||||
# Used during padding to detect platform edges with safety margin
|
||||
PLATFORM_EDGE_MIN: int = 180
|
||||
PLATFORM_EDGE_MAX: int = 220
|
||||
|
||||
|
||||
THRESHOLDS = SamplingThresholds()
|
||||
"""Singleton instance of sampling thresholds."""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Padding and Border Configuration
|
||||
# =============================================================================
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PaddingConfig:
|
||||
"""Configuration for obstacle padding and borders.
|
||||
|
||||
These values match the old implementation for consistent safety margins.
|
||||
With horizontal_scale=0.1m/cell, these translate to:
|
||||
- GOAL_PADDING (3 cells) = 0.3m around obstacles for goal positions
|
||||
- SPAWN_PADDING (8 cells) = 0.8m around obstacles for spawn positions
|
||||
|
||||
The larger spawn padding accounts for:
|
||||
- Robot body dimensions (~0.5m x 0.3m for quadrupeds)
|
||||
- Random yaw orientation (diagonal ~0.58m requires ~0.3m radius)
|
||||
- Platform edge clearance (prevent falling when spawning near stair edges)
|
||||
- Controller startup imprecision
|
||||
"""
|
||||
|
||||
# Padding around obstacles for goals (in height field cells)
|
||||
# 5 cells * 0.1m/cell = 0.5m padding (robot just needs to reach)
|
||||
GOAL_PADDING: int = 5
|
||||
|
||||
# Larger padding for spawn positions (in height field cells)
|
||||
# 6 cells * 0.1m/cell = 0.6m padding
|
||||
# This accounts for:
|
||||
# - Robot body radius (~0.5m) with worst-case orientation
|
||||
# - Additional safety margin for platform edges
|
||||
# - Controller startup behavior
|
||||
SPAWN_PADDING: int = 6
|
||||
|
||||
# Border exclusion (keep away from terrain edges)
|
||||
# Old implementation used 2 layers (inner + outer), we use 2 cells
|
||||
BORDER_CELLS: int = 2
|
||||
|
||||
# Pillar expanded footprint size (5x5 cells around pillar center)
|
||||
PILLAR_FOOTPRINT: int = 5
|
||||
|
||||
# Pillar safe margin from cell edges
|
||||
PILLAR_EDGE_MARGIN: int = 2
|
||||
|
||||
# Height transition detection threshold (height field units)
|
||||
# Marks cells with height differences >= this value as transitions
|
||||
HEIGHT_TRANSITION_THRESHOLD: int = 100
|
||||
|
||||
# Padding around detected height transitions (in cells)
|
||||
HEIGHT_TRANSITION_PADDING: int = 1
|
||||
|
||||
|
||||
PADDING = PaddingConfig()
|
||||
"""Singleton instance of padding configuration."""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Stair Configuration
|
||||
# =============================================================================
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StairConfig:
|
||||
"""Configuration for stair generation."""
|
||||
|
||||
# Number of steps in a staircase
|
||||
NUM_STEPS: int = 5
|
||||
|
||||
# Height of each step in meters
|
||||
STEP_HEIGHT_METERS: float = 0.2
|
||||
|
||||
# Grid size for 3x3 stair structures
|
||||
STAIR_GRID_SIZE: int = 3
|
||||
|
||||
# Single cell size in pixels (same as CELL_PIXELS)
|
||||
SINGLE_CELL_PIXELS: int = 20
|
||||
|
||||
@property
|
||||
def step_height_units(self) -> float:
|
||||
"""Step height in discretized units."""
|
||||
return self.STEP_HEIGHT_METERS / VERTICAL_SCALE
|
||||
|
||||
@property
|
||||
def step_resolution(self) -> int:
|
||||
"""Pixels per step (cell_pixels / num_steps)."""
|
||||
return self.SINGLE_CELL_PIXELS // self.NUM_STEPS
|
||||
|
||||
|
||||
STAIRS = StairConfig()
|
||||
"""Singleton instance of stair configuration."""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Obstacle Structure Types
|
||||
# =============================================================================
|
||||
|
||||
class ObstacleType:
|
||||
"""Enumeration of obstacle structure types for random generation."""
|
||||
|
||||
PILLAR = 0
|
||||
BAR = 1
|
||||
CROSS = 2
|
||||
SHIFTED_BLOCK = 3
|
||||
|
||||
# Total number of obstacle types
|
||||
NUM_TYPES = 4
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Obstacle Generation Parameters
|
||||
# =============================================================================
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ObstacleConfig:
|
||||
"""Configuration for random obstacle generation.
|
||||
|
||||
Controls the randomness parameters for obstacle shapes, sizes, and pit ratios.
|
||||
All probabilities are in [0, 1] range.
|
||||
"""
|
||||
|
||||
# Height scale range (multiplier for wall_height)
|
||||
SCALE_MIN: float = 0.5
|
||||
SCALE_MAX: float = 1.5
|
||||
|
||||
# Thickness range (margin cells from obstacle center)
|
||||
# Higher thickness = thinner pillar, thicker bar
|
||||
# Range 7-9: pillars 2-6 pixels (0.2-0.6m), bars 7-9 pixels (0.7-0.9m)
|
||||
THICKNESS_MIN: int = 7
|
||||
THICKNESS_MAX: int = 10 # exclusive upper bound
|
||||
|
||||
# Default pit probability when not specified
|
||||
DEFAULT_PIT_PROB: float = 0.15
|
||||
|
||||
# Pit environment specific settings
|
||||
PITS_BAR_RATIO: float = 0.6 # Ratio of bars vs random obstacles in pit terrain
|
||||
PITS_BAR_PIT_PROB: float = 0.75 # Pit probability for bar obstacles
|
||||
PITS_RANDOM_PIT_PROB: float = 0.5 # Pit probability for random obstacles
|
||||
|
||||
# Bridge configuration for pit terrain
|
||||
BRIDGE_COUNT_MIN: int = 3
|
||||
BRIDGE_COUNT_MAX: int = 6 # exclusive upper bound
|
||||
|
||||
# Pit terrain layout (cell indices from grid edge)
|
||||
PITS_TRENCH_ROW_OFFSET: int = 2 # Pit trenches at rows [offset, grid_h - offset - 1]
|
||||
PITS_EDGE_MARGIN: int = 2 # Margin from grid edges for middle obstacles
|
||||
|
||||
# Obstacle density multipliers (applied to difficulty)
|
||||
NON_MAZE_DENSITY: float = 0.5 # For non-maze terrain (increased from 0.35)
|
||||
PITS_DENSITY: float = 0.6 # For pit terrain middle area
|
||||
STAIRS_OBSTACLE_DENSITY: float = 0.35 # For random obstacles in stair terrain
|
||||
STAIRS_PLACEMENT_PROB: float = 0.75 # Probability of placing stairs at valid locations
|
||||
|
||||
# Non-maze terrain pillar weight (higher = more pillars)
|
||||
# Default uniform is 0.25 (1/4 types), 0.5 means ~50% pillars
|
||||
NON_MAZE_PILLAR_WEIGHT: float = 0.5
|
||||
|
||||
|
||||
OBSTACLES = ObstacleConfig()
|
||||
"""Singleton instance of obstacle configuration."""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Helper Functions
|
||||
# =============================================================================
|
||||
|
||||
def height_to_meters(height_value: int) -> float:
|
||||
"""Convert discretized height value to meters."""
|
||||
return height_value * VERTICAL_SCALE
|
||||
|
||||
|
||||
def meters_to_height(meters: float) -> int:
|
||||
"""Convert meters to discretized height value."""
|
||||
return int(meters / VERTICAL_SCALE)
|
||||
|
||||
|
||||
def is_valid_ground(height: int) -> bool:
|
||||
"""Check if height value represents valid ground."""
|
||||
return THRESHOLDS.GROUND_MIN <= height <= THRESHOLDS.GROUND_MAX
|
||||
|
||||
|
||||
def is_valid_platform(height: int) -> bool:
|
||||
"""Check if height value represents valid platform."""
|
||||
return THRESHOLDS.PLATFORM_MIN <= height <= THRESHOLDS.PLATFORM_MAX
|
||||
|
||||
|
||||
def is_valid_goal_position(height: int) -> bool:
|
||||
"""Check if height value is valid for goal/spawn placement."""
|
||||
return is_valid_ground(height) or is_valid_platform(height)
|
||||
|
||||
|
||||
def is_obstacle(height: int) -> bool:
|
||||
"""Check if height value represents an obstacle (wall or pit)."""
|
||||
return height > THRESHOLDS.WALL_THRESHOLD or height < THRESHOLDS.PIT_THRESHOLD
|
||||
|
||||
|
||||
def is_pit(height: int) -> bool:
|
||||
"""Check if height value represents a pit."""
|
||||
return height < THRESHOLDS.PIT_THRESHOLD
|
||||
|
||||
|
||||
def is_wall(height: int) -> bool:
|
||||
"""Check if height value represents a wall."""
|
||||
return height > THRESHOLDS.WALL_THRESHOLD
|
||||
|
||||
|
||||
def cell_to_pixels(cell_idx: int) -> Tuple[int, int]:
|
||||
"""Convert cell index to pixel range.
|
||||
|
||||
Args:
|
||||
cell_idx: Cell index in maze grid.
|
||||
|
||||
Returns:
|
||||
Tuple of (start_pixel, end_pixel).
|
||||
"""
|
||||
start = cell_idx * CELL_PIXELS
|
||||
end = (cell_idx + 1) * CELL_PIXELS
|
||||
return start, end
|
||||
@@ -0,0 +1,3 @@
|
||||
[build-system]
|
||||
requires = ["setuptools", "toml"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
@@ -0,0 +1,221 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Depth-ablation diagnostic for a trained navigation policy.
|
||||
|
||||
Question this answers: *does the policy actually use the depth-camera input,
|
||||
or has it learned to ignore it (perception-blind)?*
|
||||
|
||||
Method
|
||||
------
|
||||
1. Load the trained actor-critic exactly like ``play.py``.
|
||||
2. Locate the ``depth_image`` slice inside the concatenated *policy* observation.
|
||||
3. Roll the env out with the REAL on-policy actions, and at every step also
|
||||
compute a counterfactual action where the depth slice is zeroed. Compare.
|
||||
4. Also track the depth-feature statistics (std over envs/time) so we can tell
|
||||
a *dead sensor* (constant features) apart from a policy that *ignores*
|
||||
informative features.
|
||||
|
||||
Verdict logic (printed at the end)
|
||||
----------------------------------
|
||||
* depth features ~constant (std ~ 0) -> SENSOR/ENCODER dead (no signal)
|
||||
* features vary, zeroing barely changes action -> LEARNED-BLIND (ignores depth)
|
||||
* features vary, zeroing changes action a lot -> policy DOES use depth
|
||||
|
||||
Usage (inside container, mirrors play.py):
|
||||
./isaaclab.sh -p source/isaaclab_nav_task/scripts/ablate_depth.py \
|
||||
--task Isaac-Nav-PPO-Go2-Play-v0 --num_envs 16 \
|
||||
--checkpoint <path/model_6998.pt> --steps 300 --headless
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
parser = argparse.ArgumentParser(description="Depth-ablation diagnostic for a navigation policy.")
|
||||
parser.add_argument("--num_envs", type=int, default=16, help="Number of environments to simulate.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Name of the task (use a *-Play-v0 variant).")
|
||||
parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint (.pt).")
|
||||
parser.add_argument("--steps", type=int, default=300, help="Number of env steps to roll out.")
|
||||
parser.add_argument("--warmup", type=int, default=20, help="Steps to skip before collecting stats (let resets settle).")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Seed.")
|
||||
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli, hydra_args = parser.parse_known_args()
|
||||
|
||||
# Cameras are mandatory for the depth observation.
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
import gymnasium as gym
|
||||
import torch
|
||||
|
||||
from rsl_rl.runners import OnPolicyRunner
|
||||
|
||||
import isaaclab_tasks # noqa: F401
|
||||
import isaaclab_nav_task # noqa: F401
|
||||
|
||||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
|
||||
|
||||
|
||||
def _find_depth_slice(env) -> tuple[int, int, list[tuple[str, int]]]:
|
||||
"""Return (start, end, layout) of the ``depth_image`` term in the policy obs.
|
||||
|
||||
layout is a list of (term_name, flat_dim) in concatenation order, for printing.
|
||||
"""
|
||||
om = env.unwrapped.observation_manager
|
||||
names = om.active_terms["policy"]
|
||||
dims = om.group_obs_term_dim["policy"] # list of shape-tuples per term
|
||||
|
||||
layout: list[tuple[str, int]] = []
|
||||
offset = 0
|
||||
depth_start = depth_end = -1
|
||||
for name, shape in zip(names, dims):
|
||||
flat = 1
|
||||
for s in shape:
|
||||
flat *= s
|
||||
layout.append((name, flat))
|
||||
if name == "depth_image":
|
||||
depth_start, depth_end = offset, offset + flat
|
||||
offset += flat
|
||||
if depth_start < 0:
|
||||
raise RuntimeError(
|
||||
f"No 'depth_image' term found in policy obs group. Terms = {names}"
|
||||
)
|
||||
return depth_start, depth_end, layout
|
||||
|
||||
|
||||
def main():
|
||||
spec = gym.spec(args_cli.task)
|
||||
env_cfg_class = spec.kwargs.get("env_cfg_entry_point")
|
||||
agent_cfg_class = spec.kwargs.get("rsl_rl_cfg_entry_point")
|
||||
|
||||
env_cfg: ManagerBasedRLEnvCfg = env_cfg_class()
|
||||
agent_cfg: RslRlOnPolicyRunnerCfg = agent_cfg_class()
|
||||
|
||||
if args_cli.num_envs is not None:
|
||||
env_cfg.scene.num_envs = args_cli.num_envs
|
||||
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None)
|
||||
env = RslRlVecEnvWrapper(env)
|
||||
|
||||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
|
||||
|
||||
# --- load checkpoint (same logic as play.py) ---
|
||||
resume_path = args_cli.checkpoint
|
||||
if not resume_path:
|
||||
raise SystemExit("[ablate] --checkpoint is required for the diagnostic.")
|
||||
print(f"[ablate] loading checkpoint: {resume_path}")
|
||||
loaded = torch.load(resume_path, map_location="cpu", weights_only=False)
|
||||
runner.alg.actor_critic.load_state_dict(loaded["model_state_dict"], strict=True)
|
||||
if runner.empirical_normalization:
|
||||
runner.obs_normalizer.load_state_dict(loaded["obs_norm_state_dict"])
|
||||
runner.critic_obs_normalizer.load_state_dict(loaded["critic_obs_norm_state_dict"])
|
||||
print(f"[ablate] checkpoint iter = {loaded.get('iter', '?')}")
|
||||
|
||||
policy = runner.get_inference_policy(device=env.unwrapped.device)
|
||||
|
||||
depth_start, depth_end, layout = _find_depth_slice(env)
|
||||
print("=" * 78)
|
||||
print("[ablate] policy observation layout (term : flat_dim):")
|
||||
for name, flat in layout:
|
||||
marker = " <-- DEPTH" if name == "depth_image" else ""
|
||||
print(f" {name:<28} {flat:>5}{marker}")
|
||||
print(f"[ablate] depth slice = [{depth_start}:{depth_end}] ({depth_end - depth_start} dims)")
|
||||
print("=" * 78)
|
||||
|
||||
obs, _ = env.get_observations()
|
||||
|
||||
# accumulators
|
||||
n = 0
|
||||
act_norm_sum = 0.0 # ||a_real||
|
||||
diff_norm_sum = 0.0 # ||a_real - a_zero||
|
||||
per_dim_abs_diff = None # mean |a_real - a_zero| per action dim
|
||||
per_dim_abs_act = None # mean |a_real| per action dim
|
||||
depth_feat_chunks = [] # collect depth slices to measure variability
|
||||
|
||||
device = env.unwrapped.device
|
||||
step = 0
|
||||
while simulation_app.is_running() and step < args_cli.steps:
|
||||
with torch.inference_mode():
|
||||
a_real = policy(obs)
|
||||
obs_zero = obs.clone()
|
||||
obs_zero[:, depth_start:depth_end] = 0.0
|
||||
a_zero = policy(obs_zero)
|
||||
|
||||
if step >= args_cli.warmup:
|
||||
d = (a_real - a_zero)
|
||||
diff_norm_sum += torch.linalg.vector_norm(d, dim=-1).mean().item()
|
||||
act_norm_sum += torch.linalg.vector_norm(a_real, dim=-1).mean().item()
|
||||
ad = d.abs().mean(dim=0)
|
||||
aa = a_real.abs().mean(dim=0)
|
||||
per_dim_abs_diff = ad if per_dim_abs_diff is None else per_dim_abs_diff + ad
|
||||
per_dim_abs_act = aa if per_dim_abs_act is None else per_dim_abs_act + aa
|
||||
depth_feat_chunks.append(obs[:, depth_start:depth_end].detach().clone())
|
||||
n += 1
|
||||
|
||||
with torch.inference_mode():
|
||||
obs, _, _, _ = env.step(a_real)
|
||||
step += 1
|
||||
|
||||
# --- summarise ---
|
||||
print("\n" + "=" * 78)
|
||||
print(f"[ablate] collected {n} steps (after {args_cli.warmup} warmup), {args_cli.num_envs} envs")
|
||||
if n == 0:
|
||||
print("[ablate] no samples collected; increase --steps.")
|
||||
else:
|
||||
mean_act = act_norm_sum / n
|
||||
mean_diff = diff_norm_sum / n
|
||||
rel = (mean_diff / mean_act) if mean_act > 0 else float("nan")
|
||||
per_dim_abs_diff = (per_dim_abs_diff / n).tolist()
|
||||
per_dim_abs_act = (per_dim_abs_act / n).tolist()
|
||||
|
||||
feats = torch.cat(depth_feat_chunks, dim=0) # (n*envs, depth_dim)
|
||||
feat_std_over_samples = feats.std(dim=0).mean().item() # variability across time/envs
|
||||
feat_abs_mean = feats.abs().mean().item()
|
||||
|
||||
names_se2 = ["vx", "vy", "omega"]
|
||||
print("\n-- ACTION SENSITIVITY TO DEPTH (real vs depth-zeroed) --")
|
||||
print(f" mean ||a_real|| = {mean_act:.4f}")
|
||||
print(f" mean ||a_real - a_zero|| = {mean_diff:.4f}")
|
||||
print(f" RELATIVE action change = {rel*100:.2f}%")
|
||||
print(" per-action-dim |diff| / |act|:")
|
||||
for i, v in enumerate(per_dim_abs_diff):
|
||||
nm = names_se2[i] if i < len(names_se2) else f"a{i}"
|
||||
denom = per_dim_abs_act[i] if per_dim_abs_act[i] > 0 else float("nan")
|
||||
print(f" {nm:<6} |diff|={v:.4f} |act|={per_dim_abs_act[i]:.4f} ratio={v/denom*100:.1f}%")
|
||||
|
||||
print("\n-- DEPTH FEATURE SIGNAL (is the camera/encoder alive?) --")
|
||||
print(f" feature |mean| = {feat_abs_mean:.4f}")
|
||||
print(f" feature std over time/envs = {feat_std_over_samples:.4f}")
|
||||
|
||||
print("\n-- VERDICT --")
|
||||
if feat_std_over_samples < 1e-3:
|
||||
print(" >> DEAD SIGNAL: depth features are ~constant. Camera/encoder is")
|
||||
print(" not producing a varying signal (sensor/render/orientation bug).")
|
||||
elif rel < 0.02:
|
||||
print(" >> LEARNED-BLIND: features vary but zeroing them changes the action")
|
||||
print(f" by only {rel*100:.2f}%. The policy effectively IGNORES depth.")
|
||||
elif rel < 0.10:
|
||||
print(f" >> WEAK USE: depth changes the action by {rel*100:.2f}% (marginal).")
|
||||
else:
|
||||
print(f" >> DEPTH IS USED: zeroing depth changes the action by {rel*100:.2f}%.")
|
||||
print(" Perception is wired in; the failure is likely a difficulty/")
|
||||
print(" generalisation gap, not blindness.")
|
||||
print("=" * 78)
|
||||
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
simulation_app.close()
|
||||
@@ -0,0 +1,325 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Instrumented collision analysis for a trained navigation policy.
|
||||
|
||||
For every env-step, this script records:
|
||||
* base_lin_vel_b[:, 0] -- forward velocity in body frame (m/s)
|
||||
* max contact-force magnitude over the *termination sensor body subset*
|
||||
(whatever ``terminations.base_contact.params['sensor_cfg'].body_ids`` is)
|
||||
* the per-term ``base_contact`` done flag
|
||||
|
||||
It maintains a per-env rolling buffer of the last K=16 steps so that, when a
|
||||
collision fires for env i, we can dump the *pre-impact trajectory* (velocity +
|
||||
peak force) for that env. The aggregated data is saved as a single .npz which
|
||||
the companion script ``plot_collisions.py`` turns into figures.
|
||||
|
||||
Outputs (under --out-dir, default ``logs/collisions/<task>__<ckpt>``):
|
||||
* collisions.npz full dump (see save_data())
|
||||
* summary.txt human-readable numbers
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--num_envs", type=int, default=32)
|
||||
parser.add_argument("--task", type=str, required=True)
|
||||
parser.add_argument("--checkpoint", type=str, required=True)
|
||||
parser.add_argument("--steps", type=int, default=4000)
|
||||
parser.add_argument("--warmup", type=int, default=20)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--lookback", type=int, default=16,
|
||||
help="How many pre-impact steps to capture per collision.")
|
||||
parser.add_argument("--out-dir", type=str, default=None)
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli, _ = parser.parse_known_args()
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from rsl_rl.runners import OnPolicyRunner
|
||||
|
||||
import isaaclab_tasks # noqa: F401
|
||||
import isaaclab_nav_task # noqa: F401
|
||||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
from isaaclab.sensors import ContactSensor
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
|
||||
|
||||
|
||||
def main():
|
||||
spec = gym.spec(args_cli.task)
|
||||
env_cfg: ManagerBasedRLEnvCfg = spec.kwargs["env_cfg_entry_point"]()
|
||||
agent_cfg: RslRlOnPolicyRunnerCfg = spec.kwargs["rsl_rl_cfg_entry_point"]()
|
||||
env_cfg.scene.num_envs = args_cli.num_envs
|
||||
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None)
|
||||
env = RslRlVecEnvWrapper(env)
|
||||
device = env.unwrapped.device
|
||||
N = args_cli.num_envs
|
||||
|
||||
# ---- load policy ----
|
||||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
|
||||
loaded = torch.load(args_cli.checkpoint, map_location="cpu", weights_only=False)
|
||||
runner.alg.actor_critic.load_state_dict(loaded["model_state_dict"], strict=True)
|
||||
if runner.empirical_normalization:
|
||||
runner.obs_normalizer.load_state_dict(loaded["obs_norm_state_dict"])
|
||||
runner.critic_obs_normalizer.load_state_dict(loaded["critic_obs_norm_state_dict"])
|
||||
print(f"[collisions] checkpoint iter = {loaded.get('iter', '?')}")
|
||||
policy = runner.get_inference_policy(device=device)
|
||||
|
||||
# ---- resolve the *exact* body subset used by the base_contact termination ----
|
||||
tm = env.unwrapped.termination_manager
|
||||
tf = tm.get_term_cfg("base_contact")
|
||||
sensor_cfg: SceneEntityCfg = tf.params["sensor_cfg"]
|
||||
threshold = float(tf.params["threshold"])
|
||||
print(f"[collisions] base_contact threshold = {threshold} N")
|
||||
print(f"[collisions] base_contact body_names spec = {sensor_cfg.body_names}")
|
||||
|
||||
# NOTE: TerminationManager already resolved this SceneEntityCfg at env init,
|
||||
# so body_ids is already populated; calling .resolve() again would raise.
|
||||
body_ids = sensor_cfg.body_ids
|
||||
contact_sensor: ContactSensor = env.unwrapped.scene.sensors[sensor_cfg.name]
|
||||
print(f"[collisions] resolved body_ids = {list(body_ids) if not isinstance(body_ids, slice) else body_ids}")
|
||||
|
||||
# episode_termination reward weight (for context in summary)
|
||||
try:
|
||||
ep_term_w = float(env.unwrapped.reward_manager.get_term_cfg("episode_termination").weight)
|
||||
except Exception:
|
||||
ep_term_w = float("nan")
|
||||
|
||||
robot = env.unwrapped.scene["robot"]
|
||||
|
||||
# ---- rolling buffers (K x N) for pre-impact dumps ----
|
||||
K = args_cli.lookback
|
||||
vel_buf = torch.zeros(K, N, device=device) # base lin vel x (body frame)
|
||||
speed_buf = torch.zeros(K, N, device=device) # |v_xy| (body frame)
|
||||
force_buf = torch.zeros(K, N, device=device) # max contact force on subset
|
||||
head = 0 # circular index
|
||||
|
||||
# ---- collected data ----
|
||||
coll_traj_vel = [] # list of (K,) numpy arrays at the impact moment
|
||||
coll_traj_speed = []
|
||||
coll_traj_force = []
|
||||
coll_impact_vel = [] # scalar at the IMPACT step (after termination flag flips)
|
||||
coll_impact_speed = []
|
||||
coll_impact_force = []
|
||||
|
||||
# global histograms over all NON-terminating steps (for context)
|
||||
all_vel_chunks = []
|
||||
all_force_chunks = []
|
||||
|
||||
n_steps_counted = 0
|
||||
n_collisions = 0
|
||||
n_success = 0
|
||||
n_timeout = 0
|
||||
|
||||
obs, _ = env.get_observations()
|
||||
step = 0
|
||||
while simulation_app.is_running() and step < args_cli.steps:
|
||||
# -------------------------------------------------------------
|
||||
# CRITICAL: measure state BEFORE env.step(), because IsaacLab
|
||||
# auto-resets terminated envs INSIDE env.step() and that wipes
|
||||
# both robot.data and contact_sensor.data for the dying envs.
|
||||
# We do NOT lose the impact data: the contact force history
|
||||
# buffer (decimation samples) accumulates within env.step(),
|
||||
# so reading it on the *next* iteration's pre-step pass would
|
||||
# be the natural way -- but at that point reset has zeroed it
|
||||
# for dying envs. Compromise: at pre-step iter t+1, read
|
||||
# sensor history (last decimation_count physics steps from
|
||||
# iter t) for the envs that just died -- still works for envs
|
||||
# whose sensor history isn't fully cleared by reset (depends
|
||||
# on IsaacLab version). Velocity, on the other hand, must
|
||||
# be sampled BEFORE the impact step; using buffer[head-1]
|
||||
# (the prior iteration's measurement) gives us exactly that.
|
||||
# -------------------------------------------------------------
|
||||
v_b = robot.data.root_lin_vel_b # (N, 3) -- state at start of iter t
|
||||
v_fwd = v_b[:, 0]
|
||||
v_xy = torch.norm(v_b[:, :2], dim=-1)
|
||||
|
||||
net_force_hist = contact_sensor.data.net_forces_w_history # (N, H, B_total, 3)
|
||||
f_mags = torch.norm(net_force_hist[:, :, body_ids, :], dim=-1)
|
||||
peak_force, _ = f_mags.reshape(N, -1).max(dim=-1) # (N,)
|
||||
|
||||
# Save what we *had* in the buffer before overwriting, in case
|
||||
# we need to look at "1 step ago" specifically.
|
||||
prev_head = (head - 1) % K
|
||||
prev_vel = vel_buf[prev_head].clone()
|
||||
prev_speed = speed_buf[prev_head].clone()
|
||||
prev_force = force_buf[prev_head].clone()
|
||||
|
||||
# Push into rolling buffer (newest at index `head`, then advance).
|
||||
vel_buf[head] = v_fwd
|
||||
speed_buf[head] = v_xy
|
||||
force_buf[head] = peak_force
|
||||
new_head = head
|
||||
head = (head + 1) % K
|
||||
|
||||
# ---- step the env (terminations + auto-reset happen in here) ----
|
||||
with torch.inference_mode():
|
||||
actions = policy(obs)
|
||||
obs, _, _, _ = env.step(actions)
|
||||
|
||||
# ---- post-step: read sensor history AGAIN for envs that just died;
|
||||
# this captures the *peak* contact force during the impact step
|
||||
# if the reset hasn't yet zeroed the history buffer.
|
||||
post_force_hist = contact_sensor.data.net_forces_w_history
|
||||
post_f_mags = torch.norm(post_force_hist[:, :, body_ids, :], dim=-1)
|
||||
post_peak_force, _ = post_f_mags.reshape(N, -1).max(dim=-1)
|
||||
|
||||
if step >= args_cli.warmup:
|
||||
n_steps_counted += 1
|
||||
|
||||
done_bc = tm.get_term("base_contact").to(torch.bool)
|
||||
done_to = tm.get_term("time_out").to(torch.bool)
|
||||
done_succ = tm.get_term("early_termination").to(torch.bool)
|
||||
|
||||
n_collisions += int(done_bc.sum().item())
|
||||
n_success += int(done_succ.sum().item())
|
||||
n_timeout += int(done_to.sum().item())
|
||||
|
||||
if done_bc.any():
|
||||
idxs = torch.where(done_bc)[0].cpu().tolist()
|
||||
# Buffer order: oldest -> newest. Newest = just pushed at
|
||||
# `new_head`. For envs that died THIS env.step(), the
|
||||
# `new_head` sample is the pre-impact state at iter start.
|
||||
order = [(head + i) % K for i in range(K)]
|
||||
for i in idxs:
|
||||
vt = vel_buf[order, i].cpu().numpy()
|
||||
st = speed_buf[order, i].cpu().numpy()
|
||||
ft = force_buf[order, i].cpu().numpy()
|
||||
coll_traj_vel.append(vt)
|
||||
coll_traj_speed.append(st)
|
||||
coll_traj_force.append(ft)
|
||||
# Pre-impact velocity = the value we measured BEFORE
|
||||
# this env.step() (which contained the actual impact).
|
||||
coll_impact_vel.append(float(v_fwd[i].item()))
|
||||
coll_impact_speed.append(float(v_xy[i].item()))
|
||||
# Impact force: take max of (a) pre-step reading, (b) post-step
|
||||
# reading, (c) buffer max over the window -- whichever is
|
||||
# largest is the best estimate of true peak force.
|
||||
win_peak = float(force_buf[:, i].max().item())
|
||||
impact_f = max(
|
||||
float(peak_force[i].item()),
|
||||
float(post_peak_force[i].item()),
|
||||
win_peak,
|
||||
)
|
||||
coll_impact_force.append(impact_f)
|
||||
|
||||
# Global samples (alive envs only).
|
||||
alive = ~tm.dones.to(torch.bool)
|
||||
if alive.any():
|
||||
all_vel_chunks.append(v_fwd[alive].cpu().numpy())
|
||||
# use post-step force for non-dying envs -- it includes
|
||||
# in-step contacts the policy then recovers from.
|
||||
all_force_chunks.append(post_peak_force[alive].cpu().numpy())
|
||||
|
||||
step += 1
|
||||
|
||||
# ---- save ----
|
||||
out_dir = args_cli.out_dir
|
||||
if out_dir is None:
|
||||
tag = os.path.splitext(os.path.basename(args_cli.checkpoint))[0]
|
||||
out_dir = os.path.join("logs", "collisions", f"{args_cli.task}__{tag}")
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
impact_vel_arr = np.array(coll_impact_vel)
|
||||
impact_speed_arr = np.array(coll_impact_speed)
|
||||
impact_force_arr = np.array(coll_impact_force)
|
||||
traj_vel_arr = np.stack(coll_traj_vel) if coll_traj_vel else np.zeros((0, K))
|
||||
traj_speed_arr = np.stack(coll_traj_speed) if coll_traj_speed else np.zeros((0, K))
|
||||
traj_force_arr = np.stack(coll_traj_force) if coll_traj_force else np.zeros((0, K))
|
||||
all_vel_arr = np.concatenate(all_vel_chunks) if all_vel_chunks else np.zeros((0,))
|
||||
all_force_arr = np.concatenate(all_force_chunks) if all_force_chunks else np.zeros((0,))
|
||||
|
||||
np.savez_compressed(
|
||||
os.path.join(out_dir, "collisions.npz"),
|
||||
impact_vel=impact_vel_arr,
|
||||
impact_speed=impact_speed_arr,
|
||||
impact_force=impact_force_arr,
|
||||
traj_vel=traj_vel_arr,
|
||||
traj_speed=traj_speed_arr,
|
||||
traj_force=traj_force_arr,
|
||||
all_vel=all_vel_arr,
|
||||
all_force=all_force_arr,
|
||||
threshold=np.array(threshold, dtype=np.float32),
|
||||
ep_term_weight=np.array(ep_term_w, dtype=np.float32),
|
||||
K=np.array(K),
|
||||
num_envs=np.array(N),
|
||||
steps=np.array(n_steps_counted),
|
||||
n_collisions=np.array(n_collisions),
|
||||
n_success=np.array(n_success),
|
||||
n_timeout=np.array(n_timeout),
|
||||
)
|
||||
|
||||
# ---- human summary ----
|
||||
def pct(x, total):
|
||||
return 100.0 * x / max(total, 1)
|
||||
|
||||
total_term = n_collisions + n_success + n_timeout
|
||||
summary = []
|
||||
summary.append(f"task : {args_cli.task}")
|
||||
summary.append(f"checkpoint : {args_cli.checkpoint}")
|
||||
summary.append(f"steps counted : {n_steps_counted} (warmup={args_cli.warmup})")
|
||||
summary.append(f"num envs : {N}")
|
||||
summary.append(f"")
|
||||
summary.append(f"base_contact threshold (training) : {threshold} N")
|
||||
summary.append(f"episode_termination reward weight : {ep_term_w}")
|
||||
summary.append(f"")
|
||||
summary.append(f"terminations : success={n_success} collision={n_collisions} timeout={n_timeout}")
|
||||
summary.append(f" success % : {pct(n_success, total_term):.1f}")
|
||||
summary.append(f" collision % : {pct(n_collisions, total_term):.1f}")
|
||||
summary.append(f" timeout % : {pct(n_timeout, total_term):.1f}")
|
||||
summary.append(f"")
|
||||
if impact_vel_arr.size > 0:
|
||||
summary.append(f"AT IMPACT (n={impact_vel_arr.size})")
|
||||
summary.append(f" v_fwd (body x) : mean={impact_vel_arr.mean():+.3f} median={np.median(impact_vel_arr):+.3f} "
|
||||
f"p10={np.percentile(impact_vel_arr,10):+.3f} p90={np.percentile(impact_vel_arr,90):+.3f} "
|
||||
f"std={impact_vel_arr.std():.3f} (m/s)")
|
||||
summary.append(f" |v_xy| : mean={impact_speed_arr.mean():.3f} median={np.median(impact_speed_arr):.3f} "
|
||||
f"p10={np.percentile(impact_speed_arr,10):.3f} p90={np.percentile(impact_speed_arr,90):.3f} (m/s)")
|
||||
summary.append(f" peak force : mean={impact_force_arr.mean():.1f} median={np.median(impact_force_arr):.1f} "
|
||||
f"p10={np.percentile(impact_force_arr,10):.1f} p90={np.percentile(impact_force_arr,90):.1f} "
|
||||
f"max={impact_force_arr.max():.1f} (N)")
|
||||
summary.append(f" force / threshold ratio : median {np.median(impact_force_arr)/threshold:.1f}x, "
|
||||
f"max {impact_force_arr.max()/threshold:.1f}x")
|
||||
n_slow = int((impact_speed_arr < 0.2).sum())
|
||||
n_fast = int((impact_speed_arr > 0.8).sum())
|
||||
summary.append(f" slow impacts (|v_xy|<0.2 m/s) : {n_slow} ({pct(n_slow, impact_speed_arr.size):.1f}%)")
|
||||
summary.append(f" fast impacts (|v_xy|>0.8 m/s) : {n_fast} ({pct(n_fast, impact_speed_arr.size):.1f}%)")
|
||||
else:
|
||||
summary.append("NO COLLISIONS in this run.")
|
||||
if all_vel_arr.size > 0:
|
||||
summary.append(f"")
|
||||
summary.append(f"GLOBAL (alive-step samples, n={all_vel_arr.size})")
|
||||
summary.append(f" v_fwd : mean={all_vel_arr.mean():+.3f} median={np.median(all_vel_arr):+.3f} "
|
||||
f"std={all_vel_arr.std():.3f} (m/s)")
|
||||
summary.append(f" peak F : mean={all_force_arr.mean():.2f} median={np.median(all_force_arr):.2f} "
|
||||
f"p99={np.percentile(all_force_arr,99):.2f} (N)")
|
||||
|
||||
txt = "\n".join(summary)
|
||||
with open(os.path.join(out_dir, "summary.txt"), "w") as f:
|
||||
f.write(txt + "\n")
|
||||
print("\n" + "=" * 78)
|
||||
print(txt)
|
||||
print("=" * 78)
|
||||
print(f"\n[collisions] saved to: {out_dir}")
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
simulation_app.close()
|
||||
@@ -0,0 +1,246 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Convert old policy checkpoints to the new positional encoding structure.
|
||||
|
||||
The old checkpoint uses:
|
||||
attn_*_net.position_embedding.penc.inv_freq
|
||||
|
||||
The new checkpoint uses:
|
||||
attn_*_net.pos_encoding
|
||||
|
||||
This script converts old checkpoints to the new format by:
|
||||
1. Removing the old position_embedding.penc.inv_freq keys
|
||||
2. Computing and adding the new pos_encoding tensors
|
||||
|
||||
Usage:
|
||||
./isaaclab.sh -p scripts/convert_checkpoint.py --input <old_checkpoint.pt> --output <new_checkpoint.pt>
|
||||
|
||||
# Or convert in-place:
|
||||
./isaaclab.sh -p scripts/convert_checkpoint.py --input <checkpoint.pt> --inplace
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _compute_positional_encoding_3d(
|
||||
channels: int, D: int, H: int, W: int, device: torch.device, dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
"""Compute 3D positional encoding for given spatial dimensions.
|
||||
|
||||
This code is adapted from: https://github.com/tatp22/multidim-positional-encoding
|
||||
|
||||
Args:
|
||||
channels: Number of channels for the encoding.
|
||||
D, H, W: Spatial dimensions (depth, height, width).
|
||||
device: Device to create tensor on.
|
||||
dtype: Data type for the encoding tensor.
|
||||
|
||||
Returns:
|
||||
Positional encoding tensor of shape (1, channels, D, H, W).
|
||||
"""
|
||||
org_channels = channels
|
||||
channels = int(math.ceil(channels / 6) * 2)
|
||||
if channels % 2:
|
||||
channels += 1
|
||||
inv_freq = 1.0 / (10000 ** (torch.arange(0, channels, 2, device=device).float() / channels))
|
||||
|
||||
def get_emb(sin_inp: torch.Tensor) -> torch.Tensor:
|
||||
emb = torch.stack((sin_inp.sin(), sin_inp.cos()), dim=-1)
|
||||
return torch.flatten(emb, -2, -1)
|
||||
|
||||
pos_x = torch.arange(D, device=device, dtype=inv_freq.dtype)
|
||||
pos_y = torch.arange(H, device=device, dtype=inv_freq.dtype)
|
||||
pos_z = torch.arange(W, device=device, dtype=inv_freq.dtype)
|
||||
sin_inp_x = torch.einsum("i,j->ij", pos_x, inv_freq)
|
||||
sin_inp_y = torch.einsum("i,j->ij", pos_y, inv_freq)
|
||||
sin_inp_z = torch.einsum("i,j->ij", pos_z, inv_freq)
|
||||
emb_x = get_emb(sin_inp_x).unsqueeze(1).unsqueeze(1)
|
||||
emb_y = get_emb(sin_inp_y).unsqueeze(1)
|
||||
emb_z = get_emb(sin_inp_z)
|
||||
emb = torch.zeros((D, H, W, channels * 3), device=device, dtype=dtype)
|
||||
emb[:, :, :, :channels] = emb_x
|
||||
emb[:, :, :, channels : 2 * channels] = emb_y
|
||||
emb[:, :, :, 2 * channels :] = emb_z
|
||||
|
||||
# Convert from (D, H, W, ch) to (1, ch, D, H, W) format (channel-first with batch dim)
|
||||
enc = emb[None, :, :, :, :org_channels] # (1, D, H, W, org_channels)
|
||||
enc = enc.permute(0, 4, 1, 2, 3) # (1, org_channels, D, H, W)
|
||||
return enc
|
||||
|
||||
|
||||
def convert_checkpoint(
|
||||
input_path: str,
|
||||
output_path: str,
|
||||
image_input_dims: tuple = (64, 5, 8),
|
||||
height_input_dims: tuple = (64, 7, 7),
|
||||
num_cameras: int = 1,
|
||||
force: bool = False,
|
||||
) -> None:
|
||||
"""Convert old checkpoint to new format.
|
||||
|
||||
Args:
|
||||
input_path: Path to the old checkpoint.
|
||||
output_path: Path to save the converted checkpoint.
|
||||
image_input_dims: Tuple of (C, H, W) for image input.
|
||||
height_input_dims: Tuple of (C, H, W) for height map input.
|
||||
num_cameras: Number of cameras (1 or 2).
|
||||
force: Force re-compute pos_encoding even if already in new format.
|
||||
"""
|
||||
print(f"Loading checkpoint from: {input_path}")
|
||||
checkpoint = torch.load(input_path, map_location="cpu", weights_only=False)
|
||||
|
||||
state_dict = checkpoint["model_state_dict"]
|
||||
|
||||
# Find old position_embedding keys to remove
|
||||
old_keys = [k for k in state_dict.keys() if "position_embedding.penc.inv_freq" in k]
|
||||
|
||||
# Check if new keys exist
|
||||
new_keys = [k for k in state_dict.keys() if ".pos_encoding" in k]
|
||||
|
||||
if not old_keys and not force:
|
||||
print("No old position_embedding keys found. Checkpoint may already be in new format.")
|
||||
if new_keys:
|
||||
print(f"Found new pos_encoding keys: {new_keys}")
|
||||
for k in new_keys:
|
||||
print(f" {k}: {state_dict[k].shape}")
|
||||
print("Checkpoint is already in new format. Use --force to re-compute pos_encoding.")
|
||||
return
|
||||
else:
|
||||
print("Warning: Neither old nor new positional encoding keys found!")
|
||||
return
|
||||
|
||||
if old_keys:
|
||||
print(f"Found {len(old_keys)} old position_embedding keys to convert:")
|
||||
for k in old_keys:
|
||||
print(f" - {k}")
|
||||
# Remove old keys
|
||||
for k in old_keys:
|
||||
del state_dict[k]
|
||||
|
||||
if new_keys and force:
|
||||
print(f"Force mode: Replacing {len(new_keys)} existing pos_encoding keys:")
|
||||
for k in new_keys:
|
||||
print(f" - {k}: {state_dict[k].shape}")
|
||||
# Remove existing pos_encoding keys to replace them
|
||||
for k in new_keys:
|
||||
del state_dict[k]
|
||||
|
||||
# Compute and add new positional encodings
|
||||
# For attn_image_net and attn_critic_image_net: spatial_dims = (num_cameras, H, W)
|
||||
# For attn_height_net: spatial_dims = (1, H, W)
|
||||
|
||||
image_channels = image_input_dims[0]
|
||||
height_channels = height_input_dims[0]
|
||||
|
||||
image_spatial_dims = (num_cameras, image_input_dims[1], image_input_dims[2])
|
||||
height_spatial_dims = (1, height_input_dims[1], height_input_dims[2])
|
||||
|
||||
print(f"\nComputing new positional encodings:")
|
||||
print(f" - image_input_dims: {image_input_dims}, num_cameras: {num_cameras}")
|
||||
print(f" - height_input_dims: {height_input_dims}")
|
||||
print(f" - image spatial_dims (D, H, W): {image_spatial_dims}")
|
||||
print(f" - height spatial_dims (D, H, W): {height_spatial_dims}")
|
||||
|
||||
# Compute positional encodings
|
||||
image_pos_enc = _compute_positional_encoding_3d(
|
||||
image_channels, *image_spatial_dims, torch.device("cpu"), torch.float32
|
||||
)
|
||||
height_pos_enc = _compute_positional_encoding_3d(
|
||||
height_channels, *height_spatial_dims, torch.device("cpu"), torch.float32
|
||||
)
|
||||
|
||||
print(f"\nNew positional encoding shapes:")
|
||||
print(f" - attn_image_net.pos_encoding: {image_pos_enc.shape}")
|
||||
print(f" - attn_height_net.pos_encoding: {height_pos_enc.shape}")
|
||||
print(f" - attn_critic_image_net.pos_encoding: {image_pos_enc.shape}")
|
||||
|
||||
# Add new keys
|
||||
state_dict["attn_image_net.pos_encoding"] = image_pos_enc
|
||||
state_dict["attn_height_net.pos_encoding"] = height_pos_enc
|
||||
state_dict["attn_critic_image_net.pos_encoding"] = image_pos_enc.clone()
|
||||
|
||||
# Save converted checkpoint
|
||||
print(f"\nSaving converted checkpoint to: {output_path}")
|
||||
torch.save(checkpoint, output_path)
|
||||
print("Conversion complete!")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert old policy checkpoints to new positional encoding structure."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input", "-i",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to the old checkpoint file."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", "-o",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to save the converted checkpoint. If not specified, uses --inplace."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--inplace",
|
||||
action="store_true",
|
||||
help="Convert the checkpoint in-place (overwrite the input file)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image-dims",
|
||||
type=str,
|
||||
default="64,5,8",
|
||||
help="Image input dimensions as C,H,W (default: 64,5,8)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height-dims",
|
||||
type=str,
|
||||
default="64,7,7",
|
||||
help="Height input dimensions as C,H,W (default: 64,7,7)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-cameras",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of cameras (default: 1)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Force re-compute pos_encoding even if already in new format."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Determine output path
|
||||
if args.output:
|
||||
output_path = args.output
|
||||
elif args.inplace:
|
||||
output_path = args.input
|
||||
else:
|
||||
# Default: add _converted suffix
|
||||
input_path = Path(args.input)
|
||||
output_path = str(input_path.parent / f"{input_path.stem}_converted{input_path.suffix}")
|
||||
|
||||
# Parse dimensions
|
||||
image_dims = tuple(int(x) for x in args.image_dims.split(","))
|
||||
height_dims = tuple(int(x) for x in args.height_dims.split(","))
|
||||
|
||||
if len(image_dims) != 3 or len(height_dims) != 3:
|
||||
raise ValueError("Dimensions must be specified as C,H,W (3 values)")
|
||||
|
||||
convert_checkpoint(
|
||||
input_path=args.input,
|
||||
output_path=output_path,
|
||||
image_input_dims=image_dims,
|
||||
height_input_dims=height_dims,
|
||||
num_cameras=args.num_cameras,
|
||||
force=args.force,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+140
@@ -0,0 +1,140 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Headless evaluator for a *TorchScript JIT* navigation policy.
|
||||
|
||||
Mirrors ``eval_terminations.py`` but instead of building an rsl_rl Runner and
|
||||
loading a ``model_*.pt`` checkpoint, it directly loads a JIT-scripted
|
||||
``ActorCriticSRU`` policy (e.g. the official upstream deployment policy
|
||||
``nav_policy_new.pt``) and rolls it out in the IsaacLab navigation env.
|
||||
|
||||
The JIT module exposes:
|
||||
forward(observations: Tensor, reset: bool=False) -> actions: Tensor
|
||||
with internal ``hidden_state`` / ``cell_state`` buffers sized
|
||||
``(num_layers, 1, hidden_size)``. We resize those buffers to
|
||||
``(num_layers, num_envs, hidden_size)`` after load so a single forward call
|
||||
handles all envs in parallel, and we per-env zero them on episode reset.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
parser = argparse.ArgumentParser(description="Evaluate a JIT navigation policy.")
|
||||
parser.add_argument("--num_envs", type=int, default=32)
|
||||
parser.add_argument("--task", type=str, required=True, help="Gym task id (use a *-Play-v0 variant).")
|
||||
parser.add_argument("--jit_policy", type=str, required=True, help="In-container path to TorchScript .pt.")
|
||||
parser.add_argument("--steps", type=int, default=2000)
|
||||
parser.add_argument("--warmup", type=int, default=20)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli, _ = parser.parse_known_args()
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
import gymnasium as gym
|
||||
import torch
|
||||
|
||||
import isaaclab_tasks # noqa: F401
|
||||
import isaaclab_nav_task # noqa: F401
|
||||
|
||||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||||
from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper
|
||||
|
||||
|
||||
def main():
|
||||
# ---- env setup ----
|
||||
spec = gym.spec(args_cli.task)
|
||||
env_cfg_class = spec.kwargs.get("env_cfg_entry_point")
|
||||
env_cfg: ManagerBasedRLEnvCfg = env_cfg_class()
|
||||
env_cfg.scene.num_envs = args_cli.num_envs
|
||||
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None)
|
||||
env = RslRlVecEnvWrapper(env)
|
||||
device = env.unwrapped.device
|
||||
num_envs = args_cli.num_envs
|
||||
|
||||
# ---- policy load ----
|
||||
print(f"[eval-jit] loading JIT policy: {args_cli.jit_policy}")
|
||||
model = torch.jit.load(args_cli.jit_policy, map_location=device)
|
||||
model.eval()
|
||||
print(f"[eval-jit] model loaded; type={type(model).__name__}")
|
||||
|
||||
# Resize hidden state buffers from (num_layers, 1, hidden_size) to
|
||||
# (num_layers, num_envs, hidden_size). Works because ``rnn`` is a generic
|
||||
# LSTM-style cell that accepts any batch dim.
|
||||
h = model.hidden_state
|
||||
c = model.cell_state
|
||||
print(f"[eval-jit] original hidden_state shape: {tuple(h.shape)}")
|
||||
model.hidden_state = torch.zeros(h.shape[0], num_envs, h.shape[2], device=device, dtype=h.dtype)
|
||||
model.cell_state = torch.zeros(c.shape[0], num_envs, c.shape[2], device=device, dtype=c.dtype)
|
||||
print(f"[eval-jit] resized hidden_state shape: {tuple(model.hidden_state.shape)}")
|
||||
|
||||
# ---- termination accounting ----
|
||||
tm = env.unwrapped.termination_manager
|
||||
term_names = list(tm.active_terms)
|
||||
print(f"[eval-jit] active termination terms: {term_names}")
|
||||
counts = {name: 0 for name in term_names}
|
||||
counts["__any__"] = 0
|
||||
|
||||
# ---- rollout ----
|
||||
obs, _ = env.get_observations()
|
||||
prev_dones = torch.zeros(num_envs, dtype=torch.bool, device=device)
|
||||
step = 0
|
||||
while simulation_app.is_running() and step < args_cli.steps:
|
||||
# Zero hidden state for envs that just terminated (got reset internally
|
||||
# by IsaacLab at the start of this step's env.step()).
|
||||
if prev_dones.any():
|
||||
mask = prev_dones.to(torch.bool)
|
||||
# buffer shape: (num_layers, num_envs, hidden_size)
|
||||
model.hidden_state[:, mask, :] = 0
|
||||
model.cell_state[:, mask, :] = 0
|
||||
|
||||
with torch.inference_mode():
|
||||
actions = model(obs, False)
|
||||
obs, _, dones, _ = env.step(actions)
|
||||
|
||||
prev_dones = dones.to(torch.bool)
|
||||
|
||||
if step >= args_cli.warmup:
|
||||
for name in term_names:
|
||||
counts[name] += int(tm.get_term(name).sum().item())
|
||||
counts["__any__"] += int(env.unwrapped.termination_manager.dones.sum().item())
|
||||
step += 1
|
||||
|
||||
# ---- report ----
|
||||
total = counts["__any__"] if counts["__any__"] > 0 else 1
|
||||
print("\n" + "=" * 78)
|
||||
print(f"[eval-jit] steps counted = {args_cli.steps - args_cli.warmup}, num_envs = {num_envs}")
|
||||
print(f"[eval-jit] total terminations observed = {counts['__any__']}")
|
||||
print("\n TERM COUNT %-of-terms")
|
||||
print(" " + "-" * 50)
|
||||
for name in term_names:
|
||||
c_n = counts[name]
|
||||
pct = 100.0 * c_n / total
|
||||
marker = ""
|
||||
if name == "early_termination":
|
||||
marker = " <- success"
|
||||
elif name in ("base_contact", "large_pitch_angle"):
|
||||
marker = " <- collision/tipover"
|
||||
elif name == "terrain_fall":
|
||||
marker = " <- PIT FALL"
|
||||
elif name == "time_out":
|
||||
marker = " <- ran out of time"
|
||||
print(f" {name:<28} {c_n:>6} {pct:>5.1f}%{marker}")
|
||||
print("=" * 78)
|
||||
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
simulation_app.close()
|
||||
@@ -0,0 +1,138 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Headless termination-breakdown evaluator for a trained navigation policy.
|
||||
|
||||
Question this answers: *for a given checkpoint on a given task/difficulty,
|
||||
what fraction of episodes end via each termination reason* -- including the
|
||||
re-thresholded ``terrain_fall`` (now -0.3m for Go2) which is the only way to
|
||||
distinguish *fell into a pit* from *bumped a wall*.
|
||||
|
||||
The script mirrors ``play.py`` for env/policy setup, then counts the
|
||||
TerminationManager's per-term done flags across a fixed number of env steps.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
parser = argparse.ArgumentParser(description="Evaluate termination breakdown for a navigation policy.")
|
||||
parser.add_argument("--num_envs", type=int, default=64, help="Number of parallel environments.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Gym task id (use a *-Play-v0 variant).")
|
||||
parser.add_argument("--checkpoint", type=str, default=None, help="In-container path to model_*.pt.")
|
||||
parser.add_argument("--steps", type=int, default=2000, help="Number of env.step() calls.")
|
||||
parser.add_argument("--warmup", type=int, default=20, help="Steps to skip before counting (let resets settle).")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Seed.")
|
||||
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli, _ = parser.parse_known_args()
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
import gymnasium as gym
|
||||
import torch
|
||||
|
||||
from rsl_rl.runners import OnPolicyRunner
|
||||
|
||||
import isaaclab_tasks # noqa: F401
|
||||
import isaaclab_nav_task # noqa: F401
|
||||
|
||||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
|
||||
|
||||
|
||||
def main():
|
||||
spec = gym.spec(args_cli.task)
|
||||
env_cfg_class = spec.kwargs.get("env_cfg_entry_point")
|
||||
agent_cfg_class = spec.kwargs.get("rsl_rl_cfg_entry_point")
|
||||
|
||||
env_cfg: ManagerBasedRLEnvCfg = env_cfg_class()
|
||||
agent_cfg: RslRlOnPolicyRunnerCfg = agent_cfg_class()
|
||||
env_cfg.scene.num_envs = args_cli.num_envs
|
||||
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None)
|
||||
env = RslRlVecEnvWrapper(env)
|
||||
|
||||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
|
||||
|
||||
if not args_cli.checkpoint:
|
||||
raise SystemExit("[eval] --checkpoint is required.")
|
||||
print(f"[eval] loading checkpoint: {args_cli.checkpoint}")
|
||||
loaded = torch.load(args_cli.checkpoint, map_location="cpu", weights_only=False)
|
||||
runner.alg.actor_critic.load_state_dict(loaded["model_state_dict"], strict=True)
|
||||
if runner.empirical_normalization:
|
||||
runner.obs_normalizer.load_state_dict(loaded["obs_norm_state_dict"])
|
||||
runner.critic_obs_normalizer.load_state_dict(loaded["critic_obs_norm_state_dict"])
|
||||
print(f"[eval] checkpoint iter = {loaded.get('iter', '?')}")
|
||||
|
||||
policy = runner.get_inference_policy(device=env.unwrapped.device)
|
||||
|
||||
# Termination manager exposes per-term done flags as torch tensors.
|
||||
tm = env.unwrapped.termination_manager
|
||||
term_names = list(tm.active_terms)
|
||||
print(f"[eval] active termination terms: {term_names}")
|
||||
|
||||
# Per-term running counters (number of envs that ended via this reason).
|
||||
counts = {name: 0 for name in term_names}
|
||||
counts["__any__"] = 0 # any termination (sanity)
|
||||
|
||||
# Print the relevant cfg knobs so the run is self-documenting.
|
||||
try:
|
||||
tf_params = env_cfg.terminations.terrain_fall.params
|
||||
print(f"[eval] terrain_fall.fall_height_threshold = "
|
||||
f"{tf_params.get('fall_height_threshold', 'default')}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
obs, _ = env.get_observations()
|
||||
step = 0
|
||||
while simulation_app.is_running() and step < args_cli.steps:
|
||||
with torch.inference_mode():
|
||||
actions = policy(obs)
|
||||
obs, _, _, _ = env.step(actions)
|
||||
|
||||
if step >= args_cli.warmup:
|
||||
# Each per-term tensor is bool[num_envs]; sum gives this-step count.
|
||||
for name in term_names:
|
||||
term_tensor = tm.get_term(name)
|
||||
counts[name] += int(term_tensor.sum().item())
|
||||
# Any-termination flag (rsl_rl wrapper resets on time_out OR done).
|
||||
any_done = env.unwrapped.termination_manager.dones
|
||||
counts["__any__"] += int(any_done.sum().item())
|
||||
step += 1
|
||||
|
||||
# -------- summarise --------
|
||||
total = counts["__any__"] if counts["__any__"] > 0 else 1
|
||||
print("\n" + "=" * 78)
|
||||
print(f"[eval] steps counted = {args_cli.steps - args_cli.warmup}, num_envs = {args_cli.num_envs}")
|
||||
print(f"[eval] total terminations observed = {counts['__any__']}")
|
||||
print("\n TERM COUNT %-of-terms")
|
||||
print(" " + "-" * 50)
|
||||
for name in term_names:
|
||||
c = counts[name]
|
||||
pct = 100.0 * c / total
|
||||
marker = ""
|
||||
if name == "early_termination":
|
||||
marker = " <- success"
|
||||
elif name in ("base_contact", "large_pitch_angle"):
|
||||
marker = " <- collision/tipover"
|
||||
elif name == "terrain_fall":
|
||||
marker = " <- PIT FALL (with new threshold)"
|
||||
elif name == "time_out":
|
||||
marker = " <- ran out of time"
|
||||
print(f" {name:<28} {c:>6} {pct:>5.1f}%{marker}")
|
||||
print("=" * 78)
|
||||
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
simulation_app.close()
|
||||
@@ -0,0 +1,277 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Play a trained navigation policy (PPO/MDPO) with automatic checkpoint loading.
|
||||
|
||||
Usage:
|
||||
python scripts/play.py --task <task_name> [options]
|
||||
|
||||
Arguments:
|
||||
--task Task name (required, typically *-Play-v0 variant)
|
||||
--checkpoint Path to model checkpoint (.pt file)
|
||||
--use_last_checkpoint Use latest checkpoint from logs (default behavior)
|
||||
--num_envs Number of parallel environments
|
||||
--video Enable video recording
|
||||
--video_length Video length in steps (default: 200)
|
||||
|
||||
Examples:
|
||||
python scripts/play.py --task Isaac-Navigation-B2W-Play-v0
|
||||
python scripts/play.py --task Isaac-Navigation-B2W-Play-v0 --checkpoint path/to/model.pt
|
||||
python scripts/play.py --task Isaac-Navigation-B2W-Play-v0 --video --num_envs 16
|
||||
|
||||
Note: Automatically finds latest checkpoint if --checkpoint not specified.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
# Add argparse arguments
|
||||
parser = argparse.ArgumentParser(description="Play a trained navigation policy with RSL-RL.")
|
||||
parser.add_argument("--video", action="store_true", default=False, help="Record videos during play.")
|
||||
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
|
||||
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
|
||||
parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint.")
|
||||
parser.add_argument("--use_last_checkpoint", action="store_true", help="Use last checkpoint from logs.")
|
||||
parser.add_argument("--export_jit", action="store_true", default=False, help="Export policy as JIT module.")
|
||||
parser.add_argument("--export_onnx", action="store_true", default=False, help="Export policy as ONNX model.")
|
||||
|
||||
# Append AppLauncher cli args
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli, hydra_args = parser.parse_known_args()
|
||||
|
||||
# Always enable cameras
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
# Launch simulation
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
# Import after launching simulation
|
||||
import gymnasium as gym
|
||||
import os
|
||||
import re
|
||||
import torch
|
||||
|
||||
from rsl_rl.runners import OnPolicyRunner
|
||||
|
||||
# Import Isaac Lab extensions
|
||||
import isaaclab_tasks # noqa: F401
|
||||
import isaaclab_nav_task # noqa: F401
|
||||
|
||||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper, export_policy_as_onnx
|
||||
|
||||
|
||||
def find_latest_checkpoint(log_path: str, checkpoint_pattern: str = "model_.*.pt") -> str:
|
||||
"""Find the latest checkpoint file in the log directory.
|
||||
|
||||
Args:
|
||||
log_path: Base log directory path
|
||||
checkpoint_pattern: Regex pattern for checkpoint files
|
||||
|
||||
Returns:
|
||||
Path to the latest checkpoint file
|
||||
"""
|
||||
# Find all run directories
|
||||
if not os.path.exists(log_path):
|
||||
raise ValueError(f"Log path does not exist: {log_path}")
|
||||
|
||||
run_dirs = []
|
||||
for entry in os.scandir(log_path):
|
||||
if entry.is_dir() and re.match(r"\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}", entry.name):
|
||||
run_dirs.append(entry.name)
|
||||
|
||||
if not run_dirs:
|
||||
raise ValueError(f"No run directories found in: {log_path}")
|
||||
|
||||
# Sort to get latest run
|
||||
run_dirs.sort()
|
||||
latest_run = run_dirs[-1]
|
||||
run_path = os.path.join(log_path, latest_run)
|
||||
|
||||
# Find checkpoint files
|
||||
checkpoint_files = []
|
||||
for f in os.listdir(run_path):
|
||||
if re.match(checkpoint_pattern, f):
|
||||
checkpoint_files.append(f)
|
||||
|
||||
if not checkpoint_files:
|
||||
raise ValueError(f"No checkpoint files matching '{checkpoint_pattern}' found in: {run_path}")
|
||||
|
||||
# Sort to get latest checkpoint
|
||||
checkpoint_files.sort(key=lambda m: f"{m:0>15}")
|
||||
latest_checkpoint = checkpoint_files[-1]
|
||||
|
||||
return os.path.join(run_path, latest_checkpoint)
|
||||
|
||||
|
||||
def load_checkpoint_with_fallback(runner: OnPolicyRunner, checkpoint_path: str, load_optimizer: bool = True):
|
||||
"""Load checkpoint with fallback for PyTorch compatibility issues.
|
||||
|
||||
Args:
|
||||
runner: RSL-RL runner instance
|
||||
checkpoint_path: Path to checkpoint file
|
||||
load_optimizer: Whether to load optimizer state
|
||||
"""
|
||||
print(f"[INFO] Loading checkpoint from: {checkpoint_path}")
|
||||
|
||||
# Load checkpoint to CPU first for compatibility
|
||||
loaded_dict = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
|
||||
|
||||
# Load model state - handle both standard algorithms (PPO) and MDPO
|
||||
if runner.is_mdpo:
|
||||
# MDPO uses two actor-critics, load same state into both
|
||||
runner.alg.actor_critic_1.load_state_dict(loaded_dict["model_state_dict"], strict=True)
|
||||
runner.alg.actor_critic_2.load_state_dict(loaded_dict["model_state_dict"], strict=True)
|
||||
else:
|
||||
# Standard algorithms use one actor-critic
|
||||
runner.alg.actor_critic.load_state_dict(loaded_dict["model_state_dict"], strict=True)
|
||||
|
||||
# Load normalizers if using empirical normalization
|
||||
if runner.empirical_normalization:
|
||||
runner.obs_normalizer.load_state_dict(loaded_dict["obs_norm_state_dict"])
|
||||
runner.critic_obs_normalizer.load_state_dict(loaded_dict["critic_obs_norm_state_dict"])
|
||||
|
||||
# Load optimizer if requested
|
||||
if load_optimizer:
|
||||
if runner.is_mdpo:
|
||||
runner.alg.optimizer_1.load_state_dict(loaded_dict["optimizer_state_dict"])
|
||||
else:
|
||||
runner.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"])
|
||||
|
||||
runner.current_learning_iteration = loaded_dict["iter"]
|
||||
print(f"[INFO] Loaded checkpoint from iteration {loaded_dict['iter']}")
|
||||
|
||||
|
||||
def export_policy_jit(runner: OnPolicyRunner, checkpoint_path: str):
|
||||
"""Export policy as JIT module to an 'export' folder next to the checkpoint.
|
||||
|
||||
Args:
|
||||
runner: RSL-RL runner instance with loaded policy
|
||||
checkpoint_path: Path to the checkpoint file (used to determine export location)
|
||||
"""
|
||||
# Determine export directory (create 'export' folder in the same directory as checkpoint)
|
||||
checkpoint_dir = os.path.dirname(checkpoint_path)
|
||||
export_dir = os.path.join(checkpoint_dir, "export")
|
||||
|
||||
# Get the actor-critic module
|
||||
if runner.is_mdpo:
|
||||
actor_critic = runner.alg.actor_critic_1
|
||||
else:
|
||||
actor_critic = runner.alg.actor_critic
|
||||
|
||||
# Get normalizer if using empirical normalization
|
||||
normalizer = runner.obs_normalizer if runner.empirical_normalization else None
|
||||
|
||||
# Export using the module's export_jit method
|
||||
print(f"[INFO] Exporting JIT policy to: {export_dir}")
|
||||
actor_critic.export_jit(path=export_dir, filename="policy.pt", normalizer=normalizer)
|
||||
print(f"[INFO] JIT export complete!")
|
||||
|
||||
|
||||
def export_policy_onnx(runner: OnPolicyRunner, checkpoint_path: str):
|
||||
"""Export policy as ONNX model to an 'export' folder next to the checkpoint.
|
||||
|
||||
Args:
|
||||
runner: RSL-RL runner instance with loaded policy
|
||||
checkpoint_path: Path to the checkpoint file (used to determine export location)
|
||||
"""
|
||||
# Determine export directory (create 'export' folder in the same directory as checkpoint)
|
||||
checkpoint_dir = os.path.dirname(checkpoint_path)
|
||||
export_dir = os.path.join(checkpoint_dir, "export")
|
||||
|
||||
# Get the actor-critic module
|
||||
if runner.is_mdpo:
|
||||
actor_critic = runner.alg.actor_critic_1
|
||||
else:
|
||||
actor_critic = runner.alg.actor_critic
|
||||
|
||||
# Get normalizer if using empirical normalization
|
||||
normalizer = runner.obs_normalizer if runner.empirical_normalization else None
|
||||
|
||||
# Check if the module has export_onnx method
|
||||
if not hasattr(actor_critic, "export_onnx"):
|
||||
raise NotImplementedError(
|
||||
f"ONNX export not implemented for {type(actor_critic).__name__}. "
|
||||
"Please add an export_onnx method to this module."
|
||||
)
|
||||
|
||||
# Export using the module's export_onnx method
|
||||
print(f"[INFO] Exporting ONNX policy to: {export_dir}")
|
||||
actor_critic.export_onnx(path=export_dir, filename="policy.onnx", normalizer=normalizer)
|
||||
print(f"[INFO] ONNX export complete!")
|
||||
|
||||
|
||||
def main():
|
||||
"""Play navigation policy with RSL-RL."""
|
||||
# Parse command-line arguments
|
||||
spec = gym.spec(args_cli.task)
|
||||
env_cfg_class = spec.kwargs.get("env_cfg_entry_point")
|
||||
agent_cfg_class = spec.kwargs.get("rsl_rl_cfg_entry_point")
|
||||
|
||||
# Instantiate the configs
|
||||
env_cfg: ManagerBasedRLEnvCfg = env_cfg_class()
|
||||
agent_cfg: RslRlOnPolicyRunnerCfg = agent_cfg_class()
|
||||
|
||||
# Override config from command line
|
||||
if args_cli.num_envs is not None:
|
||||
env_cfg.scene.num_envs = args_cli.num_envs
|
||||
|
||||
# Create the environment
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
|
||||
# Wrap the environment
|
||||
env = RslRlVecEnvWrapper(env)
|
||||
|
||||
# Get checkpoint path
|
||||
if args_cli.checkpoint:
|
||||
resume_path = args_cli.checkpoint
|
||||
else:
|
||||
# Get last checkpoint from log directory
|
||||
log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
|
||||
resume_path = find_latest_checkpoint(log_root_path, checkpoint_pattern="model_.*.pt")
|
||||
|
||||
# Create runner
|
||||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
|
||||
|
||||
# Load checkpoint with compatibility handling
|
||||
load_checkpoint_with_fallback(runner, resume_path)
|
||||
|
||||
# Export JIT if requested
|
||||
if args_cli.export_jit:
|
||||
export_policy_jit(runner, resume_path)
|
||||
|
||||
# Export ONNX if requested
|
||||
if args_cli.export_onnx:
|
||||
export_policy_onnx(runner, resume_path)
|
||||
|
||||
# Obtain policy for inference
|
||||
policy = runner.get_inference_policy(device=env.unwrapped.device)
|
||||
|
||||
# Reset environment
|
||||
obs, _ = env.get_observations()
|
||||
|
||||
# Simulate environment
|
||||
while simulation_app.is_running():
|
||||
# Run policy
|
||||
with torch.inference_mode():
|
||||
actions = policy(obs)
|
||||
# Step environment
|
||||
obs, _, _, _ = env.step(actions)
|
||||
|
||||
# Close the environment
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the main function
|
||||
main()
|
||||
# Close simulation
|
||||
simulation_app.close()
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Render figures from a collisions.npz produced by analyze_collisions.py.
|
||||
|
||||
Run on the *host* (matplotlib only, no Isaac). Writes PNGs next to the .npz.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse, os
|
||||
import numpy as np
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("npz", help="Path to collisions.npz")
|
||||
args = ap.parse_args()
|
||||
|
||||
d = np.load(args.npz)
|
||||
out_dir = os.path.dirname(os.path.abspath(args.npz))
|
||||
|
||||
thr = float(d["threshold"])
|
||||
ep_w = float(d["ep_term_weight"])
|
||||
K = int(d["K"])
|
||||
n_coll = int(d["n_collisions"])
|
||||
n_succ = int(d["n_success"])
|
||||
n_to = int(d["n_timeout"])
|
||||
|
||||
impact_vel = d["impact_vel"]
|
||||
impact_speed = d["impact_speed"]
|
||||
impact_force = d["impact_force"]
|
||||
traj_vel = d["traj_vel"]
|
||||
traj_speed = d["traj_speed"]
|
||||
traj_force = d["traj_force"]
|
||||
all_vel = d["all_vel"]
|
||||
all_force = d["all_force"]
|
||||
|
||||
# ===== Fig 1: impact velocity histogram =====
|
||||
fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
|
||||
|
||||
ax = axes[0]
|
||||
if impact_vel.size > 0:
|
||||
ax.hist(impact_vel, bins=40, color="#d9534f", alpha=0.85, edgecolor="white")
|
||||
ax.axvline(impact_vel.mean(), color="black", ls="--", lw=1.5,
|
||||
label=f"mean = {impact_vel.mean():+.2f} m/s")
|
||||
ax.axvline(0.0, color="gray", ls=":", lw=1)
|
||||
ax.set_xlabel("body-frame forward velocity at impact (m/s)")
|
||||
ax.set_ylabel("# collisions")
|
||||
ax.set_title(f"Velocity at collision (n={impact_vel.size})")
|
||||
ax.legend()
|
||||
ax.grid(alpha=0.3)
|
||||
|
||||
ax = axes[1]
|
||||
if impact_speed.size > 0:
|
||||
ax.hist(impact_speed, bins=40, color="#5bc0de", alpha=0.85, edgecolor="white")
|
||||
ax.axvline(impact_speed.mean(), color="black", ls="--", lw=1.5,
|
||||
label=f"mean = {impact_speed.mean():.2f} m/s")
|
||||
ax.axvline(0.2, color="orange", ls=":", lw=1, label="0.2 m/s (slow-impact cutoff)")
|
||||
ax.set_xlabel("planar speed |v_xy| at impact (m/s)")
|
||||
ax.set_ylabel("# collisions")
|
||||
ax.set_title("Planar speed at collision")
|
||||
ax.legend()
|
||||
ax.grid(alpha=0.3)
|
||||
fig.tight_layout()
|
||||
fig.savefig(os.path.join(out_dir, "fig1_impact_velocity.png"), dpi=130)
|
||||
plt.close(fig)
|
||||
|
||||
# ===== Fig 2: contact force histogram (log scale) with threshold =====
|
||||
fig, ax = plt.subplots(figsize=(8, 5))
|
||||
if impact_force.size > 0:
|
||||
bins = np.logspace(np.log10(max(thr * 0.5, 0.5)),
|
||||
np.log10(max(impact_force.max() * 1.2, thr * 10)), 50)
|
||||
ax.hist(impact_force, bins=bins, color="#f0ad4e", alpha=0.85, edgecolor="white",
|
||||
label=f"impact peak force (n={impact_force.size})")
|
||||
ax.axvline(thr, color="red", ls="--", lw=2, label=f"training threshold = {thr:.1f} N")
|
||||
ax.axvline(np.median(impact_force), color="black", ls=":", lw=1.5,
|
||||
label=f"median impact = {np.median(impact_force):.0f} N "
|
||||
f"({np.median(impact_force)/thr:.0f}× thr)")
|
||||
ax.set_xscale("log")
|
||||
ax.set_xlabel("peak contact force on (base, hip, thigh) at impact (N, log)")
|
||||
ax.set_ylabel("# collisions")
|
||||
ax.set_title("Impact force vs training termination threshold")
|
||||
ax.legend()
|
||||
ax.grid(alpha=0.3, which="both")
|
||||
fig.tight_layout()
|
||||
fig.savefig(os.path.join(out_dir, "fig2_impact_force.png"), dpi=130)
|
||||
plt.close(fig)
|
||||
|
||||
# ===== Fig 3: pre-impact trajectory (median + IQR) =====
|
||||
if traj_vel.shape[0] > 0:
|
||||
steps = np.arange(-K + 1, 1)
|
||||
fig, axes = plt.subplots(1, 2, figsize=(12, 4.5), sharex=True)
|
||||
|
||||
for ax, data, ylabel, title, color in [
|
||||
(axes[0], traj_speed, "planar speed |v_xy| (m/s)", "Speed leading up to collision", "#5bc0de"),
|
||||
(axes[1], traj_force, "peak contact force (N)", "Force leading up to collision", "#f0ad4e"),
|
||||
]:
|
||||
med = np.median(data, axis=0)
|
||||
p25 = np.percentile(data, 25, axis=0)
|
||||
p75 = np.percentile(data, 75, axis=0)
|
||||
ax.fill_between(steps, p25, p75, color=color, alpha=0.35, label="IQR (25-75%)")
|
||||
ax.plot(steps, med, color=color, lw=2.2, label="median")
|
||||
ax.axvline(0, color="red", ls="--", lw=1.5, label="impact step")
|
||||
ax.set_xlabel("env-step relative to impact")
|
||||
ax.set_ylabel(ylabel)
|
||||
ax.set_title(f"{title} (n={data.shape[0]})")
|
||||
ax.legend()
|
||||
ax.grid(alpha=0.3)
|
||||
if axes[1].get_ylim()[1] > 5 * thr:
|
||||
axes[1].set_yscale("symlog", linthresh=max(thr, 1.0))
|
||||
axes[1].axhline(thr, color="red", ls=":", lw=1)
|
||||
fig.tight_layout()
|
||||
fig.savefig(os.path.join(out_dir, "fig3_pre_impact.png"), dpi=130)
|
||||
plt.close(fig)
|
||||
|
||||
# ===== Fig 4: global v_fwd distribution (alive steps) vs at-impact =====
|
||||
fig, ax = plt.subplots(figsize=(8, 5))
|
||||
if all_vel.size > 0:
|
||||
ax.hist(all_vel, bins=60, density=True, color="#5cb85c", alpha=0.55,
|
||||
edgecolor="white", label=f"alive steps (n={all_vel.size})")
|
||||
if impact_vel.size > 0:
|
||||
ax.hist(impact_vel, bins=40, density=True, color="#d9534f", alpha=0.55,
|
||||
edgecolor="white", label=f"at impact (n={impact_vel.size})")
|
||||
ax.axvline(0, color="gray", ls=":", lw=1)
|
||||
ax.set_xlabel("body-frame forward velocity v_fwd (m/s)")
|
||||
ax.set_ylabel("density")
|
||||
ax.set_title("Forward velocity: normal navigation vs at collision")
|
||||
ax.legend()
|
||||
ax.grid(alpha=0.3)
|
||||
fig.tight_layout()
|
||||
fig.savefig(os.path.join(out_dir, "fig4_vfwd_normal_vs_impact.png"), dpi=130)
|
||||
plt.close(fig)
|
||||
|
||||
# ===== short text summary =====
|
||||
lines = []
|
||||
lines.append(f"threshold = {thr} N ep_term_w = {ep_w}")
|
||||
lines.append(f"terminations: success={n_succ} collision={n_coll} timeout={n_to}")
|
||||
if impact_force.size > 0:
|
||||
lines.append(f"impact force: median {np.median(impact_force):.1f} N "
|
||||
f"({np.median(impact_force)/thr:.0f}x threshold), "
|
||||
f"p90 {np.percentile(impact_force,90):.1f}, max {impact_force.max():.1f}")
|
||||
if impact_speed.size > 0:
|
||||
slow = (impact_speed < 0.2).mean() * 100
|
||||
fast = (impact_speed > 0.8).mean() * 100
|
||||
lines.append(f"impact |v_xy|: median {np.median(impact_speed):.2f} m/s "
|
||||
f"slow<0.2: {slow:.0f}% fast>0.8: {fast:.0f}%")
|
||||
print("\n".join(lines))
|
||||
print(f"\nfigures saved into: {out_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,150 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Train a navigation policy using RSL-RL (PPO/MDPO algorithms).
|
||||
|
||||
Usage:
|
||||
python scripts/train.py --task <task_name> --num_envs <num> [options]
|
||||
|
||||
Arguments:
|
||||
--task Task name (required)
|
||||
--num_envs Number of parallel environments
|
||||
--seed Random seed
|
||||
--max_iterations Training iterations
|
||||
--run_name Custom run name for logging
|
||||
--video Enable video recording
|
||||
--video_length Video length in steps (default: 200)
|
||||
--video_interval Recording interval in steps (default: 2000)
|
||||
|
||||
Examples:
|
||||
python scripts/train.py --task Isaac-Navigation-B2W-v0 --num_envs 2048
|
||||
python scripts/train.py --task Isaac-Navigation-B2W-v0 --video --seed 42
|
||||
|
||||
Logs saved to: logs/rsl_rl/<experiment_name>/<timestamp>/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
# Add the parent directory to the path so we can import from the extension
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
# Add argparse arguments
|
||||
parser = argparse.ArgumentParser(description="Train a navigation policy with RSL-RL.")
|
||||
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
|
||||
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
|
||||
parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
|
||||
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
|
||||
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
|
||||
parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.")
|
||||
parser.add_argument("--run_name", type=str, default=None, help="Name of the wandb run (appended to log directory).")
|
||||
parser.add_argument("--checkpoint", type=str, default=None, help="Path to a checkpoint (.pt) to warm-start from before training begins.")
|
||||
parser.add_argument("--load_optimizer", action="store_true", default=False, help="Also load the optimizer state from the checkpoint (default: only load model weights, recommended for cross-robot warm-start).")
|
||||
|
||||
# Append AppLauncher cli args
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli, hydra_args = parser.parse_known_args()
|
||||
|
||||
# always enable cameras to record video
|
||||
if args_cli.video:
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
# Launch simulation
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
# Import after launching simulation
|
||||
import gymnasium as gym
|
||||
import os
|
||||
import torch
|
||||
from datetime import datetime
|
||||
|
||||
from rsl_rl.runners import OnPolicyRunner
|
||||
|
||||
# Import Isaac Lab extensions
|
||||
import isaaclab_tasks # noqa: F401
|
||||
import isaaclab_nav_task # noqa: F401
|
||||
|
||||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||||
from isaaclab.utils.dict import print_dict
|
||||
from isaaclab.utils.io import dump_pickle, dump_yaml
|
||||
from isaaclab_tasks.utils import get_checkpoint_path
|
||||
from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
|
||||
|
||||
# Set torch backends for better performance
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
torch.backends.cudnn.deterministic = False
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def main():
|
||||
"""Train navigation policy with RSL-RL."""
|
||||
# Load the configurations from the registry
|
||||
env_cfg = load_cfg_from_registry(args_cli.task, "env_cfg_entry_point")
|
||||
agent_cfg: RslRlOnPolicyRunnerCfg = load_cfg_from_registry(args_cli.task, "rsl_rl_cfg_entry_point")
|
||||
|
||||
# Override config from command line
|
||||
if args_cli.num_envs is not None:
|
||||
env_cfg.scene.num_envs = args_cli.num_envs
|
||||
if args_cli.seed is not None:
|
||||
agent_cfg.seed = args_cli.seed
|
||||
if args_cli.max_iterations is not None:
|
||||
agent_cfg.max_iterations = args_cli.max_iterations
|
||||
if args_cli.run_name is not None:
|
||||
agent_cfg.run_name = args_cli.run_name
|
||||
|
||||
# Create the environment
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
|
||||
# Wrap the environment
|
||||
env = RslRlVecEnvWrapper(env)
|
||||
|
||||
# Specify log directory
|
||||
log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
|
||||
log_root_path = os.path.abspath(log_root_path)
|
||||
print(f"[INFO] Logging experiment in directory: {log_root_path}")
|
||||
# Specify run directory based on timestamp
|
||||
log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
if agent_cfg.run_name:
|
||||
log_dir += f"_{agent_cfg.run_name}"
|
||||
log_dir = os.path.join(log_root_path, log_dir)
|
||||
|
||||
# Create runner
|
||||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
|
||||
|
||||
# Optional: warm-start from an existing checkpoint (e.g. cross-robot fine-tuning).
|
||||
if args_cli.checkpoint is not None:
|
||||
ckpt_path = os.path.abspath(args_cli.checkpoint)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
|
||||
print(f"[INFO] Warm-starting from checkpoint: {ckpt_path}")
|
||||
print(f"[INFO] load_optimizer = {args_cli.load_optimizer}")
|
||||
runner.load(ckpt_path, load_optimizer=args_cli.load_optimizer)
|
||||
|
||||
# Write git state to log
|
||||
runner.add_git_repo_to_log(__file__)
|
||||
# Save configuration
|
||||
dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg)
|
||||
dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg)
|
||||
dump_pickle(os.path.join(log_dir, "params", "env.pkl"), env_cfg)
|
||||
dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), agent_cfg)
|
||||
|
||||
# Run training
|
||||
runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True)
|
||||
|
||||
# Close the environment
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the main function
|
||||
main()
|
||||
# Close simulation
|
||||
simulation_app.close()
|
||||
@@ -0,0 +1,65 @@
|
||||
# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Installation script for the 'isaaclab_nav_task' python package.
|
||||
|
||||
IsaacLab task extension for SRU (Spatially-enhanced Recurrent Unit) visual navigation project.
|
||||
Provides hierarchical control architecture, maze terrain generation with curriculum learning,
|
||||
and depth-based reinforcement learning for legged robot navigation.
|
||||
"""
|
||||
|
||||
import os
|
||||
import toml
|
||||
|
||||
from setuptools import setup
|
||||
|
||||
# Obtain the extension data from the extension.toml file
|
||||
EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__))
|
||||
# Read the extension.toml file
|
||||
EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml"))
|
||||
|
||||
# Minimum dependencies required prior to installation
|
||||
INSTALL_REQUIRES = [
|
||||
# generic
|
||||
"numpy<2",
|
||||
"torch>=2.5.1",
|
||||
"torchvision>=0.14.1",
|
||||
# io
|
||||
"h5py",
|
||||
# visualization
|
||||
"tensorboard",
|
||||
]
|
||||
|
||||
PYTORCH_INDEX_URL = ["https://download.pytorch.org/whl/cu128"]
|
||||
|
||||
# Installation operation
|
||||
setup(
|
||||
name="isaaclab_nav_task",
|
||||
author="Fan Yang, Per Frivik",
|
||||
author_email="[email protected], [email protected]",
|
||||
maintainer="Fan Yang, Per Frivik",
|
||||
maintainer_email="[email protected], [email protected]",
|
||||
url=EXTENSION_TOML_DATA["package"]["repository"],
|
||||
version=EXTENSION_TOML_DATA["package"]["version"],
|
||||
description=EXTENSION_TOML_DATA["package"]["description"],
|
||||
keywords=EXTENSION_TOML_DATA["package"]["keywords"],
|
||||
include_package_data=True,
|
||||
python_requires=">=3.10",
|
||||
install_requires=INSTALL_REQUIRES,
|
||||
dependency_links=PYTORCH_INDEX_URL,
|
||||
packages=["isaaclab_nav_task"],
|
||||
classifiers=[
|
||||
"Natural Language :: English",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Isaac Sim :: 4.5.0",
|
||||
"Isaac Lab :: 2.1.1",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Intended Audience :: Science/Research",
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
"Topic :: Scientific/Engineering :: Robotics",
|
||||
],
|
||||
zip_safe=False,
|
||||
license="MIT",
|
||||
)
|
||||
Reference in New Issue
Block a user