diff --git a/Train/mount/sru-navigation-sim/LICENSE b/Train/mount/sru-navigation-sim/LICENSE new file mode 100644 index 0000000..3262f7f --- /dev/null +++ b/Train/mount/sru-navigation-sim/LICENSE @@ -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. diff --git a/Train/mount/sru-navigation-sim/config/extension.toml b/Train/mount/sru-navigation-sim/config/extension.toml new file mode 100644 index 0000000..99d7637 --- /dev/null +++ b/Train/mount/sru-navigation-sim/config/extension.toml @@ -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" diff --git a/Train/mount/sru-navigation-sim/docs/README.md b/Train/mount/sru-navigation-sim/docs/README.md new file mode 100644 index 0000000..f93f30e --- /dev/null +++ b/Train/mount/sru-navigation-sim/docs/README.md @@ -0,0 +1,482 @@ +# IsaacLab Navigation Extension - SRU Project + +[![Paper](https://img.shields.io/badge/IJRR-2025-blue)](https://journals.sagepub.com/home/ijr) +[![Website](https://img.shields.io/badge/Project-Website-green)](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 diff --git a/Train/mount/sru-navigation-sim/docs/TERRAIN_AND_GOALS.md b/Train/mount/sru-navigation-sim/docs/TERRAIN_AND_GOALS.md new file mode 100644 index 0000000..e2d0765 --- /dev/null +++ b/Train/mount/sru-navigation-sim/docs/TERRAIN_AND_GOALS.md @@ -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. diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/__init__.py new file mode 100644 index 0000000..c3fd681 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/__init__.py @@ -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) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/__init__.py new file mode 100644 index 0000000..f0bc128 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/__init__.py @@ -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 * diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/__init__.py new file mode 100644 index 0000000..e5e2203 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/__init__.py @@ -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"] diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/aow_d.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/aow_d.py new file mode 100644 index 0000000..04a6c59 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/aow_d.py @@ -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.""" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/b2w.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/b2w.py new file mode 100644 index 0000000..4d2d312 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/b2w.py @@ -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.""" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/instanceable_meshes.usd b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/instanceable_meshes.usd new file mode 100644 index 0000000..67d3639 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/instanceable_meshes.usd differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/Battery_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/Battery_baked.jpg new file mode 100644 index 0000000..1b6c624 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/Battery_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/base_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/base_baked.jpg new file mode 100644 index 0000000..e9500f9 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/base_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/bottom_shell_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/bottom_shell_baked.jpg new file mode 100644 index 0000000..2b84c6a Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/bottom_shell_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/docking_socket_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/docking_socket_baked.jpg new file mode 100644 index 0000000..ac17496 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/docking_socket_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/drive_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/drive_baked.jpg new file mode 100644 index 0000000..64d65d8 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/drive_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/face_shell_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/face_shell_baked.jpg new file mode 100644 index 0000000..9767773 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/face_shell_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/hip_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/hip_baked.jpg new file mode 100644 index 0000000..1b2ea7b Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/hip_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/shank_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/shank_baked.jpg new file mode 100644 index 0000000..43bae90 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/shank_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/thigh_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/thigh_baked.jpg new file mode 100644 index 0000000..d6db0d9 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/thigh_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/top_shell_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/top_shell_baked.jpg new file mode 100644 index 0000000..441ba62 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/top_shell_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/velodyne_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/velodyne_baked.jpg new file mode 100644 index 0000000..e9786c8 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/velodyne_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/velodyne_guard_baked.jpg b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/velodyne_guard_baked.jpg new file mode 100644 index 0000000..569d40e Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/Props/materials/velodyne_guard_baked.jpg differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/aow_d.usd b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/aow_d.usd new file mode 100644 index 0000000..2f14df0 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/aow_d.usd differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/broken_aow_d.usd b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/broken_aow_d.usd new file mode 100644 index 0000000..51ef808 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/broken_aow_d.usd differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/config.yaml b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/config.yaml new file mode 100644 index 0000000..c49b37c --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/config.yaml @@ -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. +## diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/old_working_aow_d.usd b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/old_working_aow_d.usd new file mode 100644 index 0000000..61b89f6 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/AoW-D/old_working_aow_d.usd differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/Props/instanceable_meshes.usd b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/Props/instanceable_meshes.usd new file mode 100644 index 0000000..a328fbf Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/Props/instanceable_meshes.usd differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/Props/test.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/Props/test.py new file mode 100644 index 0000000..e1b4fa8 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/Props/test.py @@ -0,0 +1,5 @@ +# Copyright (c) 2022-2025, Fan Yang and Per Frivik, ETH Zurich. +# All rights reserved. +# +# SPDX-License-Identifier: MIT +print("hi") \ No newline at end of file diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/b2w_rsl.usd b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/b2w_rsl.usd new file mode 100644 index 0000000..89139ec Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/b2w_rsl.usd differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/b2w_wider_rsl.usd b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/b2w_wider_rsl.usd new file mode 100644 index 0000000..89139ec Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/b2w_wider_rsl.usd differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/config.yaml b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/config.yaml new file mode 100644 index 0000000..e0ca4a7 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/config.yaml @@ -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. +## diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/table.usd b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/table.usd new file mode 100644 index 0000000..6d29112 Binary files /dev/null and b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/data/Robots/B2W/table.usd differ diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/go2.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/go2.py new file mode 100644 index 0000000..291255e --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/assets/go2.py @@ -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.""" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/__init__.py new file mode 100644 index 0000000..7fce7a5 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/__init__.py @@ -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 * diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/__init__.py new file mode 100644 index 0000000..90b933c --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/__init__.py @@ -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, + }, +) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/agents/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/agents/__init__.py new file mode 100644 index 0000000..e267b97 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/agents/__init__.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/agents/rsl_rl_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/agents/rsl_rl_cfg.py new file mode 100644 index 0000000..cbeed6a --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/agents/rsl_rl_cfg.py @@ -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" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/navigation_env_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/navigation_env_cfg.py new file mode 100644 index 0000000..5abbfd2 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/aow_d/navigation_env_cfg.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/__init__.py new file mode 100644 index 0000000..39ba82a --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/__init__.py @@ -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, + }, +) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/agents/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/agents/__init__.py new file mode 100644 index 0000000..97c9703 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/agents/__init__.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/agents/rsl_rl_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/agents/rsl_rl_cfg.py new file mode 100644 index 0000000..34b7189 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/agents/rsl_rl_cfg.py @@ -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" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/navigation_env_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/navigation_env_cfg.py new file mode 100644 index 0000000..2222a3c --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/b2w/navigation_env_cfg.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/__init__.py new file mode 100644 index 0000000..7db8aa8 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/__init__.py @@ -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, + }, +) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/agents/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/agents/__init__.py new file mode 100644 index 0000000..7f31efe --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/agents/__init__.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/agents/rsl_rl_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/agents/rsl_rl_cfg.py new file mode 100644 index 0000000..6445cc6 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/agents/rsl_rl_cfg.py @@ -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" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/navigation_env_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/navigation_env_cfg.py new file mode 100644 index 0000000..994d213 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/go2/navigation_env_cfg.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/rl_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/rl_cfg.py new file mode 100644 index 0000000..585ca9b --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/config/rl_cfg.py @@ -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).""" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/__init__.py new file mode 100644 index 0000000..cb3b432 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/__init__.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/curriculums.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/curriculums.py new file mode 100644 index 0000000..6ef4e3f --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/curriculums.py @@ -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)) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/custom_noise.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/custom_noise.py new file mode 100644 index 0000000..f672149 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/custom_noise.py @@ -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.""" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/delay_manager.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/delay_manager.py new file mode 100644 index 0000000..6415d0b --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/delay_manager.py @@ -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) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_noise_encoder.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_noise_encoder.py new file mode 100644 index 0000000..f5859ac --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_noise_encoder.py @@ -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) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_utils/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_utils/__init__.py new file mode 100644 index 0000000..678a229 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_utils/__init__.py @@ -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", +] diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_utils/camera_config.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_utils/camera_config.py new file mode 100644 index 0000000..e343872 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_utils/camera_config.py @@ -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}" + ) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_utils/depth_noise_encoder.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_utils/depth_noise_encoder.py new file mode 100644 index 0000000..7f5a5fd --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/depth_utils/depth_noise_encoder.py @@ -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) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/events.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/events.py new file mode 100644 index 0000000..373a7bd --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/events.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/heightscan_utils/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/heightscan_utils/__init__.py new file mode 100644 index 0000000..d5f0f16 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/heightscan_utils/__init__.py @@ -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"] diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/heightscan_utils/heightscan_encoder.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/heightscan_utils/heightscan_encoder.py new file mode 100644 index 0000000..1d0d581 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/heightscan_utils/heightscan_encoder.py @@ -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()) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/math_utils.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/math_utils.py new file mode 100644 index 0000000..4124c11 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/math_utils.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/__init__.py new file mode 100644 index 0000000..1a258bb --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/__init__.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/actions/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/actions/__init__.py new file mode 100644 index 0000000..947624b --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/actions/__init__.py @@ -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"] \ No newline at end of file diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/actions/navigation_se2_actions.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/actions/navigation_se2_actions.py new file mode 100644 index 0000000..c3a5807 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/actions/navigation_se2_actions.py @@ -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) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/actions/navigation_se2_actions_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/actions/navigation_se2_actions_cfg.py new file mode 100644 index 0000000..dd2d17a --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/actions/navigation_se2_actions_cfg.py @@ -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""" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/goal_commands.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/goal_commands.py new file mode 100644 index 0000000..d0a6a49 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/goal_commands.py @@ -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) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/goal_commands_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/goal_commands_cfg.py new file mode 100644 index 0000000..3c74047 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/navigation/goal_commands_cfg.py @@ -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. + """ + diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/observations.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/observations.py new file mode 100644 index 0000000..e9f62ca --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/observations.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/rewards.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/rewards.py new file mode 100644 index 0000000..5d3c3c8 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/rewards.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/terminations.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/terminations.py new file mode 100644 index 0000000..5e3883f --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/mdp/terminations.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/navigation_env.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/navigation_env.py new file mode 100644 index 0000000..bd778f8 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/navigation_env.py @@ -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() diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/navigation_env_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/navigation_env_cfg.py new file mode 100644 index 0000000..a4e63c6 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/navigation/navigation_env_cfg.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/__init__.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/__init__.py new file mode 100644 index 0000000..1d6d35f --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/__init__.py @@ -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", +] diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/hf_terrains_maze.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/hf_terrains_maze.py new file mode 100644 index 0000000..0363cfc --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/hf_terrains_maze.py @@ -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) diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/hf_terrains_maze_cfg.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/hf_terrains_maze_cfg.py new file mode 100644 index 0000000..43cc1cf --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/hf_terrains_maze_cfg.py @@ -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). + """ + + diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/maze_config.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/maze_config.py new file mode 100644 index 0000000..a4516df --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/maze_config.py @@ -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.""" diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/patches.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/patches.py new file mode 100644 index 0000000..a9ba5c2 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/patches.py @@ -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=: 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 diff --git a/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/terrain_constants.py b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/terrain_constants.py new file mode 100644 index 0000000..65adb74 --- /dev/null +++ b/Train/mount/sru-navigation-sim/isaaclab_nav_task/terrains/terrain_constants.py @@ -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 diff --git a/Train/mount/sru-navigation-sim/pyproject.toml b/Train/mount/sru-navigation-sim/pyproject.toml new file mode 100644 index 0000000..8a0620f --- /dev/null +++ b/Train/mount/sru-navigation-sim/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "toml"] +build-backend = "setuptools.build_meta" diff --git a/Train/mount/sru-navigation-sim/scripts/ablate_depth.py b/Train/mount/sru-navigation-sim/scripts/ablate_depth.py new file mode 100644 index 0000000..9d69735 --- /dev/null +++ b/Train/mount/sru-navigation-sim/scripts/ablate_depth.py @@ -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 --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() diff --git a/Train/mount/sru-navigation-sim/scripts/analyze_collisions.py b/Train/mount/sru-navigation-sim/scripts/analyze_collisions.py new file mode 100755 index 0000000..26f5486 --- /dev/null +++ b/Train/mount/sru-navigation-sim/scripts/analyze_collisions.py @@ -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/__``): + * 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() diff --git a/Train/mount/sru-navigation-sim/scripts/convert_checkpoint.py b/Train/mount/sru-navigation-sim/scripts/convert_checkpoint.py new file mode 100644 index 0000000..31e55d6 --- /dev/null +++ b/Train/mount/sru-navigation-sim/scripts/convert_checkpoint.py @@ -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 --output + + # Or convert in-place: + ./isaaclab.sh -p scripts/convert_checkpoint.py --input --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() diff --git a/Train/mount/sru-navigation-sim/scripts/eval_jit_policy.py b/Train/mount/sru-navigation-sim/scripts/eval_jit_policy.py new file mode 100755 index 0000000..4d2cd90 --- /dev/null +++ b/Train/mount/sru-navigation-sim/scripts/eval_jit_policy.py @@ -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() diff --git a/Train/mount/sru-navigation-sim/scripts/eval_terminations.py b/Train/mount/sru-navigation-sim/scripts/eval_terminations.py new file mode 100644 index 0000000..5134883 --- /dev/null +++ b/Train/mount/sru-navigation-sim/scripts/eval_terminations.py @@ -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() diff --git a/Train/mount/sru-navigation-sim/scripts/play.py b/Train/mount/sru-navigation-sim/scripts/play.py new file mode 100644 index 0000000..c28b680 --- /dev/null +++ b/Train/mount/sru-navigation-sim/scripts/play.py @@ -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 [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() diff --git a/Train/mount/sru-navigation-sim/scripts/plot_collisions.py b/Train/mount/sru-navigation-sim/scripts/plot_collisions.py new file mode 100755 index 0000000..424dfa1 --- /dev/null +++ b/Train/mount/sru-navigation-sim/scripts/plot_collisions.py @@ -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() diff --git a/Train/mount/sru-navigation-sim/scripts/train.py b/Train/mount/sru-navigation-sim/scripts/train.py new file mode 100644 index 0000000..ab28157 --- /dev/null +++ b/Train/mount/sru-navigation-sim/scripts/train.py @@ -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 --num_envs [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/// +""" + +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() diff --git a/Train/mount/sru-navigation-sim/setup.py b/Train/mount/sru-navigation-sim/setup.py new file mode 100644 index 0000000..d92fe9a --- /dev/null +++ b/Train/mount/sru-navigation-sim/setup.py @@ -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="fanyang1@ethz.ch, pfrivik@ethz.ch", + maintainer="Fan Yang, Per Frivik", + maintainer_email="fanyang1@ethz.ch, pfrivik@ethz.ch", + 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", +)