diff --git a/README.md b/README.md index 5e3a021..9531417 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ [![License](https://img.shields.io/github/license/DexForce/EmbodiChain)](LICENSE) [![GitHub Pages](https://img.shields.io/badge/GitHub%20Pages-docs-blue?logo=github&logoColor=white)](https://dexforce.github.io/EmbodiChain/introduction.html) -[![Python](https://img.shields.io/badge/python-3.10-blue.svg)](https://docs.python.org/3/whatsnew/3.10.html) +[![Python](https://img.shields.io/badge/python-3.10%20|%203.11-blue.svg)](https://docs.python.org/3/whatsnew/3.10.html) [![Version](https://img.shields.io/badge/version-0.0.1-blue.svg)](https://github.com/DexForce/EmbodiChain/releases) --- diff --git a/configs/gym/special/simple_task_1.json b/configs/gym/special/simple_task_1.json new file mode 100644 index 0000000..4f10a7e --- /dev/null +++ b/configs/gym/special/simple_task_1.json @@ -0,0 +1,106 @@ +{ + "id": "simple_task_1", + "max_episodes": 30, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 20, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "random_material": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 50, + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": 0.5, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + } + }, + "dataset": { + "save_path": "./", + "extra": { + "scene_type": "commercial", + "task_description": "Sinusoidal motion" + }, + "robot_meta": { + "robot_type": "UR10", + "arm_dofs": 6, + "control_freq": 3, + "observation": { + "vision": { + "cam_high": [] + }, + "states": ["qpos"] + }, + "action": "qpos", + "min_len_steps": 5 + }, + "instruction": { + "lang": "Perform sinusoidal joint motion" + } + } + }, + "robot": { + "uid": "UR10", + "fpath": "UniversalRobots/UR10/UR10.urdf", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [1.57079, -1.57079, 1.57079, -1.57079, -1.57079, -3.14159] + }, + "sensor": [ + { + "sensor_type": "Camera", + "uid": "cam_high", + "width": 640, + "height": 480, + "intrinsics": [488.1665344238281, 488.1665344238281, 320.0, 240.0], + "extrinsics": { + "eye": [1, 0, 3], + "target": [0, 0, 1] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.8, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + ] +} \ No newline at end of file diff --git a/configs/gym/special/simple_task_2.json b/configs/gym/special/simple_task_2.json new file mode 100644 index 0000000..1bc5d7e --- /dev/null +++ b/configs/gym/special/simple_task_2.json @@ -0,0 +1,106 @@ +{ + "id": "simple_task_2", + "max_episodes": 30, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 20, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "random_material": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 50, + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": 0.5, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + } + }, + "dataset": { + "save_path": "./", + "extra": { + "scene_type": "commercial", + "task_description": "Spiral motion" + }, + "robot_meta": { + "robot_type": "UR10", + "arm_dofs": 6, + "control_freq": 3, + "observation": { + "vision": { + "cam_high": [] + }, + "states": ["qpos"] + }, + "action": "qpos", + "min_len_steps": 5 + }, + "instruction": { + "lang": "Perform spiral joint motion" + } + } + }, + "robot": { + "uid": "UR10", + "fpath": "UniversalRobots/UR10/UR10.urdf", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [1.57079, -1.57079, 1.57079, -1.57079, -1.57079, -3.14159] + }, + "sensor": [ + { + "sensor_type": "Camera", + "uid": "cam_high", + "width": 640, + "height": 480, + "intrinsics": [488.1665344238281, 488.1665344238281, 320.0, 240.0], + "extrinsics": { + "eye": [1, 0, 3], + "target": [0, 0, 1] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.8, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + ] +} \ No newline at end of file diff --git a/configs/gym/special/ur10_parallel.json b/configs/gym/special/ur10_parallel.json new file mode 100644 index 0000000..946c0b4 --- /dev/null +++ b/configs/gym/special/ur10_parallel.json @@ -0,0 +1,106 @@ +{ + "id": "ParallelRollout-v1", + "max_episodes": 24, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 20, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "random_material": { + "func": "randomize_visual_material", + "mode": "interval", + "interval_step": 50, + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": 0.5, + "texture_path": "CocoBackground/coco", + "base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]] + } + } + }, + "dataset": { + "save_path": "./", + "extra": { + "scene_type": "commercial", + "task_description": "Simple reach" + }, + "robot_meta": { + "robot_type": "UR10", + "arm_dofs": 6, + "control_freq": 3, + "observation": { + "vision": { + "cam_high": [] + }, + "states": ["qpos"] + }, + "action": "qpos", + "min_len_steps": 5 + }, + "instruction": { + "lang": "Pour water from bottle to cup" + } + } + }, + "robot": { + "uid": "UR10", + "fpath": "UniversalRobots/UR10/UR10.urdf", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [1.57079, -1.57079, 1.57079, -1.57079, -1.57079, -3.14159] + }, + "sensor": [ + { + "sensor_type": "Camera", + "uid": "cam_high", + "width": 640, + "height": 480, + "intrinsics": [488.1665344238281, 488.1665344238281, 320.0, 240.0], + "extrinsics": { + "eye": [1, 0, 3], + "target": [0, 0, 1] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply", + "compute_uv": true + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + ] +} \ No newline at end of file diff --git a/docs/source/_static/workspace_analyzer/joint_workspace.jpg b/docs/source/_static/workspace_analyzer/joint_workspace.jpg new file mode 100644 index 0000000..130afba Binary files /dev/null and b/docs/source/_static/workspace_analyzer/joint_workspace.jpg differ diff --git a/docs/source/_static/workspace_analyzer/plane_workspace1.jpg b/docs/source/_static/workspace_analyzer/plane_workspace1.jpg new file mode 100644 index 0000000..14db5a7 Binary files /dev/null and b/docs/source/_static/workspace_analyzer/plane_workspace1.jpg differ diff --git a/docs/source/_static/workspace_analyzer/plane_workspace2.jpg b/docs/source/_static/workspace_analyzer/plane_workspace2.jpg new file mode 100644 index 0000000..b3165db Binary files /dev/null and b/docs/source/_static/workspace_analyzer/plane_workspace2.jpg differ diff --git a/docs/source/_static/workspace_analyzer/pose_workspace1.jpg b/docs/source/_static/workspace_analyzer/pose_workspace1.jpg new file mode 100644 index 0000000..5d2dc0b Binary files /dev/null and b/docs/source/_static/workspace_analyzer/pose_workspace1.jpg differ diff --git a/docs/source/_static/workspace_analyzer/pose_workspace2.jpg b/docs/source/_static/workspace_analyzer/pose_workspace2.jpg new file mode 100644 index 0000000..ae553a4 Binary files /dev/null and b/docs/source/_static/workspace_analyzer/pose_workspace2.jpg differ diff --git a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.managers.rst b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.managers.rst index afc9df9..4d0f18b 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.managers.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.managers.rst @@ -114,13 +114,19 @@ Randomization .. autosummary:: - rendering + physics + visual spatial - Rendering + Physics + ~~~~~~~~~~~~~~~~~~~~~~~ + .. automodule:: embodichain.lab.gym.envs.managers.randomization.physics + :members: + + Visual ~~~~~~~~~~~~~~~~~~~~~~~ - .. automodule:: embodichain.lab.gym.envs.managers.randomization.rendering + .. automodule:: embodichain.lab.gym.envs.managers.randomization.visual :members: Spatial diff --git a/docs/source/features/workspace_analyzer/caches.md b/docs/source/features/workspace_analyzer/caches.md new file mode 100644 index 0000000..40e70a9 --- /dev/null +++ b/docs/source/features/workspace_analyzer/caches.md @@ -0,0 +1,212 @@ +# Workspace Analyzer Caches + +The Workspace Analyzer Caches module provides simple and efficient caching mechanisms for storing and reusing workspace analysis data. + +## Table of Contents + +- [Overview](#overview) +- [Cache Types](#cache-types) +- [Usage Examples](#usage-examples) +- [Best Practices](#best-practices) +- [API Reference](#api-reference) + +## Overview + +The caches module provides simple caching mechanisms for workspace analysis data to improve performance: + +- **Memory Cache**: Fast in-memory storage for small to medium datasets +- **Disk Cache**: Persistent storage for large datasets or long-term use +- **Simple API**: Easy-to-use interface for both cache types + +### When to Use + +- **Memory Cache**: When you have sufficient RAM and need fast access +- **Disk Cache**: When working with large datasets or need to persist results between sessions + + +## Cache Types + +### Memory Cache + +Fast in-memory caching for medium-sized datasets: + +```python +from embodichain.lab.sim.utility.workspace_analyzer.caches import MemoryCache + +# Create memory cache +cache = MemoryCache() + +# Add data +positions = [...] # Your pose data +cache.add(positions) + +# Get all data +all_data = cache.get_all() +``` + +### Disk Cache + +Persistent disk caching for large datasets: + +```python +from embodichain.lab.sim.utility.workspace_analyzer.caches import DiskCache + +# Create disk cache +cache = DiskCache(save_dir="./my_cache") + +# Add data +positions = [...] # Your pose data +cache.add(positions) +cache.flush() # Save to disk + +# Get data +all_data = cache.get_all() +``` + +### Cache Manager + +Simplified factory pattern for creating caches: + +```python +from embodichain.lab.sim.utility.workspace_analyzer.caches import CacheManager + +# Create memory cache +memory_cache = CacheManager.create_cache("memory") + +# Create disk cache +disk_cache = CacheManager.create_cache("disk", save_dir="./cache") +``` + +## Usage Examples + +### Basic Usage + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.caches import MemoryCache, DiskCache + +# Memory cache example +memory_cache = MemoryCache() + +# Add pose data +poses = [np.eye(4) for _ in range(1000)] # Example data +memory_cache.add(poses) + +# Get data +cached_poses = memory_cache.get_all() +print(f"Cached {len(cached_poses)} poses") + +# Disk cache example +disk_cache = DiskCache(save_dir="./workspace_cache") +disk_cache.add(poses) +disk_cache.flush() # Save to disk + +# Reload data +reloaded_poses = disk_cache.get_all() +``` + +### Real-world Usage + +```python +# Cache usage in workspace analysis +def analyze_workspace_with_cache(robot_config, use_cache=True): + if use_cache: + # Use disk cache to save results + cache = DiskCache(save_dir=f"./cache_{robot_config.name}") + + # Check if cached data exists + if cache.get_batch_count() > 0: + print("Loading data from cache...") + return cache.get_all() + + # Generate new workspace data + print("Generating new workspace data...") + poses = generate_workspace_poses(robot_config) + + if use_cache: + # Save results + cache.add(poses) + cache.flush() + print(f"Cached {len(poses)} poses to disk") + + return poses +``` + +## Best Practices + +### Choosing Cache Type + +- **Small datasets (< 100k poses)**: Use `MemoryCache` +- **Large datasets (> 100k poses)**: Use `DiskCache` +- **Need persistence**: Use `DiskCache` +- **Temporary computation**: Use `MemoryCache` + +### Simple Selection Function + +```python +def choose_cache(data_size, need_persistence=False): + """Choose cache type based on data size and requirements""" + if need_persistence or data_size > 100000: + return DiskCache() + else: + return MemoryCache() + +# Usage example +cache = choose_cache(data_size=50000, need_persistence=True) +``` + + + +## API Reference + +### Basic Methods + +All cache classes support these basic operations: + +```python +# Add data +cache.add(poses_list) + +# Get all data +all_data = cache.get_all() + +# Clear cache +cache.clear() + +# Flush to disk (DiskCache only) +cache.flush() +``` + +### Creating Caches + +```python +# Memory cache +memory_cache = MemoryCache() + +# Disk cache +disk_cache = DiskCache(save_dir="./my_cache") + +# Using factory method +cache = CacheManager.create_cache("memory") # or "disk" +``` + + + +## Quick Start + +```python +# Simplest usage +from embodichain.lab.sim.utility.workspace_analyzer.caches import MemoryCache, DiskCache + +# For small datasets +cache = MemoryCache() + +# For large datasets or persistence needed +cache = DiskCache() + +# Add data +cache.add(your_poses) + +# Get data +result = cache.get_all() +``` diff --git a/docs/source/features/workspace_analyzer/configs.md b/docs/source/features/workspace_analyzer/configs.md new file mode 100644 index 0000000..5888261 --- /dev/null +++ b/docs/source/features/workspace_analyzer/configs.md @@ -0,0 +1,648 @@ +# Workspace Analyzer Configurations + +The **Workspace Analyzer Configs** module provides a comprehensive configuration system for robotic workspace analysis. This module offers type-safe, well-documented configuration classes that control every aspect of workspace analysis, from sampling strategies to visualization options. + +Based on the `WorkspaceAnalyzerConfig` dataclass and its sub-configurations, this system provides modular, extensible configuration management for all aspects of workspace analysis. + +## Table of Contents + +- [Overview](#overview) +- [Configuration Types](#configuration-types) +- [Usage Examples](#usage-examples) +- [Best Practices](#best-practices) +- [API Reference](#api-reference) + +## Overview + +The configs module enables fine-grained control over: + +- **Sampling Strategies**: Control how joint space is sampled +- **Caching Behavior**: Configure result caching and storage +- **Metric Calculation**: Define which metrics to compute and how +- **Visualization Options**: Customize visualization appearance and behavior +- **Dimensional Constraints**: Set workspace bounds and exclusion zones + +### Key Features + +- **Type Safety**: All configurations use dataclasses with type hints +- **Validation**: Built-in validation and sensible defaults +- **Extensibility**: Easy to extend with custom configuration options +- **Documentation**: Comprehensive docstrings for all parameters +- **Modularity**: Separate configs for different analysis aspects + +## Configuration Types + +### 1. Sampling Configuration + +Controls how the robot's joint space is sampled for workspace analysis. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + SamplingConfig, SamplingStrategy +) + +# Basic uniform sampling +config = SamplingConfig( + strategy=SamplingStrategy.UNIFORM, + num_samples=10000, + grid_resolution=20, + seed=42 +) + +# Random sampling with larger batch size +random_config = SamplingConfig( + strategy=SamplingStrategy.RANDOM, + num_samples=50000, + batch_size=5000, + seed=123 +) +``` + +#### Available Sampling Strategies + +- **UNIFORM**: Uniform grid sampling across joint space +- **RANDOM**: Pseudo-random sampling with specified seed +- **HALTON**: Quasi-random Halton sequence (better coverage) +- **SOBOL**: Quasi-random Sobol sequence (low-discrepancy) +- **LATIN_HYPERCUBE**: Latin Hypercube Sampling (space-filling) +- **IMPORTANCE**: Importance sampling with custom weight function +- **GAUSSIAN**: Gaussian distribution around specified mean + +#### Advanced Sampling Examples + +```python +# Quasi-random sampling for better coverage +halton_config = SamplingConfig( + strategy=SamplingStrategy.HALTON, + num_samples=25000, + batch_size=2500 +) + +# Gaussian sampling around robot's nominal pose +gaussian_config = SamplingConfig( + strategy=SamplingStrategy.GAUSSIAN, + num_samples=15000, + gaussian_mean=0.0, # Center of joint range + gaussian_std=0.3 # 30% of range as std dev +) + +# Importance sampling with custom weight function +def weight_function(joint_values): + # Higher weight for configurations closer to zero + return np.exp(-np.sum(joint_values**2)) + +importance_config = SamplingConfig( + strategy=SamplingStrategy.IMPORTANCE, + num_samples=20000, + importance_weight_func=weight_function +) +``` + +### 2. Cache Configuration + +Manages caching of analysis results for improved performance. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import CacheConfig +from pathlib import Path + +# Basic caching configuration +cache_config = CacheConfig( + enabled=True, + cache_dir=Path("./workspace_cache"), + use_hash=True, + compression=True +) + +# High-performance caching setup +performance_cache = CacheConfig( + enabled=True, + cache_dir=Path("/fast_storage/workspace_cache"), + use_hash=True, + compression=False, # Faster access, more storage + max_cache_size_mb=5000, + cache_format="npz" +) +``` + +#### Cache Configuration Options + +- **enabled**: Enable/disable caching entirely +- **cache_dir**: Custom cache directory (default: system cache) +- **use_hash**: Hash-based cache keys for consistency +- **compression**: Compress cache files to save space +- **max_cache_size_mb**: Automatic cleanup of old cache files +- **cache_format**: Storage format ('npz', 'pkl', 'h5') + +### 3. Metric Configuration + +Defines which workspace metrics to compute and their parameters. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + MetricConfig, MetricType, ReachabilityConfig, + ManipulabilityConfig, DensityConfig +) + +# Comprehensive metric analysis +metric_config = MetricConfig( + enabled_metrics=[MetricType.ALL], + reachability=ReachabilityConfig( + voxel_size=0.005, # High resolution + min_points_per_voxel=2, + compute_coverage=True + ), + manipulability=ManipulabilityConfig( + jacobian_threshold=0.001, + compute_isotropy=True, + compute_heatmap=True + ), + density=DensityConfig( + radius=0.03, + k_neighbors=50, + compute_distribution=True + ), + save_results=True, + output_format="json" +) +``` + +#### Metric Types + +**Reachability Metrics**: + +- Workspace volume calculation +- Coverage percentage +- Reachable point density + +```python +reachability_config = ReachabilityConfig( + voxel_size=0.01, # Voxel resolution for volume + min_points_per_voxel=1, # Occupancy threshold + compute_coverage=True # Coverage vs bounding box +) +``` + +**Manipulability Metrics**: + +- Manipulability index throughout workspace +- Isotropy analysis +- Dexterity heatmaps + +```python +manipulability_config = ManipulabilityConfig( + jacobian_threshold=0.01, # Minimum valid manipulability + compute_isotropy=True, # Condition number analysis + compute_heatmap=False # Generate spatial heatmap +) +``` + +**Density Metrics**: + +- Local point density +- Distribution statistics +- Clustering analysis + +```python +density_config = DensityConfig( + radius=0.05, # Local neighborhood radius + k_neighbors=30, # Number of neighbors to consider + compute_distribution=True # Statistical distribution +) +``` + +### 4. Visualization Configuration + +Controls the appearance and behavior of workspace visualizations. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + VisualizationConfig, VisualizationType +) + +# High-quality point cloud visualization +viz_config = VisualizationConfig( + enabled=True, + vis_type=VisualizationType.POINT_CLOUD, + point_size=3.0, + alpha=0.8, + color_by_distance=True, + voxel_size=0.02, + is_voxel_down=True, + show_unreachable_points=True +) + +# Sphere visualization for presentations +sphere_viz_config = VisualizationConfig( + vis_type=VisualizationType.SPHERE, + sphere_radius=0.008, + sphere_resolution=15, + alpha=0.6, + show_unreachable_points=False +) +``` + +#### Visualization Options + +- **vis_type**: Visualization method (POINT_CLOUD, VOXEL, SPHERE, MESH, HEATMAP) +- **point_size**: Size of rendered points +- **alpha**: Transparency level (0.0-1.0) +- **color_by_distance**: Color-code by distance from base +- **voxel_size**: Downsampling resolution +- **show_unreachable_points**: Display unreachable regions + +### 5. Dimension Constraints + +Defines workspace bounds, exclusion zones, and physical constraints. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import DimensionConstraint +import numpy as np + +# Basic workspace bounds +dimension_config = DimensionConstraint( + min_bounds=np.array([-1.0, -1.0, 0.0]), # [x_min, y_min, z_min] + max_bounds=np.array([1.0, 1.0, 2.0]), # [x_max, y_max, z_max] + joint_limits_scale=0.95, # Use 95% of joint range + ground_height=0.0 +) + +# Complex workspace with exclusion zones +complex_constraints = DimensionConstraint( + min_bounds=np.array([-2.0, -1.5, -0.5]), + max_bounds=np.array([2.0, 1.5, 2.5]), + joint_limits_scale=0.9, + exclude_zones=[ + # Obstacle 1: table + (np.array([0.3, -0.5, 0.0]), np.array([1.2, 0.5, 0.8])), + # Obstacle 2: wall + (np.array([1.8, -1.5, 0.0]), np.array([2.0, 1.5, 2.5])) + ], + ground_height=-0.1, + enforce_collision_free=True, + self_collision_check=True +) +``` + +#### Constraint Options + +- **min_bounds/max_bounds**: Cartesian workspace limits +- **joint_limits_scale**: Scale factor for joint range usage +- **exclude_zones**: List of forbidden regions +- **ground_height**: Floor level for filtering +- **enforce_collision_free**: Enable collision checking +- **self_collision_check**: Check robot self-collisions + +## Usage Examples + +### Basic Workspace Analysis Setup + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import * + +# Complete configuration for basic analysis +sampling_config = SamplingConfig( + strategy=SamplingStrategy.UNIFORM, + num_samples=15000, + grid_resolution=25, + batch_size=1500 +) + +cache_config = CacheConfig( + enabled=True, + compression=True, + max_cache_size_mb=2000 +) + +metric_config = MetricConfig( + enabled_metrics=[MetricType.REACHABILITY, MetricType.DENSITY], + save_results=True +) + +viz_config = VisualizationConfig( + vis_type=VisualizationType.POINT_CLOUD, + point_size=2.5, + color_by_distance=True +) + +constraints = DimensionConstraint( + min_bounds=np.array([-1.2, -1.2, 0.0]), + max_bounds=np.array([1.2, 1.2, 1.8]), + joint_limits_scale=0.9 +) +``` + +### High-Performance Configuration + +```python +# Configuration optimized for speed and large datasets +fast_sampling = SamplingConfig( + strategy=SamplingStrategy.HALTON, # Better than random + num_samples=100000, # Large dataset + batch_size=10000, # Large batches + seed=42 +) + +fast_cache = CacheConfig( + enabled=True, + compression=False, # Faster access + cache_format="npz", # Efficient format + max_cache_size_mb=10000 # Large cache +) + +minimal_metrics = MetricConfig( + enabled_metrics=[MetricType.REACHABILITY], # Only essential + reachability=ReachabilityConfig( + voxel_size=0.02, # Lower resolution + compute_coverage=False # Skip expensive computation + ), + save_results=False # Don't save to disk +) + +efficient_viz = VisualizationConfig( + vis_type=VisualizationType.VOXEL, # Faster than spheres + voxel_size=0.03, # Aggressive downsampling + is_voxel_down=True, + show_unreachable_points=False # Reduce complexity +) +``` + +### Research Configuration + +```python +# Configuration for detailed research analysis +research_sampling = SamplingConfig( + strategy=SamplingStrategy.SOBOL, # Low-discrepancy sequence + num_samples=50000, + batch_size=2500 +) + +persistent_cache = CacheConfig( + enabled=True, + cache_dir=Path("./research_cache"), + use_hash=True, + compression=True, + cache_format="h5" # Good for large datasets +) + +comprehensive_metrics = MetricConfig( + enabled_metrics=[MetricType.ALL], + reachability=ReachabilityConfig( + voxel_size=0.005, # High resolution + compute_coverage=True + ), + manipulability=ManipulabilityConfig( + jacobian_threshold=0.001, + compute_isotropy=True, + compute_heatmap=True # Generate heatmaps + ), + density=DensityConfig( + radius=0.02, # Fine-grained density + k_neighbors=50, + compute_distribution=True + ), + save_results=True, + output_format="json" +) + +publication_viz = VisualizationConfig( + vis_type=VisualizationType.SPHERE, + sphere_radius=0.006, + sphere_resolution=20, # High quality spheres + alpha=0.7, + color_by_distance=True, + show_unreachable_points=True +) + +detailed_constraints = DimensionConstraint( + min_bounds=np.array([-1.5, -1.5, -0.2]), + max_bounds=np.array([1.5, 1.5, 2.0]), + joint_limits_scale=1.0, # Full joint range + exclude_zones=[], # No exclusions for research + enforce_collision_free=True, + self_collision_check=True +) +``` + +### Configuration Composition + +```python +# Create configurations programmatically +def create_analysis_config( + resolution: str = "medium", + enable_cache: bool = True, + visualization_type: str = "point_cloud" +) -> tuple: + """Create configuration set based on analysis requirements.""" + + # Resolution-dependent settings + resolution_params = { + "low": {"samples": 5000, "voxel": 0.05, "batch": 1000}, + "medium": {"samples": 20000, "voxel": 0.02, "batch": 2000}, + "high": {"samples": 80000, "voxel": 0.01, "batch": 5000} + } + params = resolution_params[resolution] + + sampling_config = SamplingConfig( + strategy=SamplingStrategy.HALTON, + num_samples=params["samples"], + batch_size=params["batch"] + ) + + cache_config = CacheConfig(enabled=enable_cache) + + viz_config = VisualizationConfig( + vis_type=getattr(VisualizationType, visualization_type.upper()), + voxel_size=params["voxel"] + ) + + return sampling_config, cache_config, viz_config + +# Usage +low_res_configs = create_analysis_config("low", True, "voxel") +high_res_configs = create_analysis_config("high", False, "sphere") +``` + +## Best Practices + +### Configuration Management + +1. **Use Type Hints**: Leverage the type safety provided by dataclasses +2. **Validate Parameters**: Check parameter ranges before analysis +3. **Document Choices**: Comment configuration choices for reproducibility +4. **Version Configs**: Save configuration alongside results +5. **Test Different Settings**: Experiment with sampling strategies + +### Performance Optimization + +```python +# Performance tips based on use case +def optimize_for_speed(): + return SamplingConfig( + strategy=SamplingStrategy.RANDOM, # Fastest generation + batch_size=10000, # Large batches + num_samples=20000 # Moderate resolution + ) + +def optimize_for_accuracy(): + return SamplingConfig( + strategy=SamplingStrategy.SOBOL, # Better coverage + batch_size=2000, # Moderate batches + num_samples=100000 # High resolution + ) + +def optimize_for_memory(): + return SamplingConfig( + batch_size=500, # Small batches + num_samples=15000 # Moderate total + ) +``` + +### Configuration Validation + +```python +def validate_config(config): + """Validate configuration parameters.""" + if config.num_samples <= 0: + raise ValueError("num_samples must be positive") + + if config.batch_size > config.num_samples: + raise ValueError("batch_size cannot exceed num_samples") + + if hasattr(config, 'voxel_size') and config.voxel_size <= 0: + raise ValueError("voxel_size must be positive") + +# Usage +try: + validate_config(my_config) +except ValueError as e: + print(f"Configuration error: {e}") +``` + +## API Reference + +### SamplingConfig + +Configuration for joint space sampling strategies. + +#### Parameters + +- **strategy**: `SamplingStrategy` - Sampling method to use +- **num_samples**: `int` - Total number of samples to generate (default: 1000) +- **grid_resolution**: `int` - Grid resolution for uniform sampling (default: 10) +- **batch_size**: `int` - Batch size for processing (default: 1000) +- **seed**: `int` - Random seed for reproducibility (default: 42) +- **importance_weight_func**: `Optional[Callable]` - Weight function for importance sampling +- **gaussian_mean**: `Optional[float]` - Mean for Gaussian sampling +- **gaussian_std**: `Optional[float]` - Standard deviation for Gaussian sampling + +### CacheConfig + +Configuration for result caching and storage. + +#### CacheConfig Parameters + +- **enabled**: `bool` - Enable caching (default: True) +- **cache_dir**: `Optional[Path]` - Cache directory path +- **use_hash**: `bool` - Use hash-based cache keys (default: True) +- **compression**: `bool` - Compress cache files (default: True) +- **max_cache_size_mb**: `int` - Maximum cache size in MB (default: 1000) +- **cache_format**: `str` - Cache file format: 'npz', 'pkl', 'h5' (default: 'npz') + +### MetricConfig + +Configuration for workspace analysis metrics. + +#### MetricConfig Parameters + +- **enabled_metrics**: `List[MetricType]` - List of metrics to compute +- **reachability**: `ReachabilityConfig` - Reachability analysis settings +- **manipulability**: `ManipulabilityConfig` - Manipulability analysis settings +- **density**: `DensityConfig` - Density analysis settings +- **save_results**: `bool` - Save results to file (default: True) +- **output_format**: `str` - Output format: 'json', 'yaml', 'pkl' (default: 'json') + +### VisualizationConfig + +Configuration for visualization appearance and behavior. + +#### VisualizationConfig Parameters + +- **enabled**: `bool` - Enable visualization (default: True) +- **vis_type**: `VisualizationType` - Visualization type (default: POINT_CLOUD) +- **voxel_size**: `float` - Voxel downsampling size (default: 0.05) +- **point_size**: `float` - Point size in visualization (default: 4.0) +- **alpha**: `float` - Transparency level 0.0-1.0 (default: 0.5) +- **color_by_distance**: `bool` - Color by distance from base (default: True) +- **sphere_radius**: `float` - Sphere radius for sphere visualization (default: 0.005) +- **show_unreachable_points**: `bool` - Show unreachable points (default: False) + +### DimensionConstraint + +Configuration for workspace bounds and constraints. + +#### DimensionConstraint Parameters + +- **min_bounds**: `Optional[np.ndarray]` - Minimum workspace bounds [x, y, z] +- **max_bounds**: `Optional[np.ndarray]` - Maximum workspace bounds [x, y, z] +- **joint_limits_scale**: `float` - Joint range scale factor (default: 1.0) +- **exclude_zones**: `List[Tuple[np.ndarray, np.ndarray]]` - Exclusion zones +- **ground_height**: `float` - Ground plane height (default: 0.0) +- **enforce_collision_free**: `bool` - Enable collision checking (default: False) +- **self_collision_check**: `bool` - Check self-collisions (default: False) + +## Configuration Examples + +### Minimal Configuration + +```python +# Simplest possible configuration +from embodichain.lab.sim.utility.workspace_analyzer.configs import * + +config = SamplingConfig() # Uses defaults +# Result: uniform sampling, 1000 samples, grid resolution 10 +``` + +### Production Configuration + +```python +# Production-ready configuration +production_sampling = SamplingConfig( + strategy=SamplingStrategy.HALTON, + num_samples=25000, + batch_size=2500, + seed=42 +) + +production_cache = CacheConfig( + enabled=True, + cache_dir=Path("/app/cache"), + compression=True, + max_cache_size_mb=5000 +) + +production_metrics = MetricConfig( + enabled_metrics=[MetricType.REACHABILITY, MetricType.MANIPULABILITY], + save_results=True, + output_format="json" +) +``` + +### Debug Configuration + +```python +# Configuration for debugging and development +debug_config = SamplingConfig( + strategy=SamplingStrategy.UNIFORM, + num_samples=500, # Small for fast testing + batch_size=100, # Small batches + grid_resolution=8 # Low resolution +) + +debug_viz = VisualizationConfig( + vis_type=VisualizationType.POINT_CLOUD, + point_size=5.0, # Large points for visibility + show_unreachable_points=True, # Show all data + alpha=0.9 # High opacity +) +``` diff --git a/docs/source/features/workspace_analyzer/constraints.md b/docs/source/features/workspace_analyzer/constraints.md new file mode 100644 index 0000000..0f71148 --- /dev/null +++ b/docs/source/features/workspace_analyzer/constraints.md @@ -0,0 +1,415 @@ +# Workspace Analyzer Constraints + +The **Workspace Analyzer Constraints** module provides comprehensive constraint checking and validation for robotic workspace analysis. This module ensures that workspace analysis respects physical limitations, safety boundaries, and operational constraints. + +## Table of Contents + +- [Overview](#overview) +- [Constraint Types](#constraint-types) +- [Usage Examples](#usage-examples) +- [Best Practices](#best-practices) +- [API Reference](#api-reference) + +## Overview + +The constraints module provides basic constraint checking for workspace analysis: + +- **Spatial Boundaries**: Define workspace limits using min/max bounds +- **Exclusion Zones**: Avoid specified rectangular obstacle regions +- **Ground Constraints**: Handle floor plane height limits +- **Configuration-based Setup**: Easy configuration through DimensionConstraint + +### Key Features + +- **Simple Boundary Checking**: Min/max bounds validation +- **Basic Obstacle Avoidance**: Rectangular exclusion zone filtering +- **Multi-Backend Support**: Works with both NumPy arrays and PyTorch tensors +- **Configuration Integration**: Uses DimensionConstraint config objects + +### TODO - Not Yet Implemented + +- **Joint Limit Constraints**: Physical joint limitation checking +- **Advanced Collision Detection**: Robot mesh-based collision checking with environment +- **Self-Collision Detection**: Robot self-collision validation +- **Complex Safety Margins**: Advanced safety margin calculations +- **Dynamic Constraint Updates**: Real-time constraint modification +- **Performance Statistics**: Constraint checking metrics + +**Note**: The current `check_collision()` method only performs simple rectangular exclusion zone checking, not true geometric collision detection. + +## Constraint Types + +### 1. Spatial Boundary Constraints + +Define workspace boundaries using min/max bounds: + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Basic workspace boundaries +checker = WorkspaceConstraintChecker( + min_bounds=np.array([-1.0, -1.0, 0.0]), # [x_min, y_min, z_min] + max_bounds=np.array([1.0, 1.0, 2.0]), # [x_max, y_max, z_max] + ground_height=0.0 +) +``` + +### 2. Exclusion Zone Constraints + +Avoid specified rectangular obstacle regions (simple collision avoidance): + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Workspace with exclusion zones +checker = WorkspaceConstraintChecker( + min_bounds=np.array([-2.0, -1.5, 0.0]), + max_bounds=np.array([2.0, 1.5, 2.0]), + exclude_zones=[ + # Table obstacle + (np.array([0.3, -0.5, 0.0]), np.array([1.2, 0.5, 0.8])), + # Wall section + (np.array([1.8, -1.5, 0.0]), np.array([2.0, 1.5, 2.0])) + ] +) +``` + +### 3. Configuration-based Setup (Recommended) + +Use DimensionConstraint for easy configuration: + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.configs.dimension_constraint import DimensionConstraint +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Create configuration +config = DimensionConstraint( + min_bounds=np.array([-1.0, -1.0, 0.0]), + max_bounds=np.array([1.0, 1.0, 2.0]), + ground_height=0.0, + exclude_zones=[ + (np.array([0.2, 0.2, 0.0]), np.array([0.4, 0.4, 0.2])) + ] +) + +# Create checker from config +checker = WorkspaceConstraintChecker.from_config(config) +``` + +### TODO - Planned Features + +#### Advanced Constraint Features (Not Implemented) + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.configs.dimension_constraint import DimensionConstraint + +# TODO: These parameters exist in DimensionConstraint but are not implemented +config = DimensionConstraint( + joint_limits_scale=0.95, # TODO: Joint limit checking not implemented + enforce_collision_free=True, # TODO: Mesh-based collision not implemented + self_collision_check=True # TODO: Self-collision checking not implemented +) +``` + +**Current Implementation Status**: + +- ✅ **Simple exclusion zones**: Rectangular bounding box obstacle avoidance +- ❌ **Joint limits**: Physical joint limitation checking not implemented +- ❌ **Mesh collision**: Advanced geometric collision detection not implemented +- ❌ **Self-collision**: Robot self-collision validation not implemented + +**Planned Features**: + +- Enforce mechanical joint limits and prevent over-extension +- True mesh-based collision detection with environment geometry +- Self-collision checking between robot links +- Singular configuration avoidance + +## Usage Examples + +### Basic Constraint Checking + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Create constraint checker +checker = WorkspaceConstraintChecker( + min_bounds=np.array([-1.0, -1.0, 0.0]), + max_bounds=np.array([1.0, 1.0, 1.5]), + ground_height=0.0 +) + +# Test points +test_points = np.array([ + [0.5, 0.5, 0.5], # Valid point + [2.0, 0.0, 0.5], # Outside x boundary + [0.0, 0.0, -0.1], # Below ground + [0.8, 0.8, 1.2] # Valid point +]) + +# Check boundaries +valid_bounds = checker.check_bounds(test_points) +print(f"Valid bounds: {valid_bounds}") # [True, False, False, True] + +# Filter valid points +valid_points = checker.filter_points(test_points) +print(f"Filtered points shape: {valid_points.shape}") # (2, 3) +``` + +### Workspace with Obstacles + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Workspace with exclusion zones +checker = WorkspaceConstraintChecker( + min_bounds=np.array([-1.5, -1.5, 0.0]), + max_bounds=np.array([1.5, 1.5, 2.0]), + exclude_zones=[ + # Table obstacle + (np.array([0.2, -0.4, 0.0]), np.array([0.8, 0.4, 0.8])), + ] +) + +# Test points +test_points = np.array([ + [0.0, 0.0, 0.5], # Valid point + [0.5, 0.0, 0.4], # Inside table (excluded) + [1.0, 1.0, 1.0] # Valid point +]) + +# Check collision (simple exclusion zone checking) +valid_collision = checker.check_collision(test_points) +print(f"Collision-free: {valid_collision}") # [True, False, True] + +# Comprehensive filtering (bounds + collision) +safe_points = checker.filter_points(test_points) +print(f"Safe points: {len(safe_points)} out of {len(test_points)}") +``` + +### Configuration-based Setup + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.configs.dimension_constraint import DimensionConstraint +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Create configuration object +config = DimensionConstraint( + min_bounds=np.array([-1.0, -1.0, 0.0]), + max_bounds=np.array([1.0, 1.0, 1.5]), + ground_height=0.0, + exclude_zones=[ + # Add table obstacle + (np.array([0.2, 0.2, 0.0]), np.array([0.6, 0.6, 0.8])) + ] +) + +# Create checker from config +checker = WorkspaceConstraintChecker.from_config(config) + +# Add more obstacles dynamically +checker.add_exclude_zone( + min_bounds=np.array([-0.3, 0.7, 0.0]), + max_bounds=np.array([0.3, 1.0, 1.2]) +) + +print(f"Total exclusion zones: {checker.get_num_exclude_zones()}") +``` + +### Dynamic Obstacle Management + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Add and remove obstacles dynamically +checker = WorkspaceConstraintChecker( + min_bounds=np.array([-1.0, -1.0, 0.0]), + max_bounds=np.array([1.0, 1.0, 1.5]) +) + +# Add obstacles +checker.add_exclude_zone( + min_bounds=np.array([0.2, 0.2, 0.0]), + max_bounds=np.array([0.4, 0.4, 0.8]) +) + +print(f"Exclusion zones: {checker.get_num_exclude_zones()}") + +# Clear all obstacles +checker.clear_exclude_zones() +print(f"After clearing: {checker.get_num_exclude_zones()} zones") +``` + +### Integration with Workspace Analysis + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.configs.dimension_constraint import DimensionConstraint +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Define workspace constraints +constraints = DimensionConstraint( + min_bounds=np.array([-0.8, -0.6, 0.0]), + max_bounds=np.array([0.8, 0.6, 1.2]), + exclude_zones=[ + # Table obstacle + (np.array([0.2, 0.1, 0.0]), np.array([0.6, 0.4, 0.8])) + ] +) + +# Create checker +checker = WorkspaceConstraintChecker.from_config(constraints) + +# Filter workspace points +workspace_points = np.random.uniform(-1, 1, size=(10000, 3)) +valid_points = checker.filter_points(workspace_points) + +print(f"Valid workspace points: {len(valid_points)}/{len(workspace_points)}") +print(f"Coverage: {len(valid_points)/len(workspace_points)*100:.1f}%") +``` + +## Best Practices + +### Choosing Bounds + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Conservative workspace bounds +def create_safe_bounds(robot_reach, safety_margin=0.1): + """Create workspace bounds with safety margins.""" + return { + 'min_bounds': np.array([-robot_reach + safety_margin] * 3), + 'max_bounds': np.array([robot_reach - safety_margin] * 3), + 'ground_height': safety_margin + } + +# Example usage +bounds = create_safe_bounds(robot_reach=1.2, safety_margin=0.1) +checker = WorkspaceConstraintChecker(**bounds) +``` + +### Validation + +```python +import numpy as np + +def validate_bounds(min_bounds, max_bounds): + """Validate constraint bounds.""" + if min_bounds is not None and max_bounds is not None: + if not np.all(min_bounds < max_bounds): + raise ValueError("min_bounds must be less than max_bounds") + + for bounds in [min_bounds, max_bounds]: + if bounds is not None and len(bounds) != 3: + raise ValueError("Bounds must have exactly 3 dimensions") + +# Use validation +min_bounds = np.array([-1, -1, 0]) +max_bounds = np.array([1, 1, 2]) +validate_bounds(min_bounds, max_bounds) +``` + +## API Reference + +### WorkspaceConstraintChecker + +Main constraint checker implementation. + +#### Constructor Parameters + +- **min_bounds**: `Optional[np.ndarray]` - Minimum bounds [x, y, z] +- **max_bounds**: `Optional[np.ndarray]` - Maximum bounds [x, y, z] +- **ground_height**: `float` - Ground plane height (default: 0.0) +- **exclude_zones**: `List[Tuple[np.ndarray, np.ndarray]]` - Exclusion zones +- **device**: `Optional[torch.device]` - PyTorch device + +#### Methods + +- **check_bounds(points)**: Check if points are within bounds +- **check_collision(points)**: Check if points avoid rectangular exclusion zones (simple collision) +- **filter_points(points)**: Filter points by all constraints +- **add_exclude_zone(min_bounds, max_bounds)**: Add exclusion zone +- **clear_exclude_zones()**: Remove all exclusion zones +- **get_num_exclude_zones()**: Get number of exclusion zones + +#### Class Methods + +- **from_config(config)**: Create checker from DimensionConstraint + +### DimensionConstraint + +Configuration dataclass for constraints. + +#### Parameters + +- **min_bounds**: `Optional[np.ndarray]` - Workspace minimum bounds +- **max_bounds**: `Optional[np.ndarray]` - Workspace maximum bounds +- **exclude_zones**: `List[Tuple[np.ndarray, np.ndarray]]` - Exclusion zones +- **ground_height**: `float` - Ground plane height (default: 0.0) +- **joint_limits_scale**: `float` - Joint limits scaling (TODO: not implemented) +- **enforce_collision_free**: `bool` - Advanced collision checking (TODO: not implemented) +- **self_collision_check**: `bool` - Self-collision check (TODO: not implemented) + +**Note**: These advanced collision detection features are defined in the config but not yet used by WorkspaceConstraintChecker. + +## Quick Start + +### Basic Setup + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Simple boundary constraints +checker = WorkspaceConstraintChecker( + min_bounds=np.array([-1, -1, 0]), + max_bounds=np.array([1, 1, 1]) +) +``` + +### With Obstacles + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Workspace with exclusion zones +checker = WorkspaceConstraintChecker( + min_bounds=np.array([-1.2, -1.2, 0.05]), + max_bounds=np.array([1.2, 1.2, 1.8]), + exclude_zones=[ + # Table obstacle + (np.array([0.3, 0.2, 0.0]), np.array([0.8, 0.7, 0.8])) + ] +) +``` + +### Using Configuration (Recommended) + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.configs.dimension_constraint import DimensionConstraint +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Create config +config = DimensionConstraint( + min_bounds=np.array([-1.0, -1.0, 0.0]), + max_bounds=np.array([1.0, 1.0, 1.5]), + exclude_zones=[ + (np.array([0.2, 0.2, 0.0]), np.array([0.4, 0.4, 0.8])) + ] +) + +# Create checker +checker = WorkspaceConstraintChecker.from_config(config) +``` diff --git a/docs/source/features/workspace_analyzer/index.rst b/docs/source/features/workspace_analyzer/index.rst new file mode 100644 index 0000000..cb4cbc5 --- /dev/null +++ b/docs/source/features/workspace_analyzer/index.rst @@ -0,0 +1,16 @@ +Workspace Analyzer +====================== + +The Workspace Analyzer provides comprehensive analysis of robot workspace characteristics, including joint space and Cartesian space analysis capabilities for understanding robot reachability, workspace volume, and performance metrics. + +.. toctree:: + :maxdepth: 2 + + Overview and Usage + Configuration System + Sampling Strategies + Visualization Options + Caching System + Constraint Handling + Metric Computation + \ No newline at end of file diff --git a/docs/source/features/workspace_analyzer/metrics.md b/docs/source/features/workspace_analyzer/metrics.md new file mode 100644 index 0000000..a03908f --- /dev/null +++ b/docs/source/features/workspace_analyzer/metrics.md @@ -0,0 +1,169 @@ +# Workspace Analyzer Metrics + +The **Workspace Analyzer Metrics** module provides comprehensive metrics and analysis tools for evaluating robotic workspace characteristics. This module implements various quantitative measures to assess workspace quality, coverage, manipulability, and performance. + +## Overview + +This module provides three main workspace analysis metrics: + +- **ReachabilityMetric**: Voxel-based volume and coverage analysis ✅ +- **ManipulabilityMetric**: Dexterity analysis (heuristic method or with Jacobians) ✅ +- **DensityMetric**: Point distribution and density analysis ✅ + +**Note**: Some advanced features like spatial heatmaps, clustering analysis, and robot-specific helpers are not yet implemented. + +## Table of Contents + +- [Metric Types](#metric-types) +- [Usage Examples](#usage-examples) +- [Configuration](#configuration) +- [Quick Reference](#quick-reference) + +## Metric Types + +### 1. ReachabilityMetric + +Computes workspace volume and coverage using voxel-based analysis. + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.metrics import ReachabilityMetric + +# Basic usage +workspace_points = np.random.uniform(-1, 1, size=(1000, 3)) +reachability = ReachabilityMetric() +results = reachability.compute(workspace_points) + +print(f"Volume: {results['volume']:.4f} m³") +print(f"Coverage: {results['coverage']:.1f}%") +``` + +**Returns**: volume, coverage percentage, voxel count, bounding box, centroid + +### 2. ManipulabilityMetric + +Analyzes dexterity using distance-based heuristic or Yoshikawa index (with Jacobians). + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.metrics import ManipulabilityMetric + +# Basic usage (heuristic method) +workspace_points = np.random.uniform(-1, 1, size=(1000, 3)) +manipulability = ManipulabilityMetric() +results = manipulability.compute(workspace_points) + +print(f"Mean manipulability: {results['mean_manipulability']:.3f}") +``` + +**Returns**: mean, std, min, max manipulability; condition numbers (if Jacobians provided) + +### 3. DensityMetric + +Computes local point density using radius-based neighborhood analysis. + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.metrics import DensityMetric + +# Basic usage +workspace_points = np.random.uniform(-1, 1, size=(1000, 3)) +density = DensityMetric() +results = density.compute(workspace_points) + +print(f"Mean density: {results['mean_density']:.2f}") +``` + +**Returns**: mean, std, min, max density; distribution histogram (if enabled) + +## Usage Examples + +### Basic Usage + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.metrics import ( + ReachabilityMetric, ManipulabilityMetric, DensityMetric +) + +# Generate sample data +workspace_points = np.random.uniform(-1, 1, size=(5000, 3)) + +# Compute all metrics +reach_results = ReachabilityMetric().compute(workspace_points) +manip_results = ManipulabilityMetric().compute(workspace_points) +density_results = DensityMetric().compute(workspace_points) + +print(f"Volume: {reach_results['volume']:.4f} m³") +print(f"Mean manipulability: {manip_results['mean_manipulability']:.3f}") +print(f"Mean density: {density_results['mean_density']:.2f}") +``` + +### With Custom Configuration + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + ReachabilityConfig, DensityConfig +) + +# Custom configuration +reach_config = ReachabilityConfig(voxel_size=0.02, compute_coverage=True) +density_config = DensityConfig(radius=0.1, compute_distribution=True) + +reachability = ReachabilityMetric(reach_config) +density = DensityMetric(density_config) + +reach_results = reachability.compute(workspace_points) +density_results = density.compute(workspace_points) +``` + +## Configuration + +All metrics can be customized using configuration classes: + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + ReachabilityConfig, ManipulabilityConfig, DensityConfig +) + +# Reachability configuration +reach_config = ReachabilityConfig( + voxel_size=0.01, # Voxel size in meters + min_points_per_voxel=1, # Minimum points for occupancy + compute_coverage=True # Enable coverage calculation +) + +# Manipulability configuration +manip_config = ManipulabilityConfig( + jacobian_threshold=0.01, # Minimum valid manipulability + compute_isotropy=True # Enable condition number analysis +) + +# Density configuration +density_config = DensityConfig( + radius=0.05, # Neighborhood radius + compute_distribution=True # Enable histogram computation +) +``` + +## API Reference + +**ReachabilityMetric**: Returns volume, coverage, num_voxels, bounding_box, centroid + +**ManipulabilityMetric**: Returns mean/std/min/max manipulability, num_valid_points, condition numbers (if available) + +**DensityMetric**: Returns mean/std/min/max density, distribution histogram (if enabled) + +## Quick Reference + +```python +# Default usage +reachability = ReachabilityMetric() +manipulability = ManipulabilityMetric() +density = DensityMetric() + +# Custom configuration +from embodichain.lab.sim.utility.workspace_analyzer.configs import ReachabilityConfig +reach_config = ReachabilityConfig(voxel_size=0.02, compute_coverage=True) +reachability = ReachabilityMetric(reach_config) +``` diff --git a/docs/source/features/workspace_analyzer/samplers.md b/docs/source/features/workspace_analyzer/samplers.md new file mode 100644 index 0000000..e24bb51 --- /dev/null +++ b/docs/source/features/workspace_analyzer/samplers.md @@ -0,0 +1,156 @@ +# Workspace Analyzer Samplers + +The samplers module provides various sampling strategies for workspace analysis, from uniform grids to quasi-random sequences. + +## Table of Contents + +- [Quick Start](#quick-start) +- [Overview](#overview) +- [Basic Usage](#basic-usage) +- [Available Samplers](#available-samplers) + - [1. UniformSampler](#1-uniformsampler) + - [2. RandomSampler](#2-randomsampler) + - [3. GaussianSampler](#3-gaussiansampler-) +- [Factory Pattern Usage](#factory-pattern-usage) +- [Integration with Workspace Analyzer](#integration-with-workspace-analyzer) +- [Quick Reference](#quick-reference) + +## Quick Start + +```python +import torch +from embodichain.lab.devices.workspace_analyzer.samplers import UniformSampler + +# Define sampling bounds +bounds = torch.tensor([[0.0, 1.0], # x-axis range + [0.0, 1.0]]) # y-axis range + +# Create uniform sampler and generate 1000 samples +sampler = UniformSampler(samples_per_dim=10) +samples = sampler.sample(1000, bounds) # Note: (num_samples, bounds) order +``` + +## Overview + +Available sampling strategies: + +- **UniformSampler**: Grid-based regular sampling ✅ +- **RandomSampler**: Pure random sampling ✅ +- **GaussianSampler**: Normal distribution sampling ✅ + +## Basic Usage + +```python +import torch +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.samplers import ( + UniformSampler, RandomSampler, GaussianSampler +) + +# Define bounds: [[min, max], [min, max], ...] +bounds = torch.tensor([[-1, 1], [-1, 1], [0, 2]], dtype=torch.float32) + +# Generate samples with implemented samplers +samples = UniformSampler(samples_per_dim=10).sample(1000, bounds) +print(f"Generated {samples.shape[0]} samples") +``` + +## Available Samplers + +### 1. UniformSampler + +Grid-based sampling with regular spacing. + +```python +import torch +from embodichain.lab.sim.utility.workspace_analyzer.samplers import UniformSampler + +bounds = torch.tensor([[-1, 1], [-1, 1], [0, 2]], dtype=torch.float32) +# Create uniform sampler +uniform_sampler = UniformSampler(seed=42, samples_per_dim=10) +samples = uniform_sampler.sample(1000, bounds) +``` + +**Best for**: Low-dimensional spaces (2-4D), systematic exploration + +### 2. RandomSampler + +Pure random sampling with uniform distribution. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.samplers import RandomSampler + +random_sampler = RandomSampler(seed=42) +samples = random_sampler.sample(1000, bounds) +``` + +**Best for**: High-dimensional spaces, baseline comparisons + +### 3. GaussianSampler ✅ + +Gaussian distribution sampling around specified mean. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.samplers import GaussianSampler + +gaussian_sampler = GaussianSampler(seed=42, std=0.2) +samples = gaussian_sampler.sample(1000, bounds) +``` + +**Best for**: Uncertainty analysis, robustness studies around workspace center + +## Factory Pattern Usage + +Use the factory to create samplers by strategy (only implemented samplers): + +```python +from embodichain.lab.sim.utility.workspace_analyzer.samplers import create_sampler +from embodichain.lab.sim.utility.workspace_analyzer.configs import SamplingStrategy + +# Create implemented samplers by strategy +uniform_sampler = create_sampler( + SamplingStrategy.UNIFORM, + seed=42, + samples_per_dim=10 +) + +random_sampler = create_sampler( + SamplingStrategy.RANDOM, + seed=42 +) + +gaussian_sampler = create_sampler( + SamplingStrategy.GAUSSIAN, + seed=42, + std=0.2 +) +``` + +## Integration with Workspace Analyzer + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import SamplingConfig, SamplingStrategy + +# Configure sampling in analyzer +sampling_config = SamplingConfig( + strategy=SamplingStrategy.UNIFORM, + num_samples=1000, + seed=42 +) + +# Analyzer will use the specified sampler automatically +``` + +## Quick Reference + +**Available Samplers**: + +- **UniformSampler**: Use for 2-4D systematic exploration +- **RandomSampler**: Baseline for any dimensional spaces +- **GaussianSampler**: Uncertainty and robustness analysis + +**Sample Size Guidelines**: + +- **Uniform**: `samples_per_dim^n_dims` (exponential growth) +- **Random**: Any size (flexible) +- **Gaussian**: Any size (flexible) diff --git a/docs/source/features/workspace_analyzer/visualizers.md b/docs/source/features/workspace_analyzer/visualizers.md new file mode 100644 index 0000000..5df8af2 --- /dev/null +++ b/docs/source/features/workspace_analyzer/visualizers.md @@ -0,0 +1,146 @@ +# Workspace Analyzer Visualizers + +The visualizers module provides visualization tools for analyzing robotic workspace data in 3D space. + +## Table of Contents + +- [Overview](#overview) +- [Visualization Types](#visualization-types) +- [Usage Examples](#usage-examples) +- [Backend Support](#backend-support) +- [Quick Reference](#quick-reference) + +## Overview + +The visualizers module enables: + +- **Workspace reachability visualization** with multiple rendering styles +- **3D point cloud, voxel, sphere, and coordinate axis representations** +- **Multiple backends**: Open3D, Matplotlib, and simulation environments +- **Factory pattern** for easy visualizer creation + +## Visualization Types + +### 1. Point Cloud Visualizer ✅ + +Fast rendering for large datasets. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.visualizers import PointCloudVisualizer + +visualizer = PointCloudVisualizer(backend='open3d', point_size=2.0) +``` + +**Best for**: Large point sets (>10k points), fast rendering + +### 2. Voxel Visualizer ✅ + +Volumetric representation for occupancy maps. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.visualizers import VoxelVisualizer + +visualizer = VoxelVisualizer(backend='open3d', voxel_size=0.01) +``` + +**Best for**: Occupancy grids, collision detection + +### 3. Sphere Visualizer ✅ + +Smooth visualization for reachability zones. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.visualizers import SphereVisualizer + +visualizer = SphereVisualizer(backend='open3d', sphere_radius=0.005) +``` + +**Best for**: Publication-quality figures, smooth appearance + +### 4. Axis Visualizer ✅ + +Coordinate frame visualization for poses and transformations. + +```python +from embodichain.lab.sim.utility.workspace_analyzer.visualizers import AxisVisualizer + +visualizer = AxisVisualizer(backend='sim_manager', sim_manager=sim, axis_length=0.15) +``` + +**Best for**: Robot poses, end-effector frames, coordinate system visualization + +## Usage Examples + +### Basic Usage + +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.visualizers import create_visualizer + +# Generate workspace data +points = np.random.rand(1000, 3) * 2 - 1 # Random points in [-1, 1]³ +colors = np.random.rand(1000, 3) # Random colors + +# Create and use visualizer +visualizer = create_visualizer('point_cloud', backend='open3d', point_size=3.0) +result = visualizer.visualize(points, colors) +visualizer.show() +``` + +### Factory Pattern + +```python +from embodichain.lab.sim.utility.workspace_analyzer.visualizers import VisualizerFactory + +factory = VisualizerFactory() +visualizer = factory.create_visualizer('voxel', backend='open3d', voxel_size=0.02) + +# Create axis visualizer +axis_viz = factory.create_visualizer('axis', backend='sim_manager', sim_manager=sim) +``` + +## Backend Support + +### Available Backends + +- **`open3d`**: Interactive 3D visualization (requires `pip install open3d`) +- **`matplotlib`**: Static figures and plots (requires `pip install matplotlib`) +- **`sim_manager`**: Integration with simulation environment +- **`data`**: Returns processed data without visualization (headless mode) + +## Quick Reference + +**Available Visualizers**: + +- **PointCloudVisualizer**: Fast rendering for large datasets +- **VoxelVisualizer**: Volumetric representation for occupancy maps +- **SphereVisualizer**: Smooth visualization for publication figures +- **AxisVisualizer**: Coordinate frame visualization for poses + +**Common Parameters**: + +- **PointCloud**: `backend`, `point_size` +- **Voxel**: `backend`, `voxel_size` +- **Sphere**: `backend`, `sphere_radius`, `sphere_resolution` +- **Axis**: `backend`, `axis_length`, `axis_size`, `sim_manager` + +**Quick Creation**: + +```python +from embodichain.lab.sim.utility.workspace_analyzer.visualizers import create_visualizer +import numpy as np + +# Create any visualizer +visualizer = create_visualizer('point_cloud', backend='open3d', point_size=2.0) +visualizer = create_visualizer('voxel', backend='open3d', voxel_size=0.01) +visualizer = create_visualizer('sphere', backend='open3d', sphere_radius=0.005) + +# Create axis visualizer for coordinate frames +axis_viz = create_visualizer('axis', backend='sim_manager', sim_manager=sim, axis_length=0.05) + +# Visualize coordinate frame at robot end-effector +pose = np.eye(4) +pose[:3, 3] = [1.0, 0.5, 1.2] # Set position +axis_viz.visualize(pose, name_prefix="end_effector_frame") +axis_viz.show() +``` diff --git a/docs/source/features/workspace_analyzer/workspace_analyzer.md b/docs/source/features/workspace_analyzer/workspace_analyzer.md new file mode 100644 index 0000000..133ee0e --- /dev/null +++ b/docs/source/features/workspace_analyzer/workspace_analyzer.md @@ -0,0 +1,228 @@ +# Workspace Analyzer + +The Workspace Analyzer is a comprehensive tool in EmbodiChain for analyzing robot workspace characteristics, reachability, and performance metrics. It provides multiple analysis modes with advanced sampling strategies, caching mechanisms, and visualization capabilities. + +## Table of Contents + +1. [Quick Start](#quick-start) +2. [Analysis Modes](#analysis-modes) +3. [Configuration](#configuration) +4. [Usage Examples](#usage-examples) +5. [Best Practices](#best-practices) + +## Quick Start + +```python +import torch +import numpy as np +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.robots import DexforceW1Cfg +from embodichain.lab.sim.utility.workspace_analyzer import ( + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, + AnalysisMode +) + +# Setup simulation +sim = SimulationManager(SimulationManagerCfg(headless=False, sim_device="cpu")) + +# Add robot +robot = sim.add_robot(DexforceW1Cfg.from_dict({ + "uid": "dexforce_w1", + "version": "v021", + "arm_kind": "industrial" +})) + +# Quick analysis with defaults +analyzer = WorkspaceAnalyzer(robot=robot, sim_manager=sim) +results = analyzer.analyze(num_samples=1000, visualize=True) +print(f"Analysis complete: {results['num_reachable']} reachable points") +``` + +## Analysis Modes + +### Joint Space Analysis + +Analyzes robot workspace based on joint configurations. + +```python +wa_joint = WorkspaceAnalyzer(robot=robot, sim_manager=sim) +results = wa_joint.analyze(num_samples=1000, visualize=True) +print(f"Valid points: {results['num_valid']}/{results['num_samples']}") +``` + +
+
+ Joint Workspace +
Joint Workspace Analysis
+
+
+ +### Cartesian Space Analysis + +Analyzes reachable positions in 3D Cartesian space. + +```python +cartesian_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.CARTESIAN_SPACE, + visualization=VisualizationConfig(show_unreachable_points=False, point_size=8.0), + control_part_name="left_arm", +) + +wa_cartesian = WorkspaceAnalyzer(robot=robot, config=cartesian_config, sim_manager=sim) +results = wa_cartesian.analyze(num_samples=1000, visualize=True) +print(f"Reachable points: {results['num_reachable']}/{results['num_samples']}") +``` + +
+
+ pose_workspace1 +
(show_unreachable_points=True)
+
+
+ pose_workspace2 +
(show_unreachable_points=False)
+
+
+ +### Plane Sampling Analysis + +Analyzes reachability within a 2D plane (e.g., table height). + +```python +plane_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.PLANE_SAMPLING, + plane_normal=torch.tensor([0.0, 0.0, 1.0]), # Horizontal plane + plane_point=torch.tensor([0.0, 0.0, 1.2]), # Height z=1.2m + visualization=VisualizationConfig(show_unreachable_points=True, point_size=8.0), + control_part_name="left_arm", +) + +wa_plane = WorkspaceAnalyzer(robot=robot, config=plane_config, sim_manager=sim) +results = wa_plane.analyze(num_samples=1500, visualize=True) +print(f"Reachable points: {results['num_reachable']}/{results['num_samples']}") +``` + +
+
+ plane_workspace1 +
(show_unreachable_points=True)
+
+
+ plane_workspace2 +
(show_unreachable_points=False)
+
+
+ +## Configuration + +### Basic Configuration + +```python +from embodichain.lab.sim.utility.workspace_analyzer.configs import VisualizationConfig + +# Basic configuration with defaults +config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.CARTESIAN_SPACE +) + +# Custom visualization +vis_config = VisualizationConfig( + show_unreachable_points=False, # Only show reachable points + point_size=8.0 # Larger points for visibility +) + +config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.CARTESIAN_SPACE, + visualization=vis_config, + control_part_name="left_arm", +) +``` + +### Configuration Parameters + +**Main Parameters:** + +- `mode`: Analysis mode (JOINT_SPACE, CARTESIAN_SPACE, PLANE_SAMPLING) +- `plane_normal`: Normal vector for plane sampling (3D vector) +- `plane_point`: Point on plane for plane sampling (3D point) + +**Visualization Options:** + +- `show_unreachable_points`: Show failed samples (True) or only reachable (False) +- `point_size`: Size of visualization points (typically 6.0-10.0) + +## Usage Examples + +### Complete Example + +```python +import torch +import numpy as np +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.robots import DexforceW1Cfg +from embodichain.lab.sim.utility.workspace_analyzer import ( + WorkspaceAnalyzer, WorkspaceAnalyzerConfig, AnalysisMode +) +from embodichain.lab.sim.utility.workspace_analyzer.configs import VisualizationConfig + +# Setup simulation +sim = SimulationManager(SimulationManagerCfg(headless=False, sim_device="cpu")) + +# Add robot +robot = sim.add_robot(DexforceW1Cfg.from_dict({ + "uid": "dexforce_w1", "version": "v021", "arm_kind": "industrial" +})) + +# 1. Joint Space Analysis +wa_joint = WorkspaceAnalyzer(robot=robot, sim_manager=sim) +results = wa_joint.analyze(num_samples=1000, visualize=True) + +# 2. Cartesian Space Analysis +cartesian_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.CARTESIAN_SPACE, + visualization=VisualizationConfig(show_unreachable_points=False), + control_part_name="left_arm", +) +wa_cartesian = WorkspaceAnalyzer(robot=robot, config=cartesian_config, sim_manager=sim) +results = wa_cartesian.analyze(num_samples=1000, visualize=True) + +# 3. Plane Sampling +plane_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.PLANE_SAMPLING, + plane_normal=torch.tensor([0.0, 0.0, 1.0]), # Horizontal plane + plane_point=torch.tensor([0.0, 0.0, 1.2]), # Height 1.2m + control_part_name="left_arm", # robot control part name +) +wa_plane = WorkspaceAnalyzer(robot=robot, config=plane_config, sim_manager=sim) +results = wa_plane.analyze(num_samples=1500, visualize=True) +``` + +## Best Practices + +**Sample Size Guidelines:** + +- Joint space: 1000-3000 samples +- Cartesian space: 1000-2000 samples +- Plane sampling: 1000-1500 samples + +**Visualization Tips:** + +- Use `show_unreachable_points=False` for clean workspace boundaries +- Use `show_unreachable_points=True` for debugging and complete coverage +- Set `point_size=8.0` for better visibility +- Use `headless=False` for visualization, `headless=True` for batch processing + +**Performance:** + +- Start with smaller sample sizes (1000) for testing +- Use CPU device for consistent results +- Disable visualization for large batch analyses + +**Results:** + +- `num_valid`: Valid joint configurations (Joint Space) +- `num_reachable`: Reachable Cartesian points (Cartesian/Plane) +- `num_samples`: Total samples tested +- `analysis_time`: Time taken for analysis +- `metrics`: Workspace volume and other metrics diff --git a/docs/source/index.rst b/docs/source/index.rst index 2058a8b..45814c4 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -32,6 +32,13 @@ Table of Contents overview/vla/index overview/rl/index +.. toctree:: + :maxdepth: 1 + :caption: Features + :glob: + + features/workspace_analyzer/index* + .. toctree:: :maxdepth: 1 :caption: Resources diff --git a/docs/source/quick_start/install.md b/docs/source/quick_start/install.md index 6389101..fd00abe 100644 --- a/docs/source/quick_start/install.md +++ b/docs/source/quick_start/install.md @@ -4,22 +4,21 @@ The following minimum system requirements are recommended to run EmbodiChain reliably. These are the tested configurations during development — other Linux distributions and versions may work but are not officially supported. -- Operating System: Linux (x86_64) - - Recommended distributions: Ubuntu 20.04 LTS or Ubuntu 22.04 LTS +- Operating System: + - Linux (x86_64): Ubuntu 20.04+ - NVIDIA GPU and drivers: - - Hardware: NVIDIA GPU with compute capability 7.0 or higher (e.g., RTX 20 series, RTX 30 series, A100, etc.) - - NVIDIA driver: 535 or higher (recommended 570) - - CUDA Toolkit: any of 11.8 — 12.8 (we test primarily with 11.8 and 12.x) + - Hardware: NVIDIA GPU with compute capability 7.0 or higher + - NVIDIA Driver: 535 or higher (recommended 570) + - Python: - - Supported Python versions: - - Python 3.10 - - Use a virtual environment (venv, virtualenv, or conda) to isolate dependencies + - 3.10 + - 3.11 Notes: -- Ensure your NVIDIA driver and CUDA toolkit versions are compatible with your chosen PyTorch wheel. +- Ensure your NVIDIA driver is compatible with your chosen PyTorch wheel. - We recommend installing PyTorch from the official PyTorch instructions for your CUDA version: https://pytorch.org/get-started/locally/ --- diff --git a/docs/source/resources/roadmap.md b/docs/source/resources/roadmap.md index f248a8f..4999e8b 100644 --- a/docs/source/resources/roadmap.md +++ b/docs/source/resources/roadmap.md @@ -20,7 +20,6 @@ Currently, EmbodiChain is under active development. Our plan for the feature roa - Motion Generation: - Add more advanced motion generation methods and examples. - Useful Tools: - - Add a robot workspace analysis tool for better visualization and sampling of robot accessible workspace. - We are working on USD support for EmbodiChain to enable better scene creation and asset management. - We will release a simple Real2Sim pipeline, which enables automatic task generation from real-world data. - Robots Integration: diff --git a/docs/source/resources/robot/cobotmagic.md b/docs/source/resources/robot/cobotmagic.md index 63eb6b6..de23dd2 100644 --- a/docs/source/resources/robot/cobotmagic.md +++ b/docs/source/resources/robot/cobotmagic.md @@ -39,9 +39,8 @@ CobotMagic is a versatile dual-arm collaborative robot developed by AgileX Robot from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import CobotMagicCfg -config = SimulationManagerCfg(headless=False, sim_device="cpu") +config = SimulationManagerCfg(headless=False, sim_device="cpu", num_envs=2) sim = SimulationManager(config) -sim.build_multiple_arenas(2) # Supports parallel simulation in multiple arenas sim.set_manual_update(False) robot = sim.add_robot(cfg=CobotMagicCfg().from_dict({})) diff --git a/docs/source/tutorial/create_scene.rst b/docs/source/tutorial/create_scene.rst index 4f8dd31..244bd93 100644 --- a/docs/source/tutorial/create_scene.rst +++ b/docs/source/tutorial/create_scene.rst @@ -31,15 +31,13 @@ Command-line arguments are parsed using ``argparse`` to allow for easy customiza .. literalinclude:: ../../../scripts/tutorials/sim/create_scene.py :language: python :start-at: # Parse command line arguments - :end-at: sim.build_multiple_arenas(args.num_envs, space=3.0) + :end-at: sim = SimulationManager(sim_cfg) There are two kinds of physics mode in :class:`SimulationManager`: - `manual`: The physics updates only when the user calls the :meth:`SimulationManager.update` function. This mode is used for robot learning tasks where precise control over simulation steps is required. Enabled by setting :meth:`SimulationManager.set_manual_update` to True. - `auto`: The physics updates in a standalone thread, which enable asynchronous rendering and physics stepping. This mode is suitable for visualizations and demos for digital twins applications. This is the default mode. -If `num_envs` is greater than 1, :meth:`SimulationManager.build_multiple_arenas` should be used to create multiple simulation arenas. - Adding objects to the scene --------------------------- diff --git a/docs/source/tutorial/create_softbody.rst b/docs/source/tutorial/create_softbody.rst index cfaf85c..d6fd5de 100644 --- a/docs/source/tutorial/create_softbody.rst +++ b/docs/source/tutorial/create_softbody.rst @@ -32,8 +32,6 @@ The first step is to configure the simulation environment. This is done using th :start-at: # Configure the simulation :end-at: print("[INFO]: Scene setup complete!") -If ``num_envs`` is greater than 1, :meth:`SimulationManager.build_multiple_arenas` should be used to create multiple simulation arenas. - Adding a soft body to the scene ------------------------------- diff --git a/docs/source/tutorial/rl.rst b/docs/source/tutorial/rl.rst index d3c1888..7be8c73 100644 --- a/docs/source/tutorial/rl.rst +++ b/docs/source/tutorial/rl.rst @@ -182,7 +182,7 @@ To start training, run: .. code-block:: bash - python embodichain/agents/rl/train.py --config configs/agents/rl/push_cube/train_config.json + python -m embodichain.agents.rl.train --config configs/agents/rl/push_cube/train_config.json Outputs ------- diff --git a/embodichain/data/handler/lerobot_data_handler.py b/embodichain/data/handler/lerobot_data_handler.py index fa04152..493c7a3 100644 --- a/embodichain/data/handler/lerobot_data_handler.py +++ b/embodichain/data/handler/lerobot_data_handler.py @@ -72,17 +72,8 @@ def _build_lerobot_features(self, use_videos: bool = True) -> Dict: } # Add state features (proprio) - state_dim = 0 - for proprio_name in SUPPORTED_PROPRIO_TYPES: - robot = self.env.robot - part = data_key_to_control_part( - robot=robot, - control_parts=robot_meta_config.get("control_parts", []), - data_key=proprio_name, - ) - if part: - indices = robot.get_joint_ids(part, remove_mimic=True) - state_dim += len(indices) + qpos = self.env.robot.get_qpos() + state_dim = qpos.shape[1] if state_dim > 0: features["observation.state"] = { @@ -123,13 +114,15 @@ def _convert_frame_to_lerobot( # Determine batch size from qpos qpos = obs["robot"][JointType.QPOS.value] num_envs = qpos.shape[0] - + frames = [] - + + all_qpos = robot.get_qpos() + # Process each environment for env_idx in range(num_envs): frame = {"task": task} - + # Add images for camera_name in extra_vision_config.keys(): if camera_name in obs["sensor"]: @@ -152,9 +145,13 @@ def _convert_frame_to_lerobot( if is_stereo: color_right_data = obs["sensor"][camera_name]["color_right"] if isinstance(color_right_data, torch.Tensor): - color_right_img = color_right_data[env_idx][:, :, :3].cpu().numpy() + color_right_img = ( + color_right_data[env_idx][:, :, :3].cpu().numpy() + ) else: - color_right_img = np.array(color_right_data)[env_idx][:, :, :3] + color_right_img = np.array(color_right_data)[env_idx][ + :, :, :3 + ] # Ensure uint8 format if ( @@ -166,23 +163,10 @@ def _convert_frame_to_lerobot( frame[get_right_name(camera_name)] = color_right_img # Add state (proprio) - state_list = [] - for proprio_name in SUPPORTED_PROPRIO_TYPES: - part = data_key_to_control_part( - robot=robot, - control_parts=robot_meta_config.get("control_parts", []), - data_key=proprio_name, - ) - if part: - indices = robot.get_joint_ids(part, remove_mimic=True) - qpos_data = qpos[env_idx][indices].cpu().numpy() - qpos_data = HandQposNormalizer.normalize_hand_qpos( - qpos_data, part, robot=robot - ) - state_list.append(qpos_data) - - if state_list: - frame["observation.state"] = np.concatenate(state_list) + # robot.get_qpos() returns shape (num_envs, num_joints) + state = all_qpos[env_idx].cpu().numpy().astype(np.float32) + + frame["observation.state"] = state # Add actions # Handle different action types @@ -202,7 +186,7 @@ def _convert_frame_to_lerobot( action_data = np.array(action)[env_idx, :arm_dofs] frame["action"] = action_data - + frames.append(frame) return frames diff --git a/embodichain/lab/gym/envs/__init__.py b/embodichain/lab/gym/envs/__init__.py index 286d782..e3d7583 100644 --- a/embodichain/lab/gym/envs/__init__.py +++ b/embodichain/lab/gym/envs/__init__.py @@ -29,3 +29,7 @@ # Reinforcement learning environments from embodichain.lab.gym.envs.tasks.rl.push_cube import PushCubeEnv + +from embodichain.lab.gym.envs.tasks.special.parallel import ParallelRolloutEnv +from embodichain.lab.gym.envs.tasks.special.simple_task_ur10_1 import SimpleTaskEnv_1 +from embodichain.lab.gym.envs.tasks.special.simple_task_ur10_2 import SimpleTaskEnv_2 diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index 244a94a..1b2e5b5 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -109,6 +109,7 @@ def __init__( self.sim_cfg = SimulationManagerCfg(headless=True) else: self.sim_cfg = self.cfg.sim_cfg + self.sim_cfg.num_envs = self.num_envs if self.cfg.seed is not None: self.cfg.seed = set_seed(self.cfg.seed) @@ -208,8 +209,6 @@ def _setup_scene(self, **kwargs): logger.log_info( f"Initializing {self.num_envs} environments on {self.sim_cfg.sim_device}." ) - if self.num_envs > 1: - self.sim.build_multiple_arenas(self.num_envs) self.robot = self._setup_robot(**kwargs) if self.robot is None: diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 2548557..6a869d8 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -156,6 +156,9 @@ def __init__(self, cfg: EmbodiedEnvCfg, **kwargs): "action_scale": 0.1, } + # Initialize total time tracker for dataset recording + self.total_time = 0.0 + for name, default in defaults.items(): value = extensions.get(name, getattr(cfg, name, default)) setattr(cfg, name, value) @@ -208,7 +211,7 @@ def _apply_functor_filter(self) -> None: from embodichain.lab.gym.envs.managers.cfg import EventCfg functors_to_remove = get_all_exported_items_from_module( - "embodichain.lab.gym.envs.managers.randomization.rendering" + "embodichain.lab.gym.envs.managers.randomization.visual" ) if self.cfg.filter_visual_rand and self.cfg.events: # Iterate through all attributes of the events object @@ -620,11 +623,25 @@ def _to_lerobot_dataset(self) -> str | None: ) action_list = self.episode_action_list - logger.log_info(f"Saving {self.num_envs} episodes with {len(obs_list)} frames each...") + logger.log_info( + f"Saving {self.num_envs} episodes with {len(obs_list)} frames each..." + ) # Get task instruction task = self.metadata["dataset"]["instruction"].get("lang", "unknown_task") + # Calculate total time for the entire dataset (all environments) + fps = self.dataset.meta.info.get("fps", 30) + current_episode_time = (len(obs_list) * self.num_envs) / fps if fps > 0 else 0 + + # Prepare extra info (same for all episodes) + extra_info = self.cfg.dataset.get("extra", {}) + episode_extra_info = extra_info.copy() + self.total_time += current_episode_time + episode_extra_info["total_time"] = self.total_time + episode_extra_info["data_type"] = "sim" + self.update_dataset_info({"extra": episode_extra_info}) + # Process each environment as a separate episode for env_idx in range(self.num_envs): # Add frames for this specific environment @@ -633,19 +650,9 @@ def _to_lerobot_dataset(self) -> str | None: # Only add the frame for this specific environment self.dataset.add_frame(frames[env_idx]) - # Save episode for this environment - extra_info = self.cfg.dataset.get("extra", {}) - fps = self.dataset.meta.info.get("fps", 30) - total_time = len(obs_list) / fps if fps > 0 else 0 - - episode_extra_info = extra_info.copy() - episode_extra_info["total_time"] = total_time - episode_extra_info["data_type"] = "sim" - episode_extra_info["env_index"] = env_idx - - self.update_dataset_info({"extra": episode_extra_info}) + # Save episode for this environment (without triggering batch encoding yet) self.dataset.save_episode() - + logger.log_info( f"Saved episode {self.curr_episode} for environment {env_idx} with {len(obs_list)} frames" ) diff --git a/embodichain/lab/gym/envs/managers/events.py b/embodichain/lab/gym/envs/managers/events.py index 0b65d49..d305232 100644 --- a/embodichain/lab/gym/envs/managers/events.py +++ b/embodichain/lab/gym/envs/managers/events.py @@ -601,6 +601,7 @@ def drop_rigid_object_group_sequentially( lower=range_low, upper=range_high, size=(num_instance, 3), + device=env.device, ) drop_pose_i = drop_pose.unsqueeze(1) drop_pose_i[:, 0, :3] = drop_pos + random_offset diff --git a/embodichain/lab/gym/envs/managers/randomization/__init__.py b/embodichain/lab/gym/envs/managers/randomization/__init__.py index 6483181..7a96537 100644 --- a/embodichain/lab/gym/envs/managers/randomization/__init__.py +++ b/embodichain/lab/gym/envs/managers/randomization/__init__.py @@ -14,8 +14,9 @@ # limitations under the License. # ---------------------------------------------------------------------------- -from .rendering import * -from .spatial import * +from .physics import * # noqa: F401, F403 +from .visual import * # noqa: F401, F403 +from .spatial import * # noqa: F401, F403 """ Randomization are all implemented as Event functors. diff --git a/embodichain/lab/gym/envs/managers/randomization/physics.py b/embodichain/lab/gym/envs/managers/randomization/physics.py new file mode 100644 index 0000000..227063f --- /dev/null +++ b/embodichain/lab/gym/envs/managers/randomization/physics.py @@ -0,0 +1,63 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING, Union, List + +from embodichain.lab.sim.objects import RigidObject, Robot +from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg +from embodichain.utils.math import sample_uniform + + +if TYPE_CHECKING: + from embodichain.lab.gym.envs import EmbodiedEnv + + +def randomize_rigid_object_mass( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, List[int]], + entity_cfg: SceneEntityCfg, + mass_range: tuple[float, float], + relative: bool = False, +) -> None: + """Randomize the mass of rigid objects in the environment. + + Args: + env (EmbodiedEnv): The environment instance. + env_ids (Union[torch.Tensor, List[int]]): The environment IDs to apply the randomization. + entity_cfg (SceneEntityCfg): The configuration for the scene entity. + mass_range (tuple[float, float]): The range (min, max) to sample the mass from. + relative (bool): Whether to apply the mass change relative to the initial mass. Defaults to False. + """ + + if entity_cfg.uid not in env.sim.get_rigid_object_uid_list(): + return + + rigid_object: RigidObject = env.sim.get_rigid_object(entity_cfg.uid) + num_instance = len(env_ids) + + sampled_masses = sample_uniform( + lower=mass_range[0], upper=mass_range[1], size=(num_instance,) + ) + + if relative: + init_mass = rigid_object.cfg.attrs.mass + init_mass = torch.full((sampled_masses.shape), init_mass, device=env.device) + sampled_masses = init_mass + sampled_masses + + rigid_object.set_mass(sampled_masses, env_ids=env_ids) diff --git a/embodichain/lab/gym/envs/managers/randomization/spatial.py b/embodichain/lab/gym/envs/managers/randomization/spatial.py index a63c5fa..0776f91 100644 --- a/embodichain/lab/gym/envs/managers/randomization/spatial.py +++ b/embodichain/lab/gym/envs/managers/randomization/spatial.py @@ -70,6 +70,7 @@ def get_random_pose( lower=pos_low, upper=pos_high, size=(num_instance, 3), + device=init_pos.device, ) if relative_position: random_value += init_pos @@ -86,6 +87,7 @@ def get_random_pose( lower=rot_low, upper=rot_high, size=(num_instance, 3), + device=init_pos.device, ) * torch.pi / 180.0 @@ -252,6 +254,7 @@ def randomize_robot_qpos( lower=torch.tensor(qpos_range[0], device=env.device), upper=torch.tensor(qpos_range[1], device=env.device), size=(num_instance, len(joint_ids)), + device=env.device, ) if relative_qpos: diff --git a/embodichain/lab/gym/envs/managers/randomization/rendering.py b/embodichain/lab/gym/envs/managers/randomization/visual.py similarity index 96% rename from embodichain/lab/gym/envs/managers/randomization/rendering.py rename to embodichain/lab/gym/envs/managers/randomization/visual.py index 535ebd6..fcddda7 100644 --- a/embodichain/lab/gym/envs/managers/randomization/rendering.py +++ b/embodichain/lab/gym/envs/managers/randomization/visual.py @@ -43,14 +43,6 @@ from embodichain.lab.gym.envs import EmbodiedEnv -__all__ = [ - "randomize_camera_extrinsics", - "randomize_light", - "randomize_camera_intrinsics", - "randomize_visual_material", -] - - def randomize_camera_extrinsics( env: EmbodiedEnv, env_ids: Union[torch.Tensor, None], @@ -448,24 +440,6 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv): self.entity_cfg.link_names = link_names self.entity.set_visual_material(mat, link_names=link_names) - @staticmethod - def gen_random_base_color_texture(width: int, height: int) -> torch.Tensor: - """Generate a random base color texture. - - Args: - width: The width of the texture. - height: The height of the texture. - - Returns: - A torch tensor representing the random base color texture with shape (height, width, 4). - """ - # Generate random RGB values - rgb = torch.ones((height, width, 3), dtype=torch.float32) - rgb *= torch.rand((1, 1, 3), dtype=torch.float32) - rgba = torch.cat((rgb, torch.ones((height, width, 1))), dim=2) - rgba = (rgba * 255).to(torch.uint8) - return rgba - def _randomize_texture(self, mat_inst: VisualMaterialInst) -> None: if len(self.textures) > 0: # Randomly select a texture from the preloaded textures diff --git a/embodichain/lab/gym/envs/tasks/special/__init__.py b/embodichain/lab/gym/envs/tasks/special/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/embodichain/lab/gym/envs/tasks/special/parallel.py b/embodichain/lab/gym/envs/tasks/special/parallel.py new file mode 100644 index 0000000..941180b --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/special/parallel.py @@ -0,0 +1,57 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["ParallelRolloutEnv"] + + +@register_env("ParallelRollout-v1", max_episode_steps=600) +class ParallelRolloutEnv(EmbodiedEnv): + """A demo environment + + Args: + EmbodiedEnv (_type_): _description_ + """ + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + def create_demo_action_list(self, *args, **kwargs): + """ + Create a demonstration action list for the current task. + + Returns: + list: A list of demo actions generated by the task. + """ + action_list = [] + for i in range(100): + action = self.action_space.sample() + action = torch.as_tensor(action, dtype=torch.float32, device=self.device) + + init_pose = self.robot.get_qpos() + action = ( + init_pose + + torch.rand_like(action, dtype=torch.float32, device=self.device) * 0.2 + - 0.1 + ) + action_list.append(action) + + return action_list diff --git a/embodichain/lab/gym/envs/tasks/special/simple_task_ur10_1.py b/embodichain/lab/gym/envs/tasks/special/simple_task_ur10_1.py new file mode 100644 index 0000000..f9c55ef --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/special/simple_task_ur10_1.py @@ -0,0 +1,87 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import math +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["SimpleTaskEnv_1"] + + +@register_env("simple_task_1", max_episode_steps=600) +class SimpleTaskEnv_1(EmbodiedEnv): + """A demo environment with sinusoidal trajectory + + Args: + EmbodiedEnv (_type_): _description_ + """ + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + def create_demo_action_list(self, *args, **kwargs): + """ + Create a demonstration action list for the current task. + + This demo creates a simple sinusoidal trajectory for the robot joints. + + Returns: + list: A list of demo actions generated by the task. + """ + action_list = [] + num_steps = 100 + + # Get initial pose + init_pose = self.robot.get_qpos() # shape: (num_envs, num_joints) + + # Create a sinusoidal trajectory + for i in range(num_steps): + # Calculate phase for sinusoidal motion + t = i / num_steps # 0 to 1 + phase = torch.full( + (init_pose.shape[0],), t * 2 * torch.pi, device=self.device + ) # repeat for num_envs + + # Create sinusoidal offsets for each joint + # Joint 0: horizontal movement + # Joint 1: vertical movement + # Other joints: smaller oscillations + offset = torch.zeros_like( + init_pose, dtype=torch.float32, device=self.device + ) + offset[:, 0] = torch.sin(phase) * 0.3 # ±0.3 rad + offset[:, 1] = torch.cos(phase) * 0.2 # ±0.2 rad + offset[:, 2] = torch.sin(phase * 2) * 0.1 # ±0.1 rad, double frequency + + # Add small random variation to make it more natural + noise = (torch.rand_like(init_pose, device=self.device) - 0.5) * 0.02 + + # Compute action + action = init_pose + offset + noise + + # Clamp to joint limits if available + if hasattr(self.robot.body_data, "qpos_limits"): + qpos_limits = self.robot.body_data.qpos_limits[0] # (num_joints, 2) + action = torch.clamp(action, qpos_limits[:, 0], qpos_limits[:, 1]) + + action_list.append(action) + + logger.log_info( + f"Generated {len(action_list)} demo actions with sinusoidal trajectory" + ) + return action_list diff --git a/embodichain/lab/gym/envs/tasks/special/simple_task_ur10_2.py b/embodichain/lab/gym/envs/tasks/special/simple_task_ur10_2.py new file mode 100644 index 0000000..680291f --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/special/simple_task_ur10_2.py @@ -0,0 +1,97 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["SimpleTaskEnv_2"] + + +@register_env("simple_task_2", max_episode_steps=600) +class SimpleTaskEnv_2(EmbodiedEnv): + """A demo environment with spiral trajectory + + Args: + EmbodiedEnv (_type_): _description_ + """ + + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + def create_demo_action_list(self, *args, **kwargs): + """ + Create a demonstration action list for the current task. + + This demo creates a spiral upward trajectory for the robot joints. + + Returns: + list: A list of demo actions generated by the task. + """ + action_list = [] + num_steps = 100 + + # Get initial pose + init_pose = self.robot.get_qpos() # shape: (num_envs, num_joints) + + # Create a spiral upward trajectory + for i in range(num_steps): + # Calculate phase for spiral motion + t = i / num_steps # 0 to 1 + phase = torch.full( + (init_pose.shape[0],), t * 4 * 3.14159, device=self.device + ) # repeat for num_envs + + # Create spiral offsets + # Radius decreases as it goes up + radius = 0.4 * (1 - t * 0.5) # Start at 0.4, end at 0.2 + + offset = torch.zeros_like( + init_pose, dtype=torch.float32, device=self.device + ) + offset[:, 0] = torch.cos(phase) * radius # X-axis circular motion + offset[:, 1] = torch.sin(phase) * radius # Y-axis circular motion + offset[:, 2] = t * 0.5 # Z-axis upward motion (0 to 0.5 rad) + + # Add wrist rotation + if init_pose.shape[1] > 3: + offset[:, 3] = phase * 0.3 # Continuous rotation + if init_pose.shape[1] > 4: + offset[:, 4] = ( + torch.sin(phase * 3) * 0.15 + ) # Triple frequency oscillation + if init_pose.shape[1] > 5: + offset[:, 5] = t * 0.2 # Gradual opening + + # Add small random variation + noise = (torch.rand_like(init_pose, device=self.device) - 0.5) * 0.01 + + # Compute action + action = init_pose + offset + noise + + # Clamp to joint limits if available + if hasattr(self.robot.body_data, "qpos_limits"): + qpos_limits = self.robot.body_data.qpos_limits[0] # (num_joints, 2) + action = torch.clamp(action, qpos_limits[:, 0], qpos_limits[:, 1]) + + action_list.append(action) + + logger.log_info( + f"Generated {len(action_list)} demo actions with spiral trajectory" + ) + return action_list diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 3864ede..4ad3b95 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -93,12 +93,12 @@ def generate_function( if not valid: break - if not debug_mode and env.is_task_success().item(): + if not debug_mode: dataset_id = f"time_{time_id}_trajectory_{trajectory_idx}" data_dict = env.to_dataset() ret.append(data_dict) - # TODO: Add data saving and online data streaming logic here. + # TODO: Add data saving and online data streaming logic here. else: log_warning(f"Task fail, Skip to next generation.") diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 1e325e0..cedb59e 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -43,12 +43,25 @@ @configclass class PhysicsCfg: gravity: np.ndarray = field(default_factory=lambda: np.array([0, 0, -9.81])) + """Gravity vector for the simulation environment.""" + bounce_threshold: float = 2.0 + """The speed threshold below which collisions will not produce bounce effects.""" + enable_pcm: bool = True + """Enable persistent contact manifold (PCM) for improved collision handling.""" + enable_tgs: bool = True + """Enable temporal gauss-seidel (TGS) solver for better stability.""" + enable_ccd: bool = False + """Enable continuous collision detection (CCD) for fast-moving objects.""" + enable_enhanced_determinism: bool = False + """Enable enhanced determinism for consistent simulation results.""" + enable_friction_every_iteration: bool = True + """Enable friction calculations at every solver iteration.""" length_tolerance: float = 0.05 """The length tolerance for the simulation. @@ -86,20 +99,28 @@ class MarkerCfg: name: str = "empty-mesh" """Name of the marker for identification purposes.""" + marker_type: Literal["axis", "line", "point"] = "axis" """Type of marker to display. Can be 'axis' (3D coordinate frame), 'line', or 'point'. (only axis supported now)""" + axis_xpos: List[np.ndarray] = None """List of 4x4 transformation matrices defining the position and orientation of each axis marker.""" + axis_size: float = 0.002 """Thickness/size of the axis lines in meters.""" + axis_len: float = 0.005 """Length of each axis arm in meters.""" + line_color: List[float] = [1, 1, 0, 1.0] """RGBA color values for the marker lines. Values should be between 0.0 and 1.0.""" + arrow_type: AxisArrowType = AxisArrowType.CONE """Type of arrow head for axis markers (e.g., CONE, ARROW, etc.).""" + corner_type: AxisCornerType = AxisCornerType.SPHERE """Type of corner/joint visualization for axis markers (e.g., SPHERE, CUBE, etc.).""" + arena_index: int = -1 """Index of the arena where the marker should be placed. -1 means all arenas.""" @@ -110,13 +131,17 @@ class GPUMemoryCfg: temp_buffer_capacity: int = 2**24 """Increase this if you get 'PxgPinnedHostLinearMemoryAllocator: overflowing initial allocation size, increase capacity to at least %.' """ + max_rigid_contact_count: int = 2**19 """Increase this if you get 'Contact buffer overflow detected'""" + max_rigid_patch_count: int = ( 2**18 ) # 81920 is DexSim default but most tasks work with 2**18 """Increase this if you get 'Patch buffer overflow detected'""" + heap_capacity: int = 2**26 + found_lost_pairs_capacity: int = ( 2**25 ) # 262144 is DexSim default but most tasks work with 2**25 @@ -135,29 +160,60 @@ class RigidBodyAttributesCfg: """ mass: float = 1.0 - # set mass to 0 will use density to calculate mass. + """Mass of the rigid body in kilograms. + + Set to 0 will use density to calculate mass. + """ + density: float = 1000.0 + """Density of the rigid body in kg/m^3.""" angular_damping: float = 0.7 + """Angular damping coefficient.""" + linear_damping: float = 0.7 + """Linear damping coefficient.""" + max_depenetration_velocity: float = 10.0 + """Maximum depenetration velocity.""" + sleep_threshold: float = 0.001 + """Threshold below which the body can go to sleep.""" + min_position_iters: int = 4 + """Minimum position iterations.""" + min_velocity_iters: int = 1 + """Minimum velocity iterations.""" max_linear_velocity: float = 1e2 + """Maximum linear velocity.""" + max_angular_velocity: float = 1e2 + """Maximum angular velocity.""" # collision properties. enable_ccd: bool = False + """Enable continuous collision detection (CCD).""" + contact_offset: float = 0.002 + """Contact offset for collision detection.""" + rest_offset: float = 0.001 + """Rest offset for collision detection.""" + enable_collision: bool = True + """Enable collision for the rigid body.""" # physics material properties. restitution: float = 0.0 + """Restitution (bounciness) coefficient.""" + dynamic_friction: float = 0.5 + """Dynamic friction coefficient.""" + static_friction: float = 0.5 + """Static friction coefficient.""" def attr(self) -> PhysicalAttr: """Convert to dexsim PhysicalAttr""" @@ -786,6 +842,9 @@ def add_component( logger.log_error(f"URDF path '{urdf_path}' does not exist.") raise FileNotFoundError(f"URDF path '{urdf_path}' does not exist.") + if transform is None: + transform = np.eye(4) + self.components[component_type] = { "urdf_path": urdf_path, "transform": np.array(transform), @@ -964,12 +1023,9 @@ class RobotCfg(ArticulationCfg): from embodichain.lab.sim.solvers import SolverCfg """Configuration for a robot asset in the simulation. - - # TODO: solver and motion planner may not be configurable inside the robot. - # But currently we put them here and could be moved if necessary. """ - control_parts: Union[Dict[str, List[str]], None] = None + control_parts: Dict[str, List[str]] | None = None """Control parts is the mapping from part name to joint names. For example, {'left_arm': ['joint1', 'joint2'], 'right_arm': ['joint3', 'joint4']} @@ -980,6 +1036,9 @@ class RobotCfg(ArticulationCfg): keys corresponding to the control parts name. - The joint names in the control parts support regular expressions, e.g., 'joint[1-6]'. After initialization of robot, the names will be expanded to a list of full joint names. + - `Robot` is a derived class of `Articulation`, with control parts support. So the `drive_pros` + in `ArticulationCfg` can use control part as key to specify the corresponding joint drive properties, + which will be overridden if these joint names are already specified. """ urdf_cfg: URDFCfg | None = None @@ -1009,6 +1068,8 @@ def from_dict(cls, init_dict: Dict[str, Union[str, float, tuple]]) -> RobotCfg: from embodichain.lab.sim.cfg import URDFCfg setattr(cfg, key, URDFCfg.from_dict(value)) + elif key == "fpath": + setattr(cfg, key, get_data_path(value)) elif is_configclass(attr): setattr( cfg, key, attr.from_dict(value) diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 00a03de..4930c36 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -797,7 +797,6 @@ def set_local_pose( # we should keep `pose_` life cycle to the end of the function. pose_ = torch.cat((quat, xyz), dim=-1) indices = self.body_data.gpu_indices[local_env_ids] - torch.cuda.synchronize(self.device) self._ps.gpu_apply_root_data( data=pose_, gpu_indices=indices, @@ -981,7 +980,6 @@ def set_qpos( indices = self.body_data.gpu_indices[local_env_ids] qpos_set = self.body_data._qpos[local_env_ids] qpos_set[:, local_joint_ids] = qpos - torch.cuda.synchronize(self.device) self._ps.gpu_apply_joint_data( data=qpos_set, gpu_indices=indices, @@ -1045,7 +1043,6 @@ def set_qvel( self.body_data.qvel qvel_set = self.body_data._qvel[local_env_ids] qvel_set[:, joint_ids] = qvel - torch.cuda.synchronize(self.device) self._ps.gpu_apply_joint_data( data=qvel_set, gpu_indices=indices, @@ -1086,7 +1083,6 @@ def set_qf( self.body_data.qf qf_set = self.body_data._qf[local_env_ids] qf_set[:, joint_ids] = qf - torch.cuda.synchronize(self.device) self._ps.gpu_apply_joint_data( data=qf_set, gpu_indices=indices, @@ -1167,13 +1163,11 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None: (len(local_env_ids), self.dof), dtype=torch.float32, device=self.device ) indices = self.body_data.gpu_indices[local_env_ids] - torch.cuda.synchronize(self.device) self._ps.gpu_apply_joint_data( data=zeros, gpu_indices=indices, data_type=ArticulationGPUAPIWriteType.JOINT_VELOCITY, ) - torch.cuda.synchronize(self.device) self._ps.gpu_apply_joint_data( data=zeros, gpu_indices=indices, diff --git a/embodichain/lab/sim/objects/light.py b/embodichain/lab/sim/objects/light.py index 7670b90..1f5b66b 100644 --- a/embodichain/lab/sim/objects/light.py +++ b/embodichain/lab/sim/objects/light.py @@ -1,3 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + import torch import numpy as np from typing import List, Sequence diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index c701236..10548e8 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/embodichain/lab/sim/objects/rigid_object.py @@ -499,6 +499,44 @@ def set_attrs( for i, env_idx in enumerate(local_env_ids): self._entities[env_idx].set_physical_attr(attrs[i].attr()) + def set_mass( + self, mass: torch.Tensor, env_ids: Sequence[int] | None = None + ) -> None: + """Set mass for the rigid object. + + Args: + mass (torch.Tensor): The mass to set with shape (N,). + env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. + """ + local_env_ids = self._all_indices if env_ids is None else env_ids + + if len(local_env_ids) != len(mass): + logger.log_error( + f"Length of env_ids {len(local_env_ids)} does not match mass length {len(mass)}." + ) + + mass = mass.cpu().numpy() + for i, env_idx in enumerate(local_env_ids): + self._entities[env_idx].get_physical_body().set_mass(mass[i]) + + def get_mass(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: + """Get mass for the rigid object. + + Args: + env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. + + Returns: + torch.Tensor: The mass of the rigid object with shape (N,). + """ + local_env_ids = self._all_indices if env_ids is None else env_ids + + masses = [] + for _, env_idx in enumerate(local_env_ids): + mass = self._entities[env_idx].get_physical_body().get_mass() + masses.append(mass) + + return torch.as_tensor(masses, dtype=torch.float32, device=self.device) + def set_visual_material( self, mat: VisualMaterial, env_ids: Sequence[int] | None = None ) -> None: diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index 9edb982..33ce435 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -67,9 +67,9 @@ def __init__( entities: List[_Articulation], device: torch.device = torch.device("cpu"), ) -> None: - super().__init__(cfg, entities, device) - self._solvers = {} + self._entities = entities + self.cfg = cfg # Initialize joint ids for control parts. self._joint_ids: Dict[str, List[int]] = {} @@ -79,6 +79,10 @@ def __init__( if self.cfg.control_parts: self._init_control_parts(self.cfg.control_parts) + super().__init__(cfg, entities, device) + + self._solvers = {} + if self.cfg.solver_cfg: self.init_solver(self.cfg.solver_cfg) @@ -182,7 +186,7 @@ def get_qvel_limits( env_ids (Sequence[int] | None): The environment ids to get the qvel limits for. If None, all environments are used. Returns: - torch.Tensor: Joint velocity limits with shape (N, dof, 2), where N is the number of environments. + torch.Tensor: Joint velocity limits with shape (N, dof), where N is the number of environments. """ local_env_ids = self._all_indices if env_ids is None else env_ids @@ -195,7 +199,7 @@ def get_qvel_limits( f"The control part '{name}' does not exist in the robot's control parts." ) part_joint_ids = self.get_joint_ids(name=name) - return qvel_limits[local_env_ids][:, part_joint_ids, :] + return qvel_limits[local_env_ids][:, part_joint_ids] def get_qf_limits( self, name: str | None = None, env_ids: Sequence[int] | None = None @@ -209,7 +213,7 @@ def get_qf_limits( env_ids (Sequence[int] | None): The environment ids to get the qf limits for. If None, all environments are used. Returns: - torch.Tensor: Joint effort limits with shape (N, dof, 2), where N is the number of environments. + torch.Tensor: Joint effort limits with shape (N, dof), where N is the number of environments. """ local_env_ids = self._all_indices if env_ids is None else env_ids @@ -222,7 +226,7 @@ def get_qf_limits( f"The control part '{name}' does not exist in the robot's control parts." ) part_joint_ids = self.get_joint_ids(name=name) - return qf_limits[local_env_ids][:, part_joint_ids, :] + return qf_limits[local_env_ids][:, part_joint_ids] def get_proprioception(self) -> Dict[str, torch.Tensor]: """Gets robot proprioception information, primarily for agent state representation in robot learning scenarios. @@ -763,7 +767,10 @@ def _init_control_parts(self, control_parts: Dict[str, List[str]]) -> None: control_parts (Dict[str, List[str]]): A dictionary where keys are control part names and values are lists of joint names or regular expressions that match joint names. """ - joint_name_to_ids = {name: i for i, name in enumerate(self.joint_names)} + joint_name_to_ids = { + name: i + for i, name in enumerate(self._entities[0].get_actived_joint_names()) + } for name, joint_names in control_parts.items(): # convert joint_names which is a regular expression to a list of joint names joint_names_expanded = [] @@ -790,6 +797,74 @@ def _init_control_parts(self, control_parts: Dict[str, List[str]]) -> None: # Initialize control groups self._control_groups = self._extract_control_groups() + def _set_default_joint_drive(self) -> None: + """Set default joint drive parameters based on the configuration.""" + import numbers + from embodichain.utils.string import resolve_matching_names_values + + drive_props = [ + ("damping", self.default_joint_damping), + ("stiffness", self.default_joint_stiffness), + ("max_effort", self.default_joint_max_effort), + ("max_velocity", self.default_joint_max_velocity), + ("friction", self.default_joint_friction), + ] + + for prop_name, default_array in drive_props: + value = getattr(self.cfg.drive_pros, prop_name, None) + if value is None: + continue + if isinstance(value, numbers.Number): + default_array[:] = value + else: + try: + control_part_dict = {} + value_copy = value.copy() + if self.control_parts: + # Extract control part and map the corresponding joint names + for key in value.keys(): + if key in self.control_parts: + control_part_dict[key] = value_copy.pop(key) + + indices, _, values = resolve_matching_names_values( + value_copy, self.joint_names + ) + + if self.control_parts: + # Add control part joints to indices and values + for part_name, part_value in control_part_dict.items(): + part_joint_names = self.control_parts[part_name] + part_indices, _, part_values = ( + resolve_matching_names_values( + {jn: part_value for jn in part_joint_names}, + self.joint_names, + ) + ) + indices.extend(part_indices) + values.extend(part_values) + + default_array[:, indices] = torch.as_tensor( + values, dtype=torch.float32, device=self.device + ) + except Exception as e: + logger.log_error(f"Failed to set {prop_name}: {e}") + + drive_pros = self.cfg.drive_pros + if isinstance(drive_pros, dict): + drive_type = drive_pros.get("drive_type", None) + else: + drive_type = getattr(drive_pros, "drive_type", None) + + # Apply drive parameters to all articulations in the batch + self.set_drive( + stiffness=self.default_joint_stiffness, + damping=self.default_joint_damping, + max_effort=self.default_joint_max_effort, + max_velocity=self.default_joint_max_velocity, + friction=self.default_joint_friction, + drive_type=drive_type, + ) + def init_solver(self, cfg: Union[SolverCfg, Dict[str, SolverCfg]]) -> None: """Initialize the kinematic solver for the robot. diff --git a/embodichain/lab/sim/robots/cobotmagic.py b/embodichain/lab/sim/robots/cobotmagic.py index ec39331..25ab38c 100644 --- a/embodichain/lab/sim/robots/cobotmagic.py +++ b/embodichain/lab/sim/robots/cobotmagic.py @@ -28,6 +28,7 @@ RigidBodyAttributesCfg, ) from embodichain.lab.sim.solvers import SolverCfg, OPWSolverCfg +from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg from embodichain.data import get_data_path from embodichain.utils import configclass from embodichain.utils import logger @@ -41,36 +42,13 @@ class CobotMagicCfg(RobotCfg): @classmethod def from_dict(cls, init_dict: Dict[str, Union[str, float, int]]) -> CobotMagicCfg: - from embodichain.lab.sim.solvers import merge_solver_cfg cfg = cls() default_cfgs = cls()._build_default_cfgs() for key, value in default_cfgs.items(): setattr(cfg, key, value) - robot_cfg = RobotCfg.from_dict(init_dict) - - # set attrs into cfg from the robot_cfg - for key, value in init_dict.items(): - if key == "solver_cfg": - # merge provided solver_cfg values into default solver config - provided_solver_cfg = init_dict.get("solver_cfg") - if provided_solver_cfg: - for part, item in provided_solver_cfg.items(): - if "class_type" in provided_solver_cfg[part]: - cfg.solver_cfg[part] = robot_cfg.solver_cfg[part] - else: - try: - merged = merge_solver_cfg( - cfg.solver_cfg, provided_solver_cfg - ) - cfg.solver_cfg = merged - except Exception: - logger.log_error( - f"Failed to merge solver_cfg, using provided config outright." - ) - else: - setattr(cfg, key, getattr(robot_cfg, key)) + cfg = merge_robot_cfg(cfg, init_dict) return cfg @@ -207,9 +185,8 @@ def build_pk_serial_chain( torch.set_printoptions(precision=5, sci_mode=False) - config = SimulationManagerCfg(headless=False, sim_device="cuda") + config = SimulationManagerCfg(headless=False, sim_device="cuda", num_envs=2) sim = SimulationManager(config) - sim.build_multiple_arenas(2) config = { "init_pos": [0.0, 0.0, 1.0], diff --git a/embodichain/lab/sim/robots/dexforce_w1/cfg.py b/embodichain/lab/sim/robots/dexforce_w1/cfg.py index 673eb60..c5588b2 100644 --- a/embodichain/lab/sim/robots/dexforce_w1/cfg.py +++ b/embodichain/lab/sim/robots/dexforce_w1/cfg.py @@ -39,6 +39,7 @@ JointDrivePropertiesCfg, RigidBodyAttributesCfg, ) +from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg from embodichain.data import get_data_path from embodichain.utils import configclass, logger @@ -49,31 +50,36 @@ class DexforceW1Cfg(RobotCfg): version: DexforceW1Version = DexforceW1Version.V021 arm_kind: DexforceW1ArmKind = DexforceW1ArmKind.INDUSTRIAL + with_default_eef: bool = True @classmethod def from_dict( - cls, init_dict: typing.Dict[str, typing.Union[str, float, tuple]] + cls, init_dict: Dict[str, str | float | tuple | dict] ) -> DexforceW1Cfg: """Initialize DexforceW1Cfg from a dictionary. Args: - init_dict (Dict[str, Union[str, float, tuple]]): Dictionary of configuration parameters. + init_dict (Dict[str, str | float | tuple | dict): Dictionary of configuration parameters. Returns: DexforceW1Cfg: An instance of DexforceW1Cfg with parameters set. """ - from embodichain.lab.sim.solvers import merge_solver_cfg init_dict_m = init_dict.copy() version = init_dict_m.get("version", "v021") arm_kind = init_dict_m.get("arm_kind", "anthropomorphic") + with_default_eef = init_dict_m.get("with_default_eef", True) init_dict_m.pop("version", None) init_dict_m.pop("arm_kind", None) + init_dict_m.pop("with_default_eef", None) + cfg: DexforceW1Cfg = cls()._build_default_cfg( - version=version, arm_kind=arm_kind + version=version, arm_kind=arm_kind, with_default_eef=with_default_eef ) - default_physics_cfgs = cls()._build_default_physics_cfgs(arm_kind=arm_kind) + default_physics_cfgs = cls()._build_default_physics_cfgs( + arm_kind=arm_kind, with_default_eef=with_default_eef + ) for key, value in default_physics_cfgs.items(): setattr(cfg, key, value) @@ -82,37 +88,12 @@ def from_dict( ) cfg.solver_cfg = default_solver_cfg - # override default values with those provided in init_dict. - robot_cfg = RobotCfg.from_dict(init_dict_m) - - # set attrs into cfg from the robot_cfg, but merge solver_cfg specially - for key, value in init_dict_m.items(): - if key == "solver_cfg": - # merge provided solver_cfg values into default solver config - provided_solver_cfg = init_dict_m.get("solver_cfg") - if provided_solver_cfg: - for part, item in provided_solver_cfg.items(): - if "class_type" in provided_solver_cfg[part]: - cfg.solver_cfg[part] = robot_cfg.solver_cfg[part] - else: - try: - merged = merge_solver_cfg( - cfg.solver_cfg, provided_solver_cfg - ) - cfg.solver_cfg = merged - except Exception: - logger.log_error( - f"Failed to merge solver_cfg, using provided config outright." - ) - else: - setattr(cfg, key, getattr(robot_cfg, key)) + cfg = merge_robot_cfg(cfg, init_dict_m) return cfg @staticmethod def _build_default_solver_cfg(is_industrial: bool) -> SolverCfg: - # TODO: maybe change default solver for DexforceW1 - from embodichain.lab.sim.solvers import PytorchSolverCfg from embodichain.lab.sim.solvers import SRSSolverCfg from embodichain.lab.sim.robots.dexforce_w1.params import ( W1ArmKineParams, @@ -129,6 +110,22 @@ def _build_default_solver_cfg(is_industrial: bool) -> SolverCfg: arm_kind=DexforceW1ArmKind.INDUSTRIAL, version=DexforceW1Version.V021, ) + left_arm_tcp = np.array( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.15], + [0.0, 0.0, 0.0, 1.0], + ] + ) + right_arm_tcp = np.array( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.15], + [0.0, 0.0, 0.0, 1.0], + ] + ) else: w1_left_arm_params = W1ArmKineParams( arm_side=DexforceW1ArmSide.LEFT, @@ -140,6 +137,22 @@ def _build_default_solver_cfg(is_industrial: bool) -> SolverCfg: arm_kind=DexforceW1ArmKind.ANTHROPOMORPHIC, version=DexforceW1Version.V021, ) + left_arm_tcp = np.array( + [ + [-1.0, 0.0, 0.0, 0.012], + [0.0, 0.0, 1.0, 0.0675], + [0.0, 1.0, 0.0, 0.127], + [0.0, 0.0, 0.0, 1.0], + ] + ) + right_arm_tcp = np.array( + [ + [1.0, 0.0, 0.0, 0.012], + [0.0, 0.0, -1.0, -0.0675], + [0.0, 1.0, 0.0, 0.127], + [0.0, 0.0, 0.0, 1.0], + ] + ) return { "right_arm": SRSSolverCfg( @@ -151,14 +164,7 @@ def _build_default_solver_cfg(is_industrial: bool) -> SolverCfg: T_b_ob=w1_right_arm_params.T_b_ob, link_lengths=w1_right_arm_params.link_lengths, rotation_directions=w1_right_arm_params.rotation_directions, - tcp=np.array( - [ - [1.0, 0.0, 0.0, 0.012], - [0.0, 0.0, -1.0, -0.0675], - [0.0, 1.0, 0.0, 0.127], - [0.0, 0.0, 0.0, 1.0], - ] - ), + tcp=right_arm_tcp, ), "left_arm": SRSSolverCfg( end_link_name="left_ee", @@ -169,65 +175,78 @@ def _build_default_solver_cfg(is_industrial: bool) -> SolverCfg: T_b_ob=w1_left_arm_params.T_b_ob, link_lengths=w1_left_arm_params.link_lengths, rotation_directions=w1_left_arm_params.rotation_directions, - tcp=np.array( - [ - [-1.0, 0.0, 0.0, 0.012], - [0.0, 0.0, 1.0, 0.0675], - [0.0, 1.0, 0.0, 0.127], - [0.0, 0.0, 0.0, 1.0], - ] - ), + tcp=left_arm_tcp, ), } @staticmethod - def _build_default_physics_cfgs(arm_kind: str) -> typing.Dict[str, typing.Any]: + def _build_default_physics_cfgs( + arm_kind: str, with_default_eef: bool = True + ) -> typing.Dict[str, typing.Any]: """Build default physics configurations for DexforceW1. Args: arm_kind: Type of arm, either "industrial" or "anthropomorphic" + with_default_eef: Whether to include default end-effector configurations Returns: Dictionary containing physics configuration parameters """ - # Define common joint patterns - arm_joints = "(RIGHT|LEFT)_J[0-9]" - body_joints = "(ANKLE|KNEE|BUTTOCK|WAIST)" - - # Hand/gripper pattern differs by arm type - hand_pattern = ( - "(LEFT|RIGHT)_FINGER[1-2]" - if arm_kind == "industrial" - else "(RIGHT|LEFT)_[A-Z|_]+" + # Define default joint drive parameters and corresponding joint name patterns + DEFAULT_EEF_JOINT_DRIVE_PARAMS = { + "stiffness": 1e2, + "damping": 1e1, + "max_effort": 1e3, + } + + DEFAULT_EEF_HAND_JOINT_NAMES = ( + "(LEFT|RIGHT)_HAND_(THUMB[12]|INDEX|MIDDLE|RING|PINKY)" ) + DEFAULT_EEF_GRIPPER_JOINT_NAMES = "(LEFT|RIGHT)_FINGER[1-2]" + + # Define common joint patterns + ARM_JOINTS = "(RIGHT|LEFT)_J[0-9]" + BODY_JOINTS = "(ANKLE|KNEE|BUTTOCK|WAIST)" + # Define physics parameters for different joint types joint_params = { "stiffness": { - arm_joints: 1e4, - hand_pattern: 1e2, - body_joints: 1e7, + ARM_JOINTS: 1e4, + BODY_JOINTS: 1e7, }, "damping": { - arm_joints: 1e3, - hand_pattern: 1e1, - body_joints: 1e4, + ARM_JOINTS: 1e3, + BODY_JOINTS: 1e4, }, "max_effort": { - arm_joints: 1e5, - hand_pattern: 1e3, - body_joints: 1e10, + ARM_JOINTS: 1e5, + BODY_JOINTS: 1e10, }, } drive_pros = JointDrivePropertiesCfg(**joint_params) + if with_default_eef: + eef_joint_names = ( + DEFAULT_EEF_HAND_JOINT_NAMES + if arm_kind == "anthropomorphic" + else DEFAULT_EEF_GRIPPER_JOINT_NAMES + ) + drive_pros.stiffness.update( + {eef_joint_names: DEFAULT_EEF_JOINT_DRIVE_PARAMS["stiffness"]} + ) + drive_pros.damping.update( + {eef_joint_names: DEFAULT_EEF_JOINT_DRIVE_PARAMS["damping"]} + ) + drive_pros.max_effort.update( + {eef_joint_names: DEFAULT_EEF_JOINT_DRIVE_PARAMS["max_effort"]} + ) + return { "min_position_iters": 32, "min_velocity_iters": 8, "drive_pros": drive_pros, - # TODO: we may use the some properties from URDF as default values - # eg. mass, friction, damping, etc. "attrs": RigidBodyAttributesCfg( mass=1.0, static_friction=0.95, @@ -243,7 +262,9 @@ def _build_default_physics_cfgs(arm_kind: str) -> typing.Dict[str, typing.Any]: @staticmethod def _build_default_cfg( - version: str = "v021", arm_kind: str = "anthropomorphic" + version: str = "v021", + arm_kind: str = "anthropomorphic", + with_default_eef: bool = True, ) -> DexforceW1Cfg: if arm_kind == "industrial": @@ -266,9 +287,11 @@ def _build_default_cfg( arm_kind=DexforceW1ArmKind(arm_kind), hand_types=hand_types, hand_versions=hand_versions, + include_hand=with_default_eef, ) cfg.version = DexforceW1Version(version) cfg.arm_kind = DexforceW1ArmKind(arm_kind) + cfg.with_default_eef = with_default_eef return cfg @@ -353,14 +376,9 @@ def build_pk_serial_chain( config = SimulationManagerCfg(headless=True, sim_device="cpu") sim = SimulationManager(config) - sim.build_multiple_arenas(1) cfg = DexforceW1Cfg.from_dict( - { - "uid": "dexforce_w1", - "version": "v021", - "arm_kind": "anthropomorphic", - } + {"uid": "dexforce_w1", "version": "v021", "arm_kind": "anthropomorphic"} ) robot = sim.add_robot(cfg=cfg) diff --git a/embodichain/lab/sim/robots/dexforce_w1/utils.py b/embodichain/lab/sim/robots/dexforce_w1/utils.py index 6340237..cdc28a6 100644 --- a/embodichain/lab/sim/robots/dexforce_w1/utils.py +++ b/embodichain/lab/sim/robots/dexforce_w1/utils.py @@ -646,6 +646,7 @@ def build_dexforce_w1_cfg( include_chassis=include_chassis, include_torso=include_torso, include_head=include_head, + include_hand=include_hand, component_versions=component_versions, ) diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index b41a1d8..83c231c 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -125,6 +125,9 @@ class SimulationManagerCfg: cpu_num: int = 1 """The number of CPU threads to use for the simulation engine.""" + num_envs: int = 1 + """The number of parallel environments (arenas) to simulate.""" + arena_space: float = 5.0 """The distance between each arena when building multiple arenas.""" @@ -243,6 +246,8 @@ def __init__( # Set physics to manual update mode by default. self.set_manual_update(True) + self._build_multiple_arenas(sim_config.num_envs) + @property def num_envs(self) -> int: """Get the number of arenas in the simulation. @@ -500,7 +505,7 @@ def close_window(self) -> None: self._world.close_window() self.is_window_opened = False - def build_multiple_arenas(self, num: int, space: float | None = None) -> None: + def _build_multiple_arenas(self, num: int, space: float | None = None) -> None: """Build multiple arenas in a grid pattern. This interface is used for vectorized simulation. diff --git a/embodichain/lab/sim/solvers/__init__.py b/embodichain/lab/sim/solvers/__init__.py index 6f6123f..051b761 100644 --- a/embodichain/lab/sim/solvers/__init__.py +++ b/embodichain/lab/sim/solvers/__init__.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -from .base_solver import SolverCfg, BaseSolver, merge_solver_cfg +from .base_solver import SolverCfg, BaseSolver from .pytorch_solver import PytorchSolverCfg, PytorchSolver from .pinocchio_solver import PinocchioSolverCfg, PinocchioSolver from .differential_solver import DifferentialSolverCfg, DifferentialSolver diff --git a/embodichain/lab/sim/solvers/base_solver.py b/embodichain/lab/sim/solvers/base_solver.py index aa31334..3817de8 100644 --- a/embodichain/lab/sim/solvers/base_solver.py +++ b/embodichain/lab/sim/solvers/base_solver.py @@ -413,56 +413,3 @@ def get_jacobian( raise ValueError( f"Invalid jac_type '{jac_type}'. Must be 'full', 'trans', or 'rot'." ) - - -def merge_solver_cfg( - default: Dict[str, SolverCfg], provided: Dict[str, Any] -) -> Dict[str, SolverCfg]: - """Merge provided solver configuration into the default solver config. - - Rules: - - For each arm key in provided, if the key exists in default, update fields provided. - - If a provided value is a dict, update attributes on the SolverCfg-like object (or dict) by setting keys. - - Primitive values or arrays/lists replace the target value. - - Unknown keys in provided create new entries in the result. - """ - - result = {} - # copy defaults shallowly - for k, v in default.items(): - result[k] = v - - for k, v in provided.items(): - if k in result: - target = result[k] - # if target has __dict__ or is a dataclass-like, set attrs - if hasattr(target, "__dict__") or isinstance(target, dict): - # if provided is a dict, set/override attributes - if isinstance(v, dict): - for sub_k, sub_v in v.items(): - # try to set attribute if possible, otherwise assign into dict - if hasattr(target, sub_k): - try: - setattr(target, sub_k, sub_v) - except Exception: - # fallback to dict assignment if object doesn't accept - try: - target[sub_k] = sub_v - except Exception: - pass - else: - try: - target[sub_k] = sub_v - except Exception: - setattr(target, sub_k, sub_v) - else: - # non-dict provided value replaces the target entirely - result[k] = v - else: - # target is a primitive, replace - result[k] = v - else: - # new solver entry provided; include as-is - result[k] = v - - return result diff --git a/embodichain/lab/sim/utility/cfg_utils.py b/embodichain/lab/sim/utility/cfg_utils.py new file mode 100644 index 0000000..1dcbe12 --- /dev/null +++ b/embodichain/lab/sim/utility/cfg_utils.py @@ -0,0 +1,170 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from embodichain.lab.sim.cfg import RobotCfg +from embodichain.lab.sim.solvers import SolverCfg +from embodichain.utils import logger + + +def merge_solver_cfg( + default: dict[str, SolverCfg], provided: dict[str, any] +) -> dict[str, SolverCfg]: + """Merge provided solver configuration into the default solver config. + + Rules: + - For each arm key in provided, if the key exists in default, update fields provided. + - If a provided value is a dict, update attributes on the SolverCfg-like object (or dict) by setting keys. + - Primitive values or arrays/lists replace the target value. + - Unknown keys in provided create new entries in the result. + """ + + result = {} + # copy defaults shallowly + for k, v in default.items(): + result[k] = v + + for k, v in provided.items(): + if k in result: + target = result[k] + # if target has __dict__ or is a dataclass-like, set attrs + if hasattr(target, "__dict__") or isinstance(target, dict): + # if provided is a dict, set/override attributes + if isinstance(v, dict): + for sub_k, sub_v in v.items(): + # try to set attribute if possible, otherwise assign into dict + if hasattr(target, sub_k): + try: + setattr(target, sub_k, sub_v) + except Exception: + # fallback to dict assignment if object doesn't accept + try: + target[sub_k] = sub_v + except Exception: + pass + else: + try: + target[sub_k] = sub_v + except Exception: + setattr(target, sub_k, sub_v) + else: + # non-dict provided value replaces the target entirely + result[k] = v + else: + # target is a primitive, replace + result[k] = v + else: + # new solver entry provided; include as-is + result[k] = v + + return result + + +def merge_robot_cfg(base_cfg: RobotCfg, override_cfg_dict: dict[str, any]) -> RobotCfg: + """Merge current robot configuration with overriding values from a dictionary. + + Args: + base_cfg (RobotCfg): The base robot configuration. + override_cfg_dict (dict[str, any]): Dictionary of overriding configuration values. + + Returns: + RobotCfg: The merged robot configuration. + """ + + robot_cfg = RobotCfg.from_dict(override_cfg_dict) + + for key, value in override_cfg_dict.items(): + if key == "solver_cfg": + # merge provided solver_cfg values into default solver config + provided_solver_cfg = override_cfg_dict.get("solver_cfg") + if provided_solver_cfg: + for part, item in provided_solver_cfg.items(): + if "class_type" in provided_solver_cfg[part]: + base_cfg.solver_cfg[part] = robot_cfg.solver_cfg[part] + else: + try: + merged = merge_solver_cfg( + base_cfg.solver_cfg, provided_solver_cfg + ) + base_cfg.solver_cfg = merged + except Exception: + logger.log_error( + f"Failed to merge solver_cfg, using provided config outright." + ) + elif key == "drive_pros": + # merge joint drive properties + user_drive_pros_dict = override_cfg_dict.get("drive_pros") + if isinstance(user_drive_pros_dict, dict): + for prop, val in user_drive_pros_dict.items(): + # Get the current value in cfg (which has defaults) + default_val = getattr(base_cfg.drive_pros, prop, None) + + if isinstance(val, dict) and isinstance(default_val, dict): + # Merge dictionaries + default_val.update(val) + else: + # Overwrite if not both dicts + setattr(base_cfg.drive_pros, prop, val) + else: + logger.log_warning( + "drive_pros should be a dictionary. Skipping drive_pros merge." + ) + elif key == "attrs": + # merge physics attributes + user_attrs_dict = override_cfg_dict.get("attrs") + if isinstance(user_attrs_dict, dict): + for attr_key, attr_val in user_attrs_dict.items(): + setattr(base_cfg.attrs, attr_key, attr_val) + else: + logger.log_warning( + "attrs should be a dictionary. Skipping attrs merge." + ) + elif key == "control_parts": + # merge control parts + user_control_parts_dict = override_cfg_dict.get("control_parts") + if isinstance(user_control_parts_dict, dict): + # Initialize control_parts if it is None to avoid TypeError on item assignment + if base_cfg.control_parts is None: + base_cfg.control_parts = {} + for part_key, part_val in user_control_parts_dict.items(): + base_cfg.control_parts[part_key] = part_val + else: + logger.log_warning( + "control_parts should be a dictionary. Skipping control_parts merge." + ) + elif key == "urdf_cfg": + if base_cfg.urdf_cfg is None: + logger.log_warning( + f"There is no defined urdf_cfg in base robot cfg. Skipping urdf_cfg merge." + ) + continue + + # merge urdf components + user_urdf_cfg = override_cfg_dict.get("urdf_cfg") + if isinstance(user_urdf_cfg, dict): + for component in user_urdf_cfg.get("components", []): + base_cfg.urdf_cfg.add_component( + component_type=component.get("component_type"), + urdf_path=component.get("urdf_path"), + transform=component.get("transform"), + ) + else: + logger.log_warning( + "urdf_cfg should be a dictionary. Skipping urdf_cfg merge." + ) + else: + setattr(base_cfg, key, getattr(robot_cfg, key)) + + return base_cfg diff --git a/embodichain/lab/sim/utility/workspace_analyzer/__init__.py b/embodichain/lab/sim/utility/workspace_analyzer/__init__.py new file mode 100644 index 0000000..a2b3e80 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/__init__.py @@ -0,0 +1,67 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +""" +Workspace Analyzer for Robotic Manipulation + +A comprehensive tool for analyzing robot workspace reachability, computing +performance metrics, and generating visualizations. + +Supports two analysis modes: + 1. Joint Space Mode: Sample joint configurations → FK → Workspace points + 2. Cartesian Space Mode: Sample Cartesian positions → IK → Verify reachability + +Example: + >>> from embodichain.lab.sim.utility.workspace_analyzer import ( + ... WorkspaceAnalyzer, + ... AnalysisMode + ... ) + >>> # Joint space analysis (default) + >>> analyzer = WorkspaceAnalyzer(robot=robot) + >>> results = analyzer.analyze(num_samples=10000) + >>> + >>> # Cartesian space analysis + >>> from embodichain.lab.sim.utility.workspace_analyzer import WorkspaceAnalyzerConfig + >>> config = WorkspaceAnalyzerConfig(mode=AnalysisMode.CARTESIAN_SPACE) + >>> analyzer = WorkspaceAnalyzer(robot=robot, config=config) + >>> results = analyzer.analyze(num_samples=10000) +""" + +from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, + AnalysisMode, +) + +# Import submodules for convenience +from embodichain.lab.sim.utility.workspace_analyzer import configs +from embodichain.lab.sim.utility.workspace_analyzer import samplers +from embodichain.lab.sim.utility.workspace_analyzer import caches +from embodichain.lab.sim.utility.workspace_analyzer import visualizers +from embodichain.lab.sim.utility.workspace_analyzer import metrics +from embodichain.lab.sim.utility.workspace_analyzer import constraints + +__all__ = [ + "WorkspaceAnalyzer", + "WorkspaceAnalyzerConfig", + "AnalysisMode", + "configs", + "samplers", + "caches", + "visualizers", + "metrics", + "constraints", +] diff --git a/embodichain/lab/sim/utility/workspace_analyzer/caches/__init__.py b/embodichain/lab/sim/utility/workspace_analyzer/caches/__init__.py new file mode 100644 index 0000000..654db95 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/caches/__init__.py @@ -0,0 +1,31 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from embodichain.lab.sim.utility.workspace_analyzer.caches.base_cache import BaseCache +from embodichain.lab.sim.utility.workspace_analyzer.caches.memory_cache import ( + MemoryCache, +) +from embodichain.lab.sim.utility.workspace_analyzer.caches.disk_cache import DiskCache +from embodichain.lab.sim.utility.workspace_analyzer.caches.cache_manager import ( + CacheManager, +) + +__all__ = [ + "BaseCache", + "MemoryCache", + "DiskCache", + "CacheManager", +] diff --git a/embodichain/lab/sim/utility/workspace_analyzer/caches/base_cache.py b/embodichain/lab/sim/utility/workspace_analyzer/caches/base_cache.py new file mode 100644 index 0000000..5da0ebe --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/caches/base_cache.py @@ -0,0 +1,76 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from abc import ABC, abstractmethod +from typing import List +import numpy as np + + +all = [ + "BaseCache", +] + + +class BaseCache(ABC): + """Abstract base class for workspace sampling cache strategies. + + Defines the interface for different caching mechanisms (memory, disk) + used during workspace analysis sampling operations. + """ + + def __init__(self, batch_size: int = 5000, save_threshold: int = 10000000): + """Initialize base cache parameters. + + Args: + batch_size: Number of samples to process in each batch + save_threshold: Number of samples to accumulate before triggering save/flush + """ + self.batch_size = batch_size + self.save_threshold = save_threshold + self._total_processed = 0 + + @abstractmethod + def add(self, poses: List[np.ndarray]) -> None: + """Add pose samples to the cache. + + Args: + poses: List of 4x4 transformation matrices + """ + pass + + @abstractmethod + def flush(self) -> None: + """Flush any pending data in the cache.""" + pass + + @abstractmethod + def get_all(self) -> List[np.ndarray] | None: + """Retrieve all cached poses. + + Returns: + List of all cached 4x4 transformation matrices, or None if unavailable + """ + pass + + @abstractmethod + def clear(self) -> None: + """Clear all cached data.""" + pass + + @property + def total_processed(self) -> int: + """Total number of poses processed by this cache.""" + return self._total_processed diff --git a/embodichain/lab/sim/utility/workspace_analyzer/caches/cache_manager.py b/embodichain/lab/sim/utility/workspace_analyzer/caches/cache_manager.py new file mode 100644 index 0000000..5f60937 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/caches/cache_manager.py @@ -0,0 +1,105 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from typing import Literal + +from embodichain.lab.sim.utility.workspace_analyzer.caches.base_cache import BaseCache +from embodichain.lab.sim.utility.workspace_analyzer.caches.memory_cache import ( + MemoryCache, +) +from embodichain.lab.sim.utility.workspace_analyzer.caches.disk_cache import DiskCache +from embodichain.lab.sim.utility.workspace_analyzer.configs.cache_config import ( + CacheConfig, +) + + +all = [ + "CacheManager", +] + + +class CacheManager: + """Factory and manager for workspace sampling caches. + + Provides a unified interface for creating and managing different + cache strategies (memory vs disk). + """ + + @staticmethod + def create_cache( + cache_mode: Literal["memory", "disk"], + save_dir: str | None = None, + batch_size: int = 5000, + save_threshold: int = 10000000, + use_cached: bool = True, + ) -> BaseCache: + """Create a cache instance based on the specified mode. + + Args: + cache_mode: Caching strategy - "memory" or "disk" + save_dir: Directory for disk cache. If None in disk mode, uses + ~/.cache/embodichain/workspace_analyzer/session_TIMESTAMP + batch_size: Number of samples per batch + save_threshold: Threshold for saving/flushing data + use_cached: Whether to use existing cached data (disk mode only) + + Returns: + Configured cache instance + + Raises: + ValueError: If cache_mode is invalid + """ + if cache_mode not in ["memory", "disk"]: + raise ValueError( + f"Invalid cache_mode '{cache_mode}'. Must be 'memory' or 'disk'" + ) + + if cache_mode == "disk": + + return DiskCache( + save_dir=save_dir, + batch_size=batch_size, + save_threshold=save_threshold, + use_cached=use_cached, + ) + else: # memory + return MemoryCache(batch_size=batch_size, save_threshold=save_threshold) + + @staticmethod + def create_cache_from_config(config: CacheConfig) -> BaseCache | None: + """Create a cache instance from a CacheConfig object. + + Args: + config: CacheConfig instance with cache settings + + Returns: + Configured cache instance if enabled, None otherwise + """ + if not config.enabled: + return None + + cache_mode = "disk" if config.cache_dir is not None else "memory" + save_dir = str(config.cache_dir) if config.cache_dir else None + + return CacheManager.create_cache( + cache_mode=cache_mode, + save_dir=save_dir, + batch_size=5000, # Default batch size + save_threshold=config.max_cache_size_mb + * 1024 + * 1024, # Convert MB to bytes + use_cached=True, + ) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/caches/cache_utils.py b/embodichain/lab/sim/utility/workspace_analyzer/caches/cache_utils.py new file mode 100644 index 0000000..a167db7 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/caches/cache_utils.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +""" +Cache management utility for workspace analyzer. + +Usage: + python cache_utils.py list # List all cache sessions + python cache_utils.py info # Show cache session info + python cache_utils.py clean # Clean specific session + python cache_utils.py clean --all # Clean all cache sessions + python cache_utils.py size # Show total cache size +""" + +import os +import sys +import shutil +import argparse +from pathlib import Path +from datetime import datetime +from embodichain.utils import logger + +# Add current directory to path for local imports +sys.path.insert(0, os.path.dirname(__file__)) + +try: + from disk_cache import DiskCache +except ImportError: + print("Error: Unable to import DiskCache") + sys.exit(1) + + +def get_cache_root(): + """Get the root cache directory.""" + return os.path.expanduser("~/.cache/embodichain/workspace_analyzer") + + +def get_dir_size(path): + """Calculate total size of a directory in bytes.""" + total = 0 + try: + for entry in os.scandir(path): + if entry.is_file(follow_symlinks=False): + total += entry.stat().st_size + elif entry.is_dir(follow_symlinks=False): + total += get_dir_size(entry.path) + except (OSError, PermissionError): + # Directory access error, return partial total + pass + return total + + +def format_size(bytes_size): + """Format bytes to human-readable size.""" + for unit in ["B", "KB", "MB", "GB"]: + if bytes_size < 1024.0: + return f"{bytes_size:.2f} {unit}" + bytes_size /= 1024.0 + return f"{bytes_size:.2f} TB" + + +def list_sessions(): + """List all cache sessions.""" + cache_root = get_cache_root() + + if not os.path.exists(cache_root): + logger.log_info("No cache sessions found.") + logger.log_info(f"Cache directory: {cache_root}") + return + + sessions = [] + for item in os.listdir(cache_root): + session_path = os.path.join(cache_root, item) + if os.path.isdir(session_path): + size = get_dir_size(session_path) + mtime = os.path.getmtime(session_path) + sessions.append( + { + "name": item, + "path": session_path, + "size": size, + "modified": datetime.fromtimestamp(mtime), + } + ) + + if not sessions: + logger.log_info("No cache sessions found.") + return + + # Sort by modification time (newest first) + sessions.sort(key=lambda x: x["modified"], reverse=True) + + logger.log_info(f"\n{'Session Name':<40} {'Size':<12} {'Last Modified'}") + logger.log_info("-" * 80) + + total_size = 0 + for session in sessions: + logger.log_info( + f"{session['name']:<40} {format_size(session['size']):<12} " + f"{session['modified'].strftime('%Y-%m-%d %H:%M:%S')}" + ) + total_size += session["size"] + + logger.log_info("-" * 80) + logger.log_info( + f"{'Total':<40} {format_size(total_size):<12} {len(sessions)} session(s)" + ) + logger.log_info(f"\nCache location: {cache_root}") + + +def show_session_info(session_name): + """Show detailed information about a cache session.""" + cache_root = get_cache_root() + session_path = os.path.join(cache_root, session_name) + + if not os.path.exists(session_path): + logger.log_info(f"Session '{session_name}' not found.") + logger.log_info(f"Use 'list' command to see available sessions.") + return + + logger.log_info(f"\nSession: {session_name}") + logger.log_info(f"Path: {session_path}") + logger.log_info(f"Size: {format_size(get_dir_size(session_path))}") + logger.log_info( + f"Modified: {datetime.fromtimestamp(os.path.getmtime(session_path))}" + ) + + # Check for batches + batches_dir = os.path.join(session_path, "batches") + if os.path.exists(batches_dir): + batch_files = [f for f in os.listdir(batches_dir) if f.endswith(".npy")] + logger.log_info(f"Batches: {len(batch_files)} file(s)") + + if batch_files: + import numpy as np + + total_poses = 0 + for batch_file in batch_files: + batch_path = os.path.join(batches_dir, batch_file) + try: + data = np.load(batch_path) + total_poses += len(data) + except Exception as e: + logger.log_warning( + f"Warning: Failed to load batch file '{batch_file}': {e}" + ) + logger.log_info(f"Total poses: {total_poses:,}") + + +def clean_session(session_name): + """Clean a specific cache session.""" + cache_root = get_cache_root() + session_path = os.path.join(cache_root, session_name) + + if not os.path.exists(session_path): + logger.log_info(f"Session '{session_name}' not found.") + return + + size = get_dir_size(session_path) + response = input(f"Delete session '{session_name}' ({format_size(size)})? [y/N]: ") + + if response.lower() == "y": + shutil.rmtree(session_path) + logger.log_info(f"✓ Deleted session '{session_name}'") + else: + logger.log_info("Cancelled.") + + +def clean_all_sessions(): + """Clean all cache sessions.""" + cache_root = get_cache_root() + + if not os.path.exists(cache_root): + logger.log_info("No cache sessions found.") + return + + total_size = get_dir_size(cache_root) + sessions = [ + d for d in os.listdir(cache_root) if os.path.isdir(os.path.join(cache_root, d)) + ] + + if not sessions: + logger.log_info("No cache sessions found.") + return + + logger.log_info( + f"Found {len(sessions)} session(s), total size: {format_size(total_size)}" + ) + response = input(f"Delete all cache sessions? [y/N]: ") + + if response.lower() == "y": + shutil.rmtree(cache_root) + logger.log_info(f"✓ Deleted all cache sessions") + else: + logger.log_info("Cancelled.") + + +def show_total_size(): + """Show total cache size.""" + cache_root = get_cache_root() + + if not os.path.exists(cache_root): + logger.log_info("No cache found.") + logger.log_info(f"Cache directory: {cache_root}") + return + + total_size = get_dir_size(cache_root) + sessions = [ + d for d in os.listdir(cache_root) if os.path.isdir(os.path.join(cache_root, d)) + ] + + logger.log_info(f"\nCache location: {cache_root}") + logger.log_info(f"Total sessions: {len(sessions)}") + logger.log_info(f"Total size: {format_size(total_size)}") + + +def main(): + parser = argparse.ArgumentParser( + description="Manage workspace analyzer cache sessions", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + %(prog)s list List all cache sessions + %(prog)s info session_20241127 Show session details + %(prog)s clean session_20241127 Clean specific session + %(prog)s clean --all Clean all sessions + %(prog)s size Show total cache size + """, + ) + + parser.add_argument( + "command", choices=["list", "info", "clean", "size"], help="Command to execute" + ) + parser.add_argument( + "session", nargs="?", help="Session name (for info/clean commands)" + ) + parser.add_argument( + "--all", action="store_true", help="Apply to all sessions (for clean command)" + ) + + args = parser.parse_args() + + if args.command == "list": + list_sessions() + elif args.command == "info": + if not args.session: + print("Error: Session name required for 'info' command") + sys.exit(1) + show_session_info(args.session) + elif args.command == "clean": + if args.all: + clean_all_sessions() + elif args.session: + clean_session(args.session) + else: + print("Error: Specify a session name or use --all flag") + sys.exit(1) + elif args.command == "size": + show_total_size() + + +if __name__ == "__main__": + main() diff --git a/embodichain/lab/sim/utility/workspace_analyzer/caches/disk_cache.py b/embodichain/lab/sim/utility/workspace_analyzer/caches/disk_cache.py new file mode 100644 index 0000000..f2f649d --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/caches/disk_cache.py @@ -0,0 +1,197 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import gc +import os +from typing import List +import numpy as np +from tqdm import tqdm + + +from embodichain.lab.sim.utility.workspace_analyzer.caches.base_cache import BaseCache + +from embodichain.utils import logger + +all = [ + "DiskCache", +] + + +class DiskCache(BaseCache): + """Disk-based cache for workspace sampling. + + Saves pose samples to disk in batches to minimize memory usage. + Suitable for large-scale sampling operations. + + Default cache location: ~/.cache/embodichain/workspace_analyzer/ + """ + + @staticmethod + def get_default_cache_dir(subdir: str = "default") -> str: + """Get default cache directory in user's home. + + Args: + subdir: Subdirectory name under workspace_analyzer cache + + Returns: + Path to cache directory: ~/.cache/embodichain/workspace_analyzer/{subdir} + """ + cache_home = os.path.expanduser("~/.cache") + cache_dir = os.path.join( + cache_home, "embodichain", "workspace_analyzer", subdir + ) + return cache_dir + + def __init__( + self, + save_dir: str | None = None, + batch_size: int = 5000, + save_threshold: int = 10000000, + use_cached: bool = True, + ): + """Initialize disk cache. + + Args: + save_dir: Directory path for saving batch files. + If None, uses ~/.cache/embodichain/workspace_analyzer/default + batch_size: Number of samples per batch + save_threshold: Number of samples to accumulate before writing to disk + use_cached: Whether to use existing cached files if available + """ + super().__init__(batch_size, save_threshold) + + # Use default cache dir if not specified + if save_dir is None: + import time + + timestamp = time.strftime("%Y%m%d_%H%M%S") + save_dir = self.get_default_cache_dir(subdir=f"session_{timestamp}") + logger.log_info(f"Using default cache directory: {save_dir}") + + self.save_dir = save_dir + self.use_cached = use_cached + self._buffer: List[np.ndarray] = [] + self._batch_count = 0 + + # Create batches directory + self._batches_dir = os.path.join(save_dir, "batches") + os.makedirs(self._batches_dir, exist_ok=True) + + # Check for existing cached data + if use_cached and self._has_cached_data(): + logger.log_info(f"Found existing cached data in {self._batches_dir}") + + def _has_cached_data(self) -> bool: + """Check if cached batch files exist. + + Returns: + True if cached .npy files exist + """ + if not os.path.exists(self._batches_dir): + return False + npy_files = [f for f in os.listdir(self._batches_dir) if f.endswith(".npy")] + return len(npy_files) > 0 + + def add(self, poses: List[np.ndarray]) -> None: + """Add poses to buffer and save to disk when threshold is reached. + + Args: + poses: List of 4x4 transformation matrices + """ + self._buffer.extend(poses) + self._total_processed += len(poses) + + # Write to disk when buffer reaches threshold + if len(self._buffer) >= self.save_threshold: + self._save_batch() + + def _save_batch(self) -> None: + """Save current buffer to disk as a batch file.""" + if not self._buffer: + return + + batch_path = os.path.join( + self._batches_dir, f"batch_{self._batch_count:04d}.npy" + ) + np.save(batch_path, np.array(self._buffer)) + logger.log_info( + f"Saved batch {self._batch_count}: " + f"{len(self._buffer)} poses -> {batch_path}" + ) + + self._batch_count += 1 + self._buffer.clear() + gc.collect() + + def flush(self) -> None: + """Flush any remaining data in buffer to disk.""" + if self._buffer: + self._save_batch() + + def get_all(self) -> List[np.ndarray] | None: + """Load and merge all batch files from disk. + + Returns: + List of all cached poses merged from batch files, or None if no data + """ + # First flush any pending data + self.flush() + + # Get all batch files + npy_files = sorted( + [f for f in os.listdir(self._batches_dir) if f.endswith(".npy")] + ) + + if not npy_files: + return None + + logger.log_info(f"Loading {len(npy_files)} batch files...") + all_poses = [] + + for npy_file in tqdm(npy_files, desc="Merging batches"): + batch_path = os.path.join(self._batches_dir, npy_file) + try: + batch_data = np.load(batch_path) + all_poses.extend(batch_data) + except Exception as e: + logger.log_warning(f"Error loading {npy_file}: {str(e)}") + + logger.log_info(f"Loaded {len(all_poses)} total poses") + return all_poses if all_poses else None + + def clear(self) -> None: + """Clear all cached data and remove batch files.""" + self._buffer.clear() + self._batch_count = 0 + self._total_processed = 0 + + # Remove all batch files + if os.path.exists(self._batches_dir): + for file in os.listdir(self._batches_dir): + if file.endswith(".npy"): + os.remove(os.path.join(self._batches_dir, file)) + + gc.collect() + + def get_batch_count(self) -> int: + """Get number of batches written to disk. + + Returns: + Number of batch files on disk + """ + if not os.path.exists(self._batches_dir): + return 0 + return len([f for f in os.listdir(self._batches_dir) if f.endswith(".npy")]) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/caches/memory_cache.py b/embodichain/lab/sim/utility/workspace_analyzer/caches/memory_cache.py new file mode 100644 index 0000000..689cc5d --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/caches/memory_cache.py @@ -0,0 +1,78 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import gc +from typing import List +import numpy as np + +from embodichain.lab.sim.utility.workspace_analyzer.caches.base_cache import BaseCache + +all = [ + "MemoryCache", +] + + +class MemoryCache(BaseCache): + """In-memory cache for workspace sampling. + + Stores all pose samples in RAM for fast access. Suitable for + smaller datasets or when memory is not a constraint. + """ + + def __init__(self, batch_size: int = 5000, save_threshold: int = 10000000): + """Initialize memory cache. + + Args: + batch_size: Number of samples per processing batch + save_threshold: Threshold for triggering garbage collection + """ + super().__init__(batch_size, save_threshold) + self._poses: List[np.ndarray] = [] + + def add(self, poses: List[np.ndarray]) -> None: + """Add poses to in-memory storage. + + Args: + poses: List of 4x4 transformation matrices + """ + self._poses.extend(poses) + self._total_processed += len(poses) + + # Trigger garbage collection periodically + if len(self._poses) % 1000 == 0: + gc.collect() + + def flush(self) -> None: + """Flush operation (no-op for memory cache, but triggers GC).""" + gc.collect() + + def get_all(self) -> List[np.ndarray] | None: + """Retrieve all cached poses. + + Returns: + List of all cached poses, or None if empty + """ + return self._poses if self._poses else None + + def clear(self) -> None: + """Clear all cached data and free memory.""" + self._poses.clear() + self._total_processed = 0 + gc.collect() + + def __len__(self) -> int: + """Return number of cached poses.""" + return len(self._poses) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/configs/__init__.py b/embodichain/lab/sim/utility/workspace_analyzer/configs/__init__.py new file mode 100644 index 0000000..eb4e06a --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/configs/__init__.py @@ -0,0 +1,52 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from embodichain.lab.sim.utility.workspace_analyzer.configs.cache_config import ( + CacheConfig, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.dimension_constraint import ( + DimensionConstraint, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.sampling_config import ( + SamplingConfig, + SamplingStrategy, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.visualization_config import ( + VisualizationConfig, + VisualizationType, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.metric_config import ( + MetricConfig, + MetricType, + ReachabilityConfig, + ManipulabilityConfig, + DensityConfig, +) + + +__all__ = [ + "CacheConfig", + "DimensionConstraint", + "SamplingConfig", + "SamplingStrategy", + "VisualizationConfig", + "VisualizationType", + "MetricConfig", + "MetricType", + "ReachabilityConfig", + "ManipulabilityConfig", + "DensityConfig", +] diff --git a/embodichain/lab/sim/utility/workspace_analyzer/configs/cache_config.py b/embodichain/lab/sim/utility/workspace_analyzer/configs/cache_config.py new file mode 100644 index 0000000..eb86830 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/configs/cache_config.py @@ -0,0 +1,41 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from dataclasses import dataclass +from pathlib import Path + + +@dataclass +class CacheConfig: + """Configuration for caching workspace analysis results.""" + + enabled: bool = True + """Whether to enable caching of analysis results.""" + + cache_dir: Path | None = None + """Directory to store cache files. If None, uses default system cache directory.""" + + use_hash: bool = True + """Whether to use hash-based cache keys (based on robot config and parameters).""" + + compression: bool = True + """Whether to compress cache files to save disk space.""" + + max_cache_size_mb: int = 1000 + """Maximum total size of cache directory in megabytes. Old files will be removed if exceeded.""" + + cache_format: str = "npz" + """Format for cache files. Options: 'npz' (numpy), 'pkl' (pickle), 'h5' (hdf5).""" diff --git a/embodichain/lab/sim/utility/workspace_analyzer/configs/dimension_constraint.py b/embodichain/lab/sim/utility/workspace_analyzer/configs/dimension_constraint.py new file mode 100644 index 0000000..9dbccba --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/configs/dimension_constraint.py @@ -0,0 +1,45 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from dataclasses import dataclass, field +from typing import List, Tuple +import numpy as np + + +@dataclass +class DimensionConstraint: + """Configuration for dimensional constraints in workspace analysis.""" + + min_bounds: np.ndarray | None = None + """Minimum bounds for workspace [x_min, y_min, z_min] in meters.""" + + max_bounds: np.ndarray | None = None + """Maximum bounds for workspace [x_max, y_max, z_max] in meters.""" + + joint_limits_scale: float = 1.0 + """Scale factor for joint limits (1.0 = use full range, 0.8 = use 80% of range).""" + + exclude_zones: List[Tuple[np.ndarray, np.ndarray]] = field(default_factory=list) + """List of excluded zones as [(min_bounds, max_bounds), ...]. Robot end-effector should avoid these regions.""" + + ground_height: float = 0.0 + """Ground plane height in meters. Points below this will be filtered out.""" + + enforce_collision_free: bool = False + """Whether to enforce collision-free constraints during analysis.""" + + self_collision_check: bool = False + """Whether to check for self-collision when analyzing workspace.""" diff --git a/embodichain/lab/sim/utility/workspace_analyzer/configs/metric_config.py b/embodichain/lab/sim/utility/workspace_analyzer/configs/metric_config.py new file mode 100644 index 0000000..f4f392e --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/configs/metric_config.py @@ -0,0 +1,114 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from dataclasses import dataclass +from typing import List +from enum import Enum + + +class MetricType(Enum): + """Types of workspace metrics.""" + + REACHABILITY = "reachability" + """Measures reachable workspace volume and coverage.""" + + MANIPULABILITY = "manipulability" + """Measures dexterity/manipulability throughout workspace.""" + + DENSITY = "density" + """Measures point density and workspace distribution.""" + + ALL = "all" + """Compute all available metrics.""" + + +@dataclass +class ReachabilityConfig: + """Configuration for reachability metric.""" + + voxel_size: float = 0.01 + """Size of voxels for volume calculation in meters.""" + + min_points_per_voxel: int = 1 + """Minimum number of points in a voxel to consider it reachable.""" + + compute_coverage: bool = True + """Whether to compute coverage percentage relative to bounding box.""" + + +@dataclass +class ManipulabilityConfig: + """Configuration for manipulability metric.""" + + jacobian_threshold: float = 0.01 + """Minimum manipulability value to consider valid.""" + + compute_isotropy: bool = True + """Whether to compute isotropy index (condition number).""" + + compute_heatmap: bool = False + """Whether to generate manipulability heatmap.""" + + +@dataclass +class DensityConfig: + """Configuration for density metric.""" + + radius: float = 0.05 + """Radius for local density estimation in meters.""" + + k_neighbors: int = 30 + """Number of nearest neighbors for density estimation.""" + + compute_distribution: bool = True + """Whether to compute density distribution statistics.""" + + +@dataclass +class MetricConfig: + """Configuration for workspace analysis metrics.""" + + enabled_metrics: List[MetricType] = None + """List of metrics to compute. If None, computes all metrics.""" + + reachability: ReachabilityConfig = None + """Configuration for reachability metric.""" + + manipulability: ManipulabilityConfig = None + """Configuration for manipulability metric.""" + + density: DensityConfig = None + """Configuration for density metric.""" + + save_results: bool = True + """Whether to save metric results to file.""" + + output_format: str = "json" + """Output format for metrics. Options: 'json', 'yaml', 'pkl'.""" + + def __post_init__(self): + """Initialize sub-configs with defaults if not provided.""" + if self.enabled_metrics is None: + self.enabled_metrics = [MetricType.ALL] + + if self.reachability is None: + self.reachability = ReachabilityConfig() + + if self.manipulability is None: + self.manipulability = ManipulabilityConfig() + + if self.density is None: + self.density = DensityConfig() diff --git a/embodichain/lab/sim/utility/workspace_analyzer/configs/sampling_config.py b/embodichain/lab/sim/utility/workspace_analyzer/configs/sampling_config.py new file mode 100644 index 0000000..4c23066 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/configs/sampling_config.py @@ -0,0 +1,81 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +from enum import Enum +from dataclasses import dataclass +from typing import Callable + + +class SamplingStrategy(Enum): + """Sampling strategy for joint space.""" + + UNIFORM = "uniform" # Uniform grid sampling + RANDOM = "random" # Random sampling + HALTON = "halton" # Quasi-random Halton sequence + SOBOL = "sobol" # Quasi-random Sobol sequence + LATIN_HYPERCUBE = "lhs" # Latin Hypercube Sampling + IMPORTANCE = "importance" # Importance sampling (requires weight function) + GAUSSIAN = "gaussian" # Gaussian (normal) distribution sampling + SPHERE = "sphere" # Sphere-constrained uniform sampling + + +@dataclass +class SamplingConfig: + """Configuration for sampling strategies in workspace analysis.""" + + strategy: SamplingStrategy = None + """Sampling strategy to use.""" + + num_samples: int = 1000 + """Number of samples to generate.""" + + grid_resolution: int = 10 + """Resolution for grid sampling (used with UNIFORM strategy).""" + + batch_size: int = 1000 + """Number of samples to process in each batch.""" + + seed: int = 42 + """Random seed for reproducibility.""" + + importance_weight_func: Callable | None = None + """Weight function for importance sampling (used with IMPORTANCE strategy).""" + + gaussian_mean: float | None = None + """Mean for Gaussian sampling (used with GAUSSIAN strategy). If None, uses center of bounds.""" + + gaussian_std: float | None = None + """Standard deviation for Gaussian sampling (used with GAUSSIAN strategy). If None, uses 1/6 of range.""" + + # Sphere sampling parameters + sphere_center_mode: str = "bounds_center" + """How to determine sphere center for SPHERE strategy. Options: 'bounds_center', 'custom', 'auto'.""" + + sphere_radius_mode: str = "inscribed" + """How to determine sphere radius for SPHERE strategy. Options: 'inscribed', 'circumscribed', 'custom'.""" + + sphere_boundary_handling: str = "reject" + """How to handle boundary violations for SPHERE strategy. Options: 'clip', 'reject', 'extend'.""" + + sphere_center: list | None = None + """Custom sphere center for SPHERE strategy (used when sphere_center_mode='custom').""" + + sphere_radius: float | None = None + """Custom sphere radius for SPHERE strategy (used when sphere_radius_mode='custom').""" + + def __post_init__(self): + """Set default strategy after initialization.""" + if self.strategy is None: + self.strategy = SamplingStrategy.UNIFORM diff --git a/embodichain/lab/sim/utility/workspace_analyzer/configs/visualization_config.py b/embodichain/lab/sim/utility/workspace_analyzer/configs/visualization_config.py new file mode 100644 index 0000000..5877420 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/configs/visualization_config.py @@ -0,0 +1,96 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from enum import Enum +from dataclasses import dataclass +from typing import Union + + +class VisualizationType(Enum): + POINT_CLOUD = "point_cloud" + VOXEL = "voxel" + SPHERE = "sphere" + AXIS = "axis" + # MESH = "mesh" + # HEATMAP = "heatmap" + + +@dataclass +class VisualizationConfig: + """Visualization configuration""" + + enabled: bool = True + """Enable or disable visualization.""" + + vis_type: Union[VisualizationType, str] = VisualizationType.POINT_CLOUD + """Type of visualization to use. Can be VisualizationType enum or string. + + Available visualization types: + - POINT_CLOUD: Fast point cloud rendering, good for large datasets + - VOXEL: Volumetric voxel grid representation for occupancy mapping + - SPHERE: Smooth sphere rendering for publication-quality figures + - AXIS: Coordinate frame visualization for poses and transformations + + Examples: + vis_type = VisualizationType.SPHERE + vis_type = "point_cloud" + vis_type = "axis" # For coordinate frames + """ + + voxel_size: float = 0.05 + """Voxel size for downsampling.""" + + nb_neighbors: int = 20 + """Number of neighbors for statistical outlier removal.""" + + std_ratio: float = 2.0 + """Standard deviation ratio for statistical outlier removal.""" + + is_voxel_down: bool = True + """Enable voxel downsampling.""" + + color_by_distance: bool = True + """Color points by distance.""" + + point_size: float = 4.0 + """Size of points in visualization.""" + + alpha: float = 0.5 + """Transparency level of points.""" + + sphere_radius: float = 0.005 + """Radius of spheres for sphere visualization.""" + + sphere_resolution: int = 10 + """Sphere mesh resolution for sphere visualization.""" + + axis_length: float = 0.05 + """Length of coordinate axes for axis visualization.""" + + axis_size: float = 0.003 + """Thickness/size of coordinate axes for axis visualization.""" + + show_unreachable_points: bool = True + """Whether to show unreachable points in Cartesian space and Plane sampling modes. + + If True, shows both reachable (green, large) and unreachable (red, small) points. + If False, only shows reachable points in Cartesian space and Plane sampling modes. + Has no effect in Joint space mode (all points are always shown as reachable). + + Note: This parameter now supports all IK-based analysis modes including: + - AnalysisMode.CARTESIAN_SPACE + - AnalysisMode.PLANE_SAMPLING + """ diff --git a/embodichain/lab/sim/utility/workspace_analyzer/constraints/__init__.py b/embodichain/lab/sim/utility/workspace_analyzer/constraints/__init__.py new file mode 100644 index 0000000..9c02f03 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/constraints/__init__.py @@ -0,0 +1,24 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from .base_constraint import BaseConstraintChecker, IConstraintChecker +from .workspace_constraint import WorkspaceConstraintChecker + +__all__ = [ + "BaseConstraintChecker", + "IConstraintChecker", + "WorkspaceConstraintChecker", +] diff --git a/embodichain/lab/sim/utility/workspace_analyzer/constraints/base_constraint.py b/embodichain/lab/sim/utility/workspace_analyzer/constraints/base_constraint.py new file mode 100644 index 0000000..c0b7a4c --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/constraints/base_constraint.py @@ -0,0 +1,199 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from abc import ABC, abstractmethod +from typing import Union +import numpy as np +import torch + +from embodichain.utils import logger + + +__all__ = [ + "IConstraintChecker", + "BaseConstraintChecker", +] + + +class IConstraintChecker: + """Interface for constraint checkers. + + This protocol defines the contract that all constraint checkers must follow. + """ + + def check_bounds( + self, points: Union[torch.Tensor, np.ndarray] + ) -> Union[torch.Tensor, np.ndarray]: + """Check if points are within bounds. + + Args: + points: Array of shape (N, 3) containing point positions. + + Returns: + Boolean array of shape (N,) indicating which points satisfy bounds. + """ + ... + + def filter_points( + self, points: Union[torch.Tensor, np.ndarray] + ) -> Union[torch.Tensor, np.ndarray]: + """Filter points to keep only those satisfying all constraints. + + Args: + points: Array of shape (N, 3) containing point positions. + + Returns: + Filtered array of shape (M, 3) where M <= N. + """ + ... + + +class BaseConstraintChecker(ABC): + """Abstract base class for workspace constraint checkers. + + Provides common functionality for checking dimensional constraints, + bounds, excluded zones, and other workspace limitations. + """ + + def __init__( + self, + min_bounds: np.ndarray | None = None, + max_bounds: np.ndarray | None = None, + ground_height: float = 0.0, + device: torch.device | None = None, + ): + """Initialize the constraint checker. + + Args: + min_bounds: Minimum bounds [x_min, y_min, z_min]. If None, no lower limit. + max_bounds: Maximum bounds [x_max, y_max, z_max]. If None, no upper limit. + ground_height: Ground plane height. Points below this are filtered. + device: PyTorch device for tensor operations. Defaults to cpu. + """ + self.min_bounds = np.array(min_bounds) if min_bounds is not None else None + self.max_bounds = np.array(max_bounds) if max_bounds is not None else None + self.ground_height = ground_height + self.device = device if device is not None else torch.device("cpu") + + # Validate bounds + if self.min_bounds is not None and self.max_bounds is not None: + if not np.all(self.min_bounds < self.max_bounds): + logger.log_warning( + "min_bounds should be less than max_bounds in all dimensions" + ) + + def check_bounds( + self, points: Union[torch.Tensor, np.ndarray] + ) -> Union[torch.Tensor, np.ndarray]: + """Check if points are within bounds. + + Args: + points: Array of shape (N, 3) containing point positions. + + Returns: + Boolean array of shape (N,) indicating which points are within bounds. + """ + is_tensor = isinstance(points, torch.Tensor) + + if is_tensor: + valid = torch.ones(len(points), dtype=torch.bool, device=points.device) + else: + valid = np.ones(len(points), dtype=bool) + + # Check minimum bounds + if self.min_bounds is not None: + if is_tensor: + min_bounds_t = torch.tensor( + self.min_bounds, dtype=points.dtype, device=points.device + ) + valid &= torch.all(points >= min_bounds_t, dim=1) + else: + valid &= np.all(points >= self.min_bounds, axis=1) + + # Check maximum bounds + if self.max_bounds is not None: + if is_tensor: + max_bounds_t = torch.tensor( + self.max_bounds, dtype=points.dtype, device=points.device + ) + valid &= torch.all(points <= max_bounds_t, dim=1) + else: + valid &= np.all(points <= self.max_bounds, axis=1) + + # Check ground height + if is_tensor: + valid &= points[:, 2] >= self.ground_height + else: + valid &= points[:, 2] >= self.ground_height + + return valid + + def filter_points( + self, points: Union[torch.Tensor, np.ndarray] + ) -> Union[torch.Tensor, np.ndarray]: + """Filter points to keep only those satisfying all constraints. + + Args: + points: Array of shape (N, 3) containing point positions. + + Returns: + Filtered array of shape (M, 3) where M <= N. + """ + valid = self.check_bounds(points) + return points[valid] + + @abstractmethod + def check_collision( + self, points: Union[torch.Tensor, np.ndarray] + ) -> Union[torch.Tensor, np.ndarray]: + """Check if points are collision-free (to be implemented by subclasses). + + Args: + points: Array of shape (N, 3) containing point positions. + + Returns: + Boolean array of shape (N,) indicating which points are collision-free. + """ + pass + + def get_bounds_volume(self) -> float: + """Calculate the volume of the bounded workspace. + + Returns: + Volume in cubic meters, or inf if unbounded. + """ + if self.min_bounds is None or self.max_bounds is None: + return float("inf") + + dimensions = self.max_bounds - self.min_bounds + return float(np.prod(dimensions)) + + def get_bounds_info(self) -> dict: + """Get information about the workspace bounds. + + Returns: + Dictionary containing bounds information. + """ + return { + "min_bounds": ( + self.min_bounds.tolist() if self.min_bounds is not None else None + ), + "max_bounds": ( + self.max_bounds.tolist() if self.max_bounds is not None else None + ), + "ground_height": self.ground_height, + "volume": self.get_bounds_volume(), + } diff --git a/embodichain/lab/sim/utility/workspace_analyzer/constraints/workspace_constraint.py b/embodichain/lab/sim/utility/workspace_analyzer/constraints/workspace_constraint.py new file mode 100644 index 0000000..0947193 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/constraints/workspace_constraint.py @@ -0,0 +1,280 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from typing import Union, List, Tuple +import numpy as np +import torch + +from embodichain.lab.sim.utility.workspace_analyzer.constraints.base_constraint import ( + BaseConstraintChecker, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.dimension_constraint import ( + DimensionConstraint, +) + + +__all__ = [ + "WorkspaceConstraintChecker", +] + + +class WorkspaceConstraintChecker(BaseConstraintChecker): + """Workspace constraint checker for robotic workspace analysis. + + Main features: + - Boundary constraint checking (prevents robot from exceeding safe working range) + - Excluded zone checking (avoids obstacles, tables, and other fixed objects) + - Ground height constraints (prevents robot collision with ground or work surface) + - Configuration-based creation for easy reuse across different scenarios + + Extends base checker with excluded zones and configuration-based setup. + """ + + def __init__( + self, + min_bounds: np.ndarray | None = None, + max_bounds: np.ndarray | None = None, + ground_height: float = 0.0, + exclude_zones: List[Tuple[np.ndarray, np.ndarray]] | None = None, + device: torch.device | None = None, + ): + """Initialize the workspace constraint checker. + + Args: + min_bounds: Minimum bounds [x_min, y_min, z_min] in meters + max_bounds: Maximum bounds [x_max, y_max, z_max] in meters + ground_height: Ground plane height in meters + exclude_zones: List of excluded zones as [(min_bounds, max_bounds), ...] + device: PyTorch device (CPU/GPU) + """ + super().__init__(min_bounds, max_bounds, ground_height, device) + self.exclude_zones = exclude_zones or [] + + @classmethod + def from_config( + cls, config: DimensionConstraint, device: torch.device | None = None + ): + """Create a constraint checker from a DimensionConstraint config (recommended approach). + + This is the recommended way to create constraint checkers. Using configuration objects + allows for easy reuse of the same constraint settings across different scenarios. + + Args: + config: Dimension constraint configuration object containing bounds, excluded zones, etc. + device: PyTorch device for tensor operations + + Returns: + Configured WorkspaceConstraintChecker instance + + Example: + >>> from embodichain.lab.sim.utility.workspace_analyzer.configs import DimensionConstraint + >>> config = DimensionConstraint( + ... min_bounds=np.array([-1, -1, 0]), + ... max_bounds=np.array([1, 1, 2]), + ... exclude_zones=[] + ... ) + >>> checker = WorkspaceConstraintChecker.from_config(config) + """ + return cls( + min_bounds=config.min_bounds, + max_bounds=config.max_bounds, + ground_height=config.ground_height, + exclude_zones=config.exclude_zones, + device=device, + ) + + def check_collision( + self, points: Union[torch.Tensor, np.ndarray] + ) -> Union[torch.Tensor, np.ndarray]: + """Check if points are not in excluded zones (collision checking). + + This method checks whether the robot end-effector positions would collide with + predefined obstacle zones such as cabinets, tables, walls, and other fixed obstacles. + + Args: + points: Array of shape (N, 3) containing 3D point positions + + Returns: + Boolean array of shape (N,) indicating which points are collision-free + True = collision-free (safe), False = collision detected (dangerous) + """ + is_tensor = isinstance(points, torch.Tensor) + + if is_tensor: + valid = torch.ones(len(points), dtype=torch.bool, device=points.device) + else: + valid = np.ones(len(points), dtype=bool) + + # Check each excluded zone + for min_zone, max_zone in self.exclude_zones: + if is_tensor: + min_zone_t = torch.tensor( + min_zone, dtype=points.dtype, device=points.device + ) + max_zone_t = torch.tensor( + max_zone, dtype=points.dtype, device=points.device + ) + # Points inside excluded zone + in_zone = torch.all(points >= min_zone_t, dim=1) & torch.all( + points <= max_zone_t, dim=1 + ) + valid &= ~in_zone # Exclude these points + else: + in_zone = np.all(points >= min_zone, axis=1) & np.all( + points <= max_zone, axis=1 + ) + valid &= ~in_zone + + return valid + + def filter_points( + self, points: Union[torch.Tensor, np.ndarray] + ) -> Union[torch.Tensor, np.ndarray]: + """Filter points to keep only those satisfying all constraints. + + This is a comprehensive filtering method that checks both boundary constraints + and collision constraints, returning only safe and reachable workspace points. + + Args: + points: Array of shape (N, 3) containing 3D point positions + + Returns: + Filtered array of shape (M, 3) where M <= N + Contains only points that satisfy all constraint conditions + """ + # First filter by bounds + valid_bounds = self.check_bounds(points) + + # Then filter by collision + valid_collision = self.check_collision(points) + + # Combine both checks + valid = valid_bounds & valid_collision + + return points[valid] + + def add_exclude_zone(self, min_bounds: np.ndarray, max_bounds: np.ndarray) -> None: + """Add an excluded zone (obstacle region) to the workspace. + + Used for dynamically adding obstacles such as temporarily placed objects, + newly installed equipment, etc. + + Args: + min_bounds: Minimum bounds of the excluded zone [x_min, y_min, z_min] in meters + max_bounds: Maximum bounds of the excluded zone [x_max, y_max, z_max] in meters + + Example: + >>> # Add a cubic obstacle + >>> checker.add_exclude_zone( + ... min_bounds=np.array([0.2, 0.2, 0.0]), + ... max_bounds=np.array([0.4, 0.4, 0.2]) + ... ) + """ + self.exclude_zones.append((np.array(min_bounds), np.array(max_bounds))) + + def clear_exclude_zones(self) -> None: + """Remove all excluded zones. + + Used to reset the workspace environment by clearing all previously + configured obstacle zones. + """ + self.exclude_zones = [] + + def get_num_exclude_zones(self) -> int: + """Get the number of current excluded zones. + + Returns: + Number of excluded zones + """ + return len(self.exclude_zones) + + +# ==================== +# Usage Examples +# ==================== +""" +Typical usage scenarios: + +1. Basic usage - Setting workspace boundaries: +```python +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.configs import DimensionConstraint +from embodichain.lab.sim.utility.workspace_analyzer.constraints import WorkspaceConstraintChecker + +# Create via configuration (recommended approach) +config = DimensionConstraint( + min_bounds=np.array([-0.8, -0.6, 0.0]), # Work range: ±80cm left/right, ±60cm front/back, height ≥ 0cm + max_bounds=np.array([0.8, 0.6, 1.5]), # Max height 150cm + ground_height=0.0, # Ground height + exclude_zones=[] # No obstacles initially +) +checker = WorkspaceConstraintChecker.from_config(config) + +# Check if points are within safe range +test_points = np.array([ + [0.5, 0.3, 0.8], # Safe point + [1.0, 0.0, 0.5], # Exceeds x boundary + [0.0, 0.0, -0.1], # Below ground level +]) +valid_mask = checker.check_bounds(test_points) +print(f"Safe points: {valid_mask}") # [True, False, False] +``` + +2. Advanced usage - Adding obstacle zones: +```python +# Add a table as an obstacle +checker.add_exclude_zone( + min_bounds=np.array([0.2, 0.1, 0.0]), + max_bounds=np.array([0.7, 0.5, 0.8]) +) + +# Check both boundaries and collisions simultaneously +test_points = np.array([ + [0.1, 0.0, 0.5], # Within boundaries, not in obstacle zone + [0.4, 0.3, 0.4], # Within boundaries, but inside table zone +]) +valid_bounds = checker.check_bounds(test_points) +valid_collision = checker.check_collision(test_points) +overall_valid = valid_bounds & valid_collision +print(f"Comprehensive safety check: {overall_valid}") # [True, False] + +# Or use the comprehensive filtering method directly +safe_points = checker.filter_points(test_points) +print(f"Filtered safe points: {safe_points}") # Only returns [0.1, 0.0, 0.5] +``` + +3. Application in robot workspace analysis: +```python +# Usage in WorkspaceAnalyzer +from embodichain.lab.sim.utility.workspace_analyzer import WorkspaceAnalyzer + +# Configure work environment with complex obstacles +config = DimensionConstraint( + min_bounds=np.array([-1.0, -1.0, 0.0]), + max_bounds=np.array([1.0, 1.0, 2.0]), + exclude_zones=[ + # Table + (np.array([0.3, 0.2, 0.0]), np.array([0.8, 0.7, 0.8])), + # Pillar + (np.array([-0.1, -0.1, 0.0]), np.array([0.1, 0.1, 2.0])), + ] +) + +# WorkspaceAnalyzer will automatically use these constraints to filter invalid workspace points +analyzer = WorkspaceAnalyzer(robot, config=analyzer_config) +results = analyzer.analyze(visualize=True) +``` +""" diff --git a/embodichain/lab/sim/utility/workspace_analyzer/metrics/__init__.py b/embodichain/lab/sim/utility/workspace_analyzer/metrics/__init__.py new file mode 100644 index 0000000..2385092 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/metrics/__init__.py @@ -0,0 +1,35 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from embodichain.lab.sim.utility.workspace_analyzer.metrics.base_metric import ( + BaseMetric, +) +from embodichain.lab.sim.utility.workspace_analyzer.metrics.reachability_metric import ( + ReachabilityMetric, +) +from embodichain.lab.sim.utility.workspace_analyzer.metrics.manipulability_metric import ( + ManipulabilityMetric, +) +from embodichain.lab.sim.utility.workspace_analyzer.metrics.density_metric import ( + DensityMetric, +) + +__all__ = [ + "BaseMetric", + "ReachabilityMetric", + "ManipulabilityMetric", + "DensityMetric", +] diff --git a/embodichain/lab/sim/utility/workspace_analyzer/metrics/base_metric.py b/embodichain/lab/sim/utility/workspace_analyzer/metrics/base_metric.py new file mode 100644 index 0000000..084326b --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/metrics/base_metric.py @@ -0,0 +1,83 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from abc import ABC, abstractmethod +from typing import Dict, Any +import numpy as np +import torch + + +class BaseMetric(ABC): + """Base class for workspace metrics. + + All metrics should inherit from this class and implement the compute method. + """ + + def __init__(self, config: Any | None = None): + """Initialize the metric. + + Args: + config: Configuration object for the metric. + """ + self.config = config + self.results: Dict[str, Any] = {} + + @abstractmethod + def compute( + self, + workspace_points: np.ndarray, + joint_configurations: np.ndarray | None = None, + **kwargs, + ) -> Dict[str, Any]: + """Compute the metric. + + Args: + workspace_points: Workspace points in Cartesian space, shape (N, 3). + joint_configurations: Joint configurations, shape (N, num_joints). + **kwargs: Additional arguments specific to the metric. + + Returns: + Dictionary containing metric results. + """ + pass + + def reset(self) -> None: + """Reset metric results.""" + self.results = {} + + def get_results(self) -> Dict[str, Any]: + """Get computed metric results. + + Returns: + Dictionary containing metric results. + """ + return self.results + + def _to_numpy(self, data: Any) -> np.ndarray: + """Convert data to numpy array. + + Args: + data: Input data (numpy array or torch tensor). + + Returns: + Numpy array. + """ + if isinstance(data, torch.Tensor): + return data.cpu().numpy() + elif isinstance(data, np.ndarray): + return data + else: + return np.array(data) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/metrics/density_metric.py b/embodichain/lab/sim/utility/workspace_analyzer/metrics/density_metric.py new file mode 100644 index 0000000..3ee9e95 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/metrics/density_metric.py @@ -0,0 +1,152 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from typing import Dict, Any +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.metrics.base_metric import ( + BaseMetric, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.metric_config import ( + DensityConfig, +) + + +class DensityMetric(BaseMetric): + """Density metric for workspace analysis. + + Computes point density and workspace distribution statistics. + """ + + def __init__(self, config: DensityConfig | None = None): + """Initialize density metric. + + Args: + config: Density configuration. + """ + super().__init__(config or DensityConfig()) + + def compute( + self, + workspace_points: np.ndarray, + joint_configurations: np.ndarray | None = None, + **kwargs, + ) -> Dict[str, Any]: + """Compute density metrics. + + Args: + workspace_points: Workspace points in Cartesian space, shape (N, 3). + joint_configurations: Joint configurations, shape (N, num_joints). + **kwargs: Additional arguments. + + Returns: + Dictionary containing: + - mean_density: Average local density + - std_density: Standard deviation + - min_density: Minimum density + - max_density: Maximum density + - density_distribution: Histogram of density values (if enabled) + """ + points = self._to_numpy(workspace_points) + + if len(points) < 2: + return { + "mean_density": 0.0, + "std_density": 0.0, + "min_density": 0.0, + "max_density": 0.0, + } + + # Compute local density for each point + densities = self._compute_local_density(points) + + self.results = { + "mean_density": float(densities.mean()), + "std_density": float(densities.std()), + "min_density": float(densities.min()), + "max_density": float(densities.max()), + } + + # Compute distribution statistics if requested + if self.config.compute_distribution: + hist, bin_edges = np.histogram(densities, bins=10) + self.results["density_distribution"] = { + "histogram": hist.tolist(), + "bin_edges": bin_edges.tolist(), + } + + return self.results + + def _compute_local_density(self, points: np.ndarray) -> np.ndarray: + """Compute local density for each point. + + Args: + points: Point cloud, shape (N, 3). + + Returns: + Local densities, shape (N,). + """ + n_points = len(points) + densities = np.zeros(n_points) + + # Use radius-based density estimation for better performance + radius = self.config.radius + + for i in range(n_points): + # Compute distances to all other points + distances = np.linalg.norm(points - points[i], axis=1) + + # Count neighbors within radius + num_neighbors = np.sum(distances <= radius) - 1 # Exclude self + + # Density = neighbors / volume of sphere + volume = (4.0 / 3.0) * np.pi * (radius**3) + densities[i] = num_neighbors / volume if volume > 0 else 0.0 + + return densities + + def _compute_knn_density(self, points: np.ndarray) -> np.ndarray: + """Compute k-nearest neighbors density. + + Alternative method using k-nearest neighbors instead of fixed radius. + + Args: + points: Point cloud, shape (N, 3). + + Returns: + Local densities, shape (N,). + """ + n_points = len(points) + k = min(self.config.k_neighbors, n_points - 1) + + if k <= 0: + return np.zeros(n_points) + + densities = np.zeros(n_points) + + for i in range(n_points): + # Compute distances to all other points + distances = np.linalg.norm(points - points[i], axis=1) + + # Find k-nearest neighbors (excluding self) + distances[i] = np.inf + knn_distances = np.partition(distances, k)[:k] + + # Density = k / volume of sphere containing k neighbors + max_distance = knn_distances.max() + volume = (4.0 / 3.0) * np.pi * (max_distance**3) + densities[i] = k / volume if volume > 0 else 0.0 + + return densities diff --git a/embodichain/lab/sim/utility/workspace_analyzer/metrics/manipulability_metric.py b/embodichain/lab/sim/utility/workspace_analyzer/metrics/manipulability_metric.py new file mode 100644 index 0000000..6080984 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/metrics/manipulability_metric.py @@ -0,0 +1,151 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from typing import Dict, Any +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.metrics.base_metric import ( + BaseMetric, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.metric_config import ( + ManipulabilityConfig, +) + + +class ManipulabilityMetric(BaseMetric): + """Manipulability metric for workspace analysis. + + Computes dexterity and manipulability measures throughout the workspace. + Note: Full implementation requires robot Jacobian computation. + """ + + def __init__(self, config: ManipulabilityConfig | None = None): + """Initialize manipulability metric. + + Args: + config: Manipulability configuration. + """ + super().__init__(config or ManipulabilityConfig()) + + def compute( + self, + workspace_points: np.ndarray, + joint_configurations: np.ndarray | None = None, + jacobians: np.ndarray | None = None, + **kwargs, + ) -> Dict[str, Any]: + """Compute manipulability metrics. + + Args: + workspace_points: Workspace points in Cartesian space, shape (N, 3). + joint_configurations: Joint configurations, shape (N, num_joints). + jacobians: Precomputed Jacobian matrices, shape (N, 6, num_joints). + **kwargs: Additional arguments. + + Returns: + Dictionary containing: + - mean_manipulability: Average manipulability index + - std_manipulability: Standard deviation + - min_manipulability: Minimum value + - max_manipulability: Maximum value + - mean_condition: Average condition number (if isotropy enabled) + """ + points = self._to_numpy(workspace_points) + + if len(points) == 0: + return { + "mean_manipulability": 0.0, + "std_manipulability": 0.0, + "min_manipulability": 0.0, + "max_manipulability": 0.0, + } + + # If Jacobians are not provided, we cannot compute true manipulability + # Return placeholder statistics + if jacobians is None: + # Estimate based on distance from centroid (simple heuristic) + centroid = points.mean(axis=0) + distances = np.linalg.norm(points - centroid, axis=1) + + # Normalize to [0, 1] range (higher manipulability near center) + max_dist = distances.max() if distances.max() > 0 else 1.0 + manipulability_scores = 1.0 - (distances / max_dist) + + # Filter by threshold + valid_mask = manipulability_scores >= self.config.jacobian_threshold + valid_scores = manipulability_scores[valid_mask] + + if len(valid_scores) == 0: + valid_scores = np.array([0.0]) + else: + # Compute true manipulability from Jacobians + manipulability_scores = self._compute_manipulability_index(jacobians) + valid_mask = manipulability_scores >= self.config.jacobian_threshold + valid_scores = manipulability_scores[valid_mask] + + self.results = { + "mean_manipulability": float(valid_scores.mean()), + "std_manipulability": float(valid_scores.std()), + "min_manipulability": float(valid_scores.min()), + "max_manipulability": float(valid_scores.max()), + "num_valid_points": int(len(valid_scores)), + } + + # Compute isotropy if requested + if self.config.compute_isotropy and jacobians is not None: + condition_numbers = self._compute_condition_numbers(jacobians) + self.results["mean_condition"] = float(condition_numbers.mean()) + self.results["std_condition"] = float(condition_numbers.std()) + + return self.results + + def _compute_manipulability_index(self, jacobians: np.ndarray) -> np.ndarray: + """Compute Yoshikawa manipulability index. + + Args: + jacobians: Jacobian matrices, shape (N, 6, num_joints). + + Returns: + Manipulability indices, shape (N,). + """ + # Manipulability index: sqrt(det(J * J^T)) + manipulability = np.zeros(len(jacobians)) + + for i, J in enumerate(jacobians): + JJT = J @ J.T + det = np.linalg.det(JJT) + manipulability[i] = np.sqrt(max(det, 0)) + + return manipulability + + def _compute_condition_numbers(self, jacobians: np.ndarray) -> np.ndarray: + """Compute condition numbers of Jacobian matrices. + + Args: + jacobians: Jacobian matrices, shape (N, 6, num_joints). + + Returns: + Condition numbers, shape (N,). + """ + condition_numbers = np.zeros(len(jacobians)) + + for i, J in enumerate(jacobians): + try: + condition_numbers[i] = np.linalg.cond(J) + except np.linalg.LinAlgError: + # Singular matrix, use infinity as condition number + condition_numbers[i] = np.inf + + return condition_numbers diff --git a/embodichain/lab/sim/utility/workspace_analyzer/metrics/reachability_metric.py b/embodichain/lab/sim/utility/workspace_analyzer/metrics/reachability_metric.py new file mode 100644 index 0000000..4414d77 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/metrics/reachability_metric.py @@ -0,0 +1,137 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from typing import Dict, Any +import numpy as np +from embodichain.lab.sim.utility.workspace_analyzer.metrics.base_metric import ( + BaseMetric, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.metric_config import ( + ReachabilityConfig, +) + + +class ReachabilityMetric(BaseMetric): + """Reachability metric for workspace analysis. + + Computes reachable workspace volume and coverage statistics. + """ + + def __init__(self, config: ReachabilityConfig | None = None): + """Initialize reachability metric. + + Args: + config: Reachability configuration. + """ + super().__init__(config or ReachabilityConfig()) + + def compute( + self, + workspace_points: np.ndarray, + joint_configurations: np.ndarray | None = None, + **kwargs, + ) -> Dict[str, Any]: + """Compute reachability metrics. + + Args: + workspace_points: Workspace points in Cartesian space, shape (N, 3). + joint_configurations: Joint configurations, shape (N, num_joints). + **kwargs: Additional arguments. + + Returns: + Dictionary containing: + - volume: Estimated reachable volume (m³) + - num_voxels: Number of occupied voxels + - coverage: Coverage percentage relative to bounding box + - bounding_box: Min and max bounds + - centroid: Center of workspace + """ + points = self._to_numpy(workspace_points) + + if len(points) == 0: + return { + "volume": 0.0, + "num_voxels": 0, + "coverage": 0.0, + "bounding_box": {"min": None, "max": None}, + "centroid": None, + } + + # Compute bounding box + min_bounds = points.min(axis=0) + max_bounds = points.max(axis=0) + + # Compute centroid + centroid = points.mean(axis=0) + + # Voxelize the workspace + voxel_size = self.config.voxel_size + voxel_grid = self._voxelize_points(points, voxel_size) + + # Count occupied voxels + num_voxels = len(voxel_grid) + + # Estimate volume + volume = num_voxels * (voxel_size**3) + + # Compute coverage if requested + coverage = 0.0 + if self.config.compute_coverage: + # Bounding box volume + bbox_volume = np.prod(max_bounds - min_bounds) + if bbox_volume > 0: + coverage = (volume / bbox_volume) * 100.0 + + self.results = { + "volume": float(volume), + "num_voxels": int(num_voxels), + "coverage": float(coverage), + "bounding_box": { + "min": min_bounds.tolist(), + "max": max_bounds.tolist(), + }, + "centroid": centroid.tolist(), + "voxel_size": voxel_size, + } + + return self.results + + def _voxelize_points( + self, points: np.ndarray, voxel_size: float + ) -> Dict[tuple, int]: + """Convert points to voxel grid. + + Args: + points: Point cloud, shape (N, 3). + voxel_size: Size of each voxel. + + Returns: + Dictionary mapping voxel coordinates to point count. + """ + # Convert points to voxel indices + voxel_indices = np.floor(points / voxel_size).astype(int) + + # Count points in each voxel + voxel_grid = {} + for idx in voxel_indices: + key = tuple(idx) + voxel_grid[key] = voxel_grid.get(key, 0) + 1 + + # Filter by minimum points threshold + min_points = self.config.min_points_per_voxel + voxel_grid = {k: v for k, v in voxel_grid.items() if v >= min_points} + + return voxel_grid diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/__init__.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/__init__.py new file mode 100644 index 0000000..b9a8106 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/__init__.py @@ -0,0 +1,64 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from embodichain.lab.sim.utility.workspace_analyzer.samplers.base_sampler import ( + BaseSampler, + ISampler, +) + +from embodichain.lab.sim.utility.workspace_analyzer.samplers.iniform_sampler import ( + UniformSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.random_sampler import ( + RandomSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.halton_sampler import ( + HaltonSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.sobol_sampler import ( + SobolSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.lhs_sampler import ( + LatinHypercubeSampler, +) + +from embodichain.lab.sim.utility.workspace_analyzer.samplers.sampler_factory import ( + SamplerFactory, + create_sampler, +) + +# Note: Geometric constraints are temporarily disabled + +# Note: PlaneSampler is intentionally not included here as it's a high-level +# specialized sampler, not a base sampler. Import it directly when needed: +# from embodichain.lab.sim.utility.workspace_analyzer.samplers.plane_sampler import PlaneSampler + +__all__ = [ + # Base samplers and interfaces + "BaseSampler", + "ISampler", + # Concrete base samplers + "UniformSampler", + "RandomSampler", + "HaltonSampler", + "SobolSampler", + "LatinHypercubeSampler", + # Note: Geometric constraints temporarily disabled + # Factory for creating base samplers + "SamplerFactory", + "create_sampler", + # Note: Constraint-based convenience functions temporarily disabled +] diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/base_sampler.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/base_sampler.py new file mode 100644 index 0000000..a489549 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/base_sampler.py @@ -0,0 +1,242 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from abc import ABC, abstractmethod +from typing import Protocol, Union, TYPE_CHECKING + +from embodichain.utils import logger + + +__all__ = [ + "ISampler", + "BaseSampler", +] + + +class ISampler(Protocol): + """Interface for all samplers. + + This protocol defines the contract that all samplers must follow. + """ + + def sample( + self, bounds: Union[torch.Tensor, np.ndarray], num_samples: int + ) -> torch.Tensor: + """Generate samples within the given bounds. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds for each dimension. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing the sampled points. + """ + ... + + def get_strategy_name(self) -> str: + """Get the name of the sampling strategy. + + Returns: + String identifier for the sampling strategy. + """ + ... + + +class BaseSampler(ABC): + """Abstract base class for all samplers. + + This class provides common functionality and enforces the implementation + of the sampling method in all derived classes. + + Attributes: + seed: Random seed for reproducibility. + rng: NumPy random number generator. + device: PyTorch device for tensor operations. + """ + + def __init__( + self, + seed: int = 42, + device: torch.device | None = None, + ): + """Initialize the base sampler. + + Args: + seed: Random seed for reproducibility. Defaults to 42. + device: PyTorch device (cpu/cuda). Defaults to cpu. + """ + self.seed = seed + self.rng = np.random.RandomState(seed) + self.device = device if device is not None else torch.device("cpu") + + # Set torch seed + torch.manual_seed(seed) + if self.device.type == "cuda": + torch.cuda.manual_seed(seed) + + def sample( + self, num_samples: int, bounds: torch.Tensor | np.ndarray | None = None + ) -> torch.Tensor: + """Generate samples within the given bounds. + + Args: + num_samples: Number of samples to generate. + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds for each dimension. + + Returns: + Tensor of shape (num_samples, n_dims) containing the sampled points. + + Raises: + ValueError: If bounds are not provided. + NotImplementedError: If the method is not implemented in the derived class. + """ + if bounds is None: + raise ValueError("bounds parameter is required") + return self._sample_from_bounds(bounds, num_samples) + + @abstractmethod + def _sample_from_bounds( + self, bounds: Union[torch.Tensor, np.ndarray], num_samples: int + ) -> torch.Tensor: + """Generate samples within the given bounds. + + This method must be implemented by all derived classes. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds for each dimension. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing the sampled points. + + Raises: + NotImplementedError: If the method is not implemented in the derived class. + """ + raise NotImplementedError( + "Subclasses must implement the _sample_from_bounds method" + ) + + @abstractmethod + def get_strategy_name(self) -> str: + """Get the name of the sampling strategy. + + Returns: + String identifier for the sampling strategy. + """ + raise NotImplementedError("Subclasses must implement get_strategy_name method") + + def _to_tensor(self, data: Union[torch.Tensor, np.ndarray]) -> torch.Tensor: + """Convert data to torch.Tensor. + + Args: + data: Input data (numpy array or torch tensor). + + Returns: + PyTorch tensor on the configured device. + """ + if isinstance(data, torch.Tensor): + return data.to(self.device) + return torch.tensor(data, dtype=torch.float32, device=self.device) + + def _to_numpy(self, data: torch.Tensor) -> np.ndarray: + """Convert tensor to numpy array. + + Args: + data: PyTorch tensor. + + Returns: + NumPy array. + """ + return data.detach().cpu().numpy() + + def _validate_bounds(self, bounds: Union[torch.Tensor, np.ndarray]) -> torch.Tensor: + """Validate the bounds array and convert to tensor. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds. + + Returns: + Validated bounds as torch.Tensor. + + Raises: + ValueError: If bounds are invalid. + """ + bounds_tensor = self._to_tensor(bounds) + + if bounds_tensor.ndim != 2 or bounds_tensor.shape[1] != 2: + raise ValueError( + f"Bounds must have shape (n_dims, 2), got {bounds_tensor.shape}" + ) + + if torch.any(bounds_tensor[:, 0] >= bounds_tensor[:, 1]): + raise ValueError( + "Lower bounds must be strictly less than upper bounds. " + f"Got bounds: {bounds_tensor}" + ) + + return bounds_tensor + + def _validate_samples(self, samples: torch.Tensor, bounds: torch.Tensor) -> None: + """Validate that samples are within bounds. + + Args: + samples: Tensor of shape (num_samples, n_dims) containing sampled points. + bounds: Tensor of shape (n_dims, 2) containing [lower, upper] bounds. + + Raises: + ValueError: If any sample is outside the bounds. + """ + lower_bounds = bounds[:, 0] + upper_bounds = bounds[:, 1] + + # Check if all samples are within bounds (with small tolerance for numerical errors) + tolerance = 1e-6 + if torch.any(samples < lower_bounds - tolerance) or torch.any( + samples > upper_bounds + tolerance + ): + out_of_bounds = torch.logical_or( + samples < lower_bounds - tolerance, samples > upper_bounds + tolerance + ) + num_violations = torch.sum(out_of_bounds).item() + logger.log_warning( + f"Found {num_violations} samples outside bounds. " + "This may be due to numerical precision issues." + ) + + def _scale_samples( + self, samples: torch.Tensor, bounds: torch.Tensor + ) -> torch.Tensor: + """Scale samples from [0, 1] to the given bounds. + + This is a utility method for samplers that generate samples in [0, 1]^n + and then scale them to the desired bounds. + + Args: + samples: Tensor of shape (num_samples, n_dims) with values in [0, 1]. + bounds: Tensor of shape (n_dims, 2) containing [lower, upper] bounds. + + Returns: + Scaled samples within the specified bounds. + """ + lower_bounds = bounds[:, 0] + upper_bounds = bounds[:, 1] + return lower_bounds + samples * (upper_bounds - lower_bounds) + + def __repr__(self) -> str: + """String representation of the sampler.""" + return f"{self.__class__.__name__}(strategy={self.get_strategy_name()}, seed={self.seed})" diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/gaussian_sampler.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/gaussian_sampler.py new file mode 100644 index 0000000..9750142 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/gaussian_sampler.py @@ -0,0 +1,265 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Union + +from embodichain.lab.sim.utility.workspace_analyzer.configs.sampling_config import ( + SamplingStrategy, +) + +from embodichain.lab.sim.utility.workspace_analyzer.samplers.base_sampler import ( + BaseSampler, +) + +__all__ = ["GaussianSampler"] + + +class GaussianSampler(BaseSampler): + """Gaussian (Normal distribution) sampler. + + This sampler generates samples from a multivariate Gaussian (normal) distribution, + which is useful for local workspace exploration around a specific point or for + generating samples with controlled spread. + + Advantages: + - Natural for local exploration + - Concentrates samples near center + - Smooth density falloff + - Good for sensitivity analysis + - Matches many physical phenomena + + Disadvantages: + - May not cover entire space well + - Can generate samples outside bounds (clipped) + - Not uniform - biased toward center + - Parameters require tuning + + Attributes: + mean: Center point of the distribution (n_dims,). If None, uses bounds center. + std: Standard deviation(s). Can be scalar (isotropic) or per-dimension (n_dims,). + clip_to_bounds: Whether to clip samples to bounds. If False, rejects out-of-bounds. + """ + + def __init__( + self, + seed: int = 42, + device: torch.device | None = None, + mean: torch.Tensor | np.ndarray | None = None, + std: Union[float, torch.Tensor, np.ndarray] = 0.3, + clip_to_bounds: bool = True, + ): + """Initialize the Gaussian sampler. + + Args: + seed: Random seed for reproducibility. Defaults to 42. + device: PyTorch device (cpu/cuda). Defaults to cpu. + mean: Mean of the distribution (center point). If None, uses midpoint of bounds. + Can be Tensor or ndarray of shape (n_dims,). + std: Standard deviation(s). Defaults to 0.3. + - Scalar: Isotropic (same std for all dimensions) + - Tensor/array (n_dims,): Per-dimension std + clip_to_bounds: Whether to clip samples to bounds. Defaults to True. + If False, regenerates out-of-bounds samples (slower but unbiased). + + """ + super().__init__(seed, device) + self.mean = self._to_tensor(mean) if mean is not None else None + + if isinstance(std, (int, float)): + self.std = float(std) + else: + self.std = self._to_tensor(std) + + self.clip_to_bounds = clip_to_bounds + + def _sample_from_bounds( + self, bounds: Union[torch.Tensor, np.ndarray], num_samples: int + ) -> torch.Tensor: + """Generate Gaussian-distributed samples within the given bounds. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing the sampled points. + + Raises: + ValueError: If bounds are invalid or num_samples is non-positive. + + Examples: + >>> # Sample around center with std=0.2 + >>> sampler = GaussianSampler(seed=42, std=0.2, clip_to_bounds=True) + >>> bounds = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]], dtype=torch.float32) + >>> samples = sampler.sample(bounds, num_samples=100) + >>> + >>> # Sample around custom point + >>> mean = torch.tensor([0.5, 0.5], dtype=torch.float32) + >>> sampler = GaussianSampler(mean=mean, std=0.15) + >>> samples = sampler.sample(bounds, num_samples=100) + """ + bounds = self._validate_bounds(bounds) + + if num_samples <= 0: + raise ValueError(f"num_samples must be positive, got {num_samples}") + + n_dims = bounds.shape[0] + + # Determine mean (center of distribution) + if self.mean is None: + # Use center of bounds + mean = (bounds[:, 0] + bounds[:, 1]) / 2.0 + else: + mean = self.mean + if mean.shape[0] != n_dims: + raise ValueError( + f"Mean dimension ({mean.shape[0]}) doesn't match bounds ({n_dims})" + ) + + # Determine std + if isinstance(self.std, float): + std = torch.full((n_dims,), self.std, device=self.device) + else: + std = self.std + if std.shape[0] != n_dims: + raise ValueError( + f"Std dimension ({std.shape[0]}) doesn't match bounds ({n_dims})" + ) + + if self.clip_to_bounds: + # Generate samples and clip + samples = self._generate_and_clip(mean, std, bounds, num_samples) + else: + # Rejection sampling to avoid clipping bias + samples = self._generate_with_rejection(mean, std, bounds, num_samples) + + # Validate samples + self._validate_samples(samples, bounds) + + return samples + + def _generate_and_clip( + self, + mean: torch.Tensor, + std: torch.Tensor, + bounds: torch.Tensor, + num_samples: int, + ) -> torch.Tensor: + """Generate Gaussian samples and clip to bounds. + + Args: + mean: Mean vector (n_dims,). + std: Standard deviation vector (n_dims,). + bounds: Bounds tensor (n_dims, 2). + num_samples: Number of samples to generate. + + Returns: + Clipped samples (num_samples, n_dims). + """ + # Generate Gaussian samples + samples = torch.randn(num_samples, mean.shape[0], device=self.device) + samples = mean + samples * std + + # Clip to bounds + lower_bounds = bounds[:, 0] + upper_bounds = bounds[:, 1] + samples = torch.clamp(samples, lower_bounds, upper_bounds) + + return samples + + def _generate_with_rejection( + self, + mean: torch.Tensor, + std: torch.Tensor, + bounds: torch.Tensor, + num_samples: int, + ) -> torch.Tensor: + """Generate Gaussian samples with rejection sampling for out-of-bounds. + + This avoids the bias introduced by clipping but is slower. + + Args: + mean: Mean vector (n_dims,). + std: Standard deviation vector (n_dims,). + bounds: Bounds tensor (n_dims, 2). + num_samples: Number of samples to generate. + + Returns: + Valid samples within bounds (num_samples, n_dims). + """ + lower_bounds = bounds[:, 0] + upper_bounds = bounds[:, 1] + + accepted_samples = [] + max_iterations = 100 + iteration = 0 + + while len(accepted_samples) < num_samples and iteration < max_iterations: + # Generate more samples than needed to reduce iterations + num_needed = num_samples - len(accepted_samples) + num_generate = max(num_needed * 2, 100) # Generate 2x to reduce rejections + + # Generate Gaussian samples + samples = torch.randn(num_generate, mean.shape[0], device=self.device) + samples = mean + samples * std + + # Check which samples are within bounds + within_bounds = torch.all( + (samples >= lower_bounds) & (samples <= upper_bounds), dim=1 + ) + + valid_samples = samples[within_bounds] + accepted_samples.append(valid_samples) + + iteration += 1 + + if ( + len(accepted_samples) == 0 + or sum(s.shape[0] for s in accepted_samples) < num_samples + ): + raise RuntimeError( + f"Failed to generate {num_samples} samples within bounds after {max_iterations} iterations. " + "Consider using clip_to_bounds=True or adjusting mean/std parameters." + ) + + # Concatenate and trim to exact size + all_samples = torch.cat(accepted_samples, dim=0) + return all_samples[:num_samples] + + def get_strategy_name(self) -> str: + """Get the name of the sampling strategy. + + Returns: + String identifier for the sampling strategy. + """ + return SamplingStrategy.GAUSSIAN.value + + def __repr__(self) -> str: + """String representation of the sampler.""" + mean_str = "auto" if self.mean is None else f"shape={self.mean.shape}" + std_str = ( + f"{self.std}" if isinstance(self.std, float) else f"shape={self.std.shape}" + ) + return ( + f"{self.__class__.__name__}(" + f"strategy={self.get_strategy_name()}, " + f"mean={mean_str}, " + f"std={std_str}, " + f"clip={self.clip_to_bounds}, " + f"seed={self.seed})" + ) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/halton_sampler.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/halton_sampler.py new file mode 100644 index 0000000..94559f2 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/halton_sampler.py @@ -0,0 +1,281 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import List, Union, TYPE_CHECKING + +from embodichain.lab.sim.utility.workspace_analyzer.configs.sampling_config import ( + SamplingStrategy, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.base_sampler import ( + BaseSampler, +) + +__all__ = ["HaltonSampler"] + + +class HaltonSampler(BaseSampler): + """Halton sequence sampler using quasi-random low-discrepancy sequences. + + The Halton sequence is a deterministic low-discrepancy sequence that provides + better coverage than random sampling. It uses coprime bases (primes) for each + dimension to generate well-distributed points. + + Advantages: + - Better uniformity than random sampling + - Deterministic (reproducible) + - Fast convergence for Monte Carlo integration + - Works well in low to medium dimensions (2-10) + + Disadvantages: + - Performance degrades in high dimensions (>10) + - Correlation between dimensions can appear + - Sequential generation (not easily parallelizable) + + Attributes: + bases: Prime bases for each dimension. If None, uses first n primes. + skip: Number of initial samples to skip (helps reduce correlation). + """ + + # First 100 prime numbers for Halton bases + _PRIMES = [ + 2, + 3, + 5, + 7, + 11, + 13, + 17, + 19, + 23, + 29, + 31, + 37, + 41, + 43, + 47, + 53, + 59, + 61, + 67, + 71, + 73, + 79, + 83, + 89, + 97, + 101, + 103, + 107, + 109, + 113, + 127, + 131, + 137, + 139, + 149, + 151, + 157, + 163, + 167, + 173, + 179, + 181, + 191, + 193, + 197, + 199, + 211, + 223, + 227, + 229, + 233, + 239, + 241, + 251, + 257, + 263, + 269, + 271, + 277, + 281, + 283, + 293, + 307, + 311, + 313, + 317, + 331, + 337, + 347, + 349, + 353, + 359, + 367, + 373, + 379, + 383, + 389, + 397, + 401, + 409, + 419, + 421, + 431, + 433, + 439, + 443, + 449, + 457, + 461, + 463, + 467, + 479, + 487, + 491, + 499, + 503, + 509, + 521, + 523, + 541, + ] + + def __init__( + self, + seed: int = 42, + device: torch.device | None = None, + bases: List[int] | None = None, + skip: int = 0, + ): + """Initialize the Halton sampler. + + Args: + seed: Random seed (used for consistency, but Halton is deterministic). + device: PyTorch device (cpu/cuda). Defaults to cpu. + bases: List of prime bases for each dimension. If None, uses first n primes. + skip: Number of initial samples to skip. Defaults to 0. + Higher values (e.g., 100-1000) can improve distribution quality. + + """ + super().__init__(seed, device) + self.bases = bases + self.skip = skip + + def sample( + self, bounds: Union[torch.Tensor, np.ndarray], num_samples: int + ) -> torch.Tensor: + """Generate Halton sequence samples within the given bounds. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing the sampled points. + + Raises: + ValueError: If bounds are invalid or num_samples is non-positive. + + Examples: + >>> sampler = HaltonSampler(skip=100) + >>> bounds = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]], dtype=torch.float32) + >>> samples = sampler.sample(bounds, num_samples=100) + >>> samples.shape + torch.Size([100, 2]) + """ + bounds = self._validate_bounds(bounds) + + if num_samples <= 0: + raise ValueError(f"num_samples must be positive, got {num_samples}") + + n_dims = bounds.shape[0] + + # Get bases for each dimension + if self.bases is None: + if n_dims > len(self._PRIMES): + raise ValueError( + f"Number of dimensions ({n_dims}) exceeds available primes ({len(self._PRIMES)}). " + "Please provide custom bases." + ) + bases = self._PRIMES[:n_dims] + else: + if len(self.bases) < n_dims: + raise ValueError( + f"Number of bases ({len(self.bases)}) is less than dimensions ({n_dims})" + ) + bases = self.bases[:n_dims] + + # Generate Halton sequence + samples_unit = np.zeros((num_samples, n_dims), dtype=np.float32) + + for dim in range(n_dims): + base = bases[dim] + for i in range(num_samples): + index = i + self.skip + 1 # Start from 1, apply skip + samples_unit[i, dim] = self._halton_number(index, base) + + # Convert to tensor and scale to bounds + samples_unit_tensor = self._to_tensor(samples_unit) + samples = self._scale_samples(samples_unit_tensor, bounds) + + # Validate samples + self._validate_samples(samples, bounds) + + return samples + + @staticmethod + def _halton_number(index: int, base: int) -> float: + """Compute a single Halton number. + + The Halton sequence is generated by reversing the base-n representation + of the index. + + Args: + index: Sequence index (starting from 1). + base: Prime base for this dimension. + + Returns: + Halton number in [0, 1]. + """ + result = 0.0 + f = 1.0 / base + i = index + + while i > 0: + result += f * (i % base) + i //= base + f /= base + + return result + + def get_strategy_name(self) -> str: + """Get the name of the sampling strategy. + + Returns: + String identifier for the sampling strategy. + """ + return SamplingStrategy.HALTON.value + + def __repr__(self) -> str: + """String representation of the sampler.""" + return ( + f"{self.__class__.__name__}(" + f"strategy={self.get_strategy_name()}, " + f"skip={self.skip}, " + f"seed={self.seed})" + ) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/importance_sampler.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/importance_sampler.py new file mode 100644 index 0000000..ee7ec24 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/importance_sampler.py @@ -0,0 +1,274 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Callable, Union, TYPE_CHECKING + +from embodichain.lab.sim.utility.workspace_analyzer.configs.sampling_config import ( + SamplingStrategy, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.base_sampler import ( + BaseSampler, +) + +from embodichain.utils import logger + +__all__ = ["ImportanceSampler"] + + +class ImportanceSampler(BaseSampler): + """Importance sampler using weighted distribution. + + Importance sampling generates samples according to a user-defined importance + (weight) function, concentrating samples in regions of interest. This is + particularly useful for workspace analysis where certain regions (e.g., task space, + reachable space) are more important than others. + + Advantages: + - Focuses samples on important regions + - More efficient for targeted analysis + - Reduces variance in Monte Carlo estimation + - Flexible - user defines importance function + + Disadvantages: + - Requires domain knowledge for weight function + - May miss unexpected important regions + - Can be computationally expensive if weight function is complex + - Quality depends heavily on weight function design + + Attributes: + weight_fn: Callable that takes positions (N, n_dims) and returns weights (N,). + Higher weights = higher probability of sampling. + num_candidates: Number of candidate samples to generate for rejection sampling. + Higher values = better approximation but slower. + method: Sampling method ('rejection' or 'transform'). + """ + + def __init__( + self, + weight_fn: Callable[[torch.Tensor], torch.Tensor], + seed: int = 42, + device: torch.device | None = None, + num_candidates: int = 10, + method: str = "rejection", + ): + """Initialize the importance sampler. + + Args: + weight_fn: Function that computes importance weights for positions. + Signature: weight_fn(positions: Tensor[N, n_dims]) -> Tensor[N] + Returns non-negative weights (higher = more important). + seed: Random seed for reproducibility. Defaults to 42. + device: PyTorch device (cpu/cuda). Defaults to cpu. + num_candidates: Number of candidates per desired sample. Defaults to 10. + Higher values improve quality but increase computation. + method: Sampling method. Defaults to 'rejection'. + 'rejection': Rejection sampling (simple, unbiased) + 'transform': Inverse transform sampling (requires normalized weights) + constraint: Optional geometric constraint for sampling (e.g., SphereConstraint). + """ + super().__init__(seed, device) + self.weight_fn = weight_fn + self.num_candidates = num_candidates + self.method = method + + if method not in ["rejection", "transform"]: + raise ValueError( + f"Invalid method '{method}'. Use 'rejection' or 'transform'." + ) + + def sample( + self, bounds: Union[torch.Tensor, np.ndarray], num_samples: int + ) -> torch.Tensor: + """Generate importance-weighted samples within the given bounds. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing the sampled points. + + Raises: + ValueError: If bounds are invalid or num_samples is non-positive. + + Examples: + >>> # Define weight function (favor center) + >>> def center_weight(positions): + ... distances = torch.sqrt(torch.sum(positions**2, dim=1)) + ... return torch.exp(-distances) # Higher weight near origin + >>> + >>> sampler = ImportanceSampler(weight_fn=center_weight, num_candidates=20) + >>> bounds = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]], dtype=torch.float32) + >>> samples = sampler.sample(bounds, num_samples=100) + """ + bounds = self._validate_bounds(bounds) + + if num_samples <= 0: + raise ValueError(f"num_samples must be positive, got {num_samples}") + + if self.method == "rejection": + samples = self._rejection_sampling(bounds, num_samples) + else: # transform + samples = self._transform_sampling(bounds, num_samples) + + # Validate samples + self._validate_samples(samples, bounds) + + return samples + + def _rejection_sampling( + self, bounds: torch.Tensor, num_samples: int + ) -> torch.Tensor: + """Rejection sampling using importance weights. + + Algorithm: + 1. Generate candidate samples uniformly + 2. Compute importance weights for all candidates + 3. Accept/reject based on weights + 4. Repeat until enough samples + + Args: + bounds: Tensor of shape (n_dims, 2) containing [lower, upper] bounds. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing accepted samples. + """ + n_dims = bounds.shape[0] + accepted_samples = [] + + while len(accepted_samples) < num_samples: + # Generate candidate samples + num_needed = num_samples - len(accepted_samples) + num_candidates_batch = max(num_needed * self.num_candidates, 100) + + candidates = torch.rand(num_candidates_batch, n_dims, device=self.device) + candidates = self._scale_samples(candidates, bounds) + + # Compute weights + weights = self.weight_fn(candidates) + + if torch.any(weights < 0): + logger.log_warning( + "Weight function returned negative values. Using absolute values." + ) + weights = torch.abs(weights) + + # Normalize weights to [0, 1] for rejection + if weights.max() > 0: + weights_normalized = weights / weights.max() + else: + logger.log_warning( + "All weights are zero. Falling back to uniform sampling." + ) + weights_normalized = torch.ones_like(weights) + + # Rejection sampling + accept_probs = torch.rand(num_candidates_batch, device=self.device) + accepted_mask = accept_probs < weights_normalized + + accepted_batch = candidates[accepted_mask] + accepted_samples.append(accepted_batch) + + # Safety check to avoid infinite loop + if ( + len(accepted_samples) > 1000 + and sum(s.shape[0] for s in accepted_samples) < num_samples * 0.1 + ): + logger.log_warning( + "Rejection sampling is very inefficient. Consider adjusting weight function " + "or increasing num_candidates." + ) + + # Concatenate and trim to exact size + all_samples = torch.cat(accepted_samples, dim=0) + return all_samples[:num_samples] + + def _transform_sampling( + self, bounds: torch.Tensor, num_samples: int + ) -> torch.Tensor: + """Inverse transform sampling using importance weights. + + This method generates candidates and then selects based on cumulative weights. + More stable than rejection sampling but requires more memory. + + Args: + bounds: Tensor of shape (n_dims, 2) containing [lower, upper] bounds. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing sampled points. + """ + n_dims = bounds.shape[0] + + # Generate candidate samples + num_candidates_total = num_samples * self.num_candidates + candidates = torch.rand(num_candidates_total, n_dims, device=self.device) + candidates = self._scale_samples(candidates, bounds) + + # Compute weights + weights = self.weight_fn(candidates) + + if torch.any(weights < 0): + logger.log_warning( + "Weight function returned negative values. Using absolute values." + ) + weights = torch.abs(weights) + + # Normalize to probability distribution + if weights.sum() > 0: + probabilities = weights / weights.sum() + else: + logger.log_warning( + "All weights are zero. Falling back to uniform sampling." + ) + probabilities = ( + torch.ones(num_candidates_total, device=self.device) + / num_candidates_total + ) + + # Sample indices according to probabilities + try: + indices = torch.multinomial(probabilities, num_samples, replacement=False) + except RuntimeError: + # If probabilities are problematic, add small epsilon + probabilities = probabilities + 1e-10 + probabilities = probabilities / probabilities.sum() + indices = torch.multinomial(probabilities, num_samples, replacement=False) + + selected_samples = candidates[indices] + + return selected_samples + + def get_strategy_name(self) -> str: + """Get the name of the sampling strategy. + + Returns: + String identifier for the sampling strategy. + """ + return SamplingStrategy.IMPORTANCE.value + + def __repr__(self) -> str: + """String representation of the sampler.""" + return ( + f"{self.__class__.__name__}(" + f"strategy={self.get_strategy_name()}, " + f"method={self.method}, " + f"num_candidates={self.num_candidates}, " + f"seed={self.seed})" + ) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/iniform_sampler.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/iniform_sampler.py new file mode 100644 index 0000000..228d4d6 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/iniform_sampler.py @@ -0,0 +1,169 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Union, TYPE_CHECKING + +from embodichain.lab.sim.utility.workspace_analyzer.configs.sampling_config import ( + SamplingStrategy, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.base_sampler import ( + BaseSampler, +) + + +from embodichain.utils import logger + +# Note: Geometric constraint imports temporarily disabled + +__all__ = ["UniformSampler"] + + +class UniformSampler(BaseSampler): + """Uniform grid sampler. + + This sampler generates samples on a regular grid within the specified bounds. + It ensures even coverage of the entire space, but suffers from the curse of + dimensionality - the number of samples grows exponentially with the number + of dimensions. + + Note: Geometric constraint sampling is temporarily disabled. + + Attributes: + samples_per_dim: Number of samples to generate per dimension. When specified, + this controls the grid density and takes precedence over num_samples. + Total grid points = samples_per_dim^n_dims. + """ + + def __init__( + self, + seed: int = 42, + samples_per_dim: int | None = None, + device: torch.device | None = None, + ): + """Initialize the uniform sampler. + + Args: + seed: Random seed for reproducibility. Defaults to 42. + samples_per_dim: Fixed number of samples per dimension. If None, + will be calculated automatically from num_samples. Defaults to None. + device: PyTorch device for tensor operations. + """ + super().__init__(seed, device) + self.samples_per_dim = samples_per_dim + + def _sample_from_bounds( + self, bounds: Union[torch.Tensor, np.ndarray], num_samples: int + ) -> torch.Tensor: + """Generate uniform grid samples within the given bounds. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds for each dimension. + num_samples: Total number of samples to generate. This is used to calculate + samples_per_dim if not explicitly provided during initialization. + Note: The actual number of samples may differ slightly from this value + to maintain a uniform grid. + + Returns: + Tensor of shape (actual_num_samples, n_dims) containing the sampled points. + The actual number of samples will be samples_per_dim^n_dims. + + Raises: + ValueError: If bounds are invalid. + + Examples: + >>> sampler = UniformSampler(samples_per_dim=3) + >>> bounds = torch.tensor([[-1, 1], [-1, 1]], dtype=torch.float32) + >>> samples = sampler.sample(bounds, num_samples=10) + >>> samples.shape + torch.Size([9, 2]) # 3^2 = 9 samples + """ + bounds = self._validate_bounds(bounds) + + n_dims = bounds.shape[0] + + # Calculate samples per dimension if not provided + if self.samples_per_dim is None: + # Compute samples_per_dim to approximate the desired num_samples + samples_per_dim = max(2, int(np.ceil(num_samples ** (1.0 / n_dims)))) + else: + samples_per_dim = self.samples_per_dim + + actual_num_samples = samples_per_dim**n_dims + + if actual_num_samples != num_samples and self.samples_per_dim is None: + logger.log_info( + f"Uniform grid: requested {num_samples} samples, " + f"generating {actual_num_samples} samples " + f"({samples_per_dim}^{n_dims}) for uniform coverage." + ) + + # Create uniform grid for each dimension + samples = self._create_grid(bounds, samples_per_dim) + + # Validate samples + self._validate_samples(samples, bounds) + + return samples + + def _create_grid(self, bounds: torch.Tensor, samples_per_dim: int) -> torch.Tensor: + """Create a uniform grid of samples. + + Args: + bounds: Tensor of shape (n_dims, 2) containing [lower, upper] bounds. + samples_per_dim: Number of samples per dimension. + + Returns: + Tensor of shape (samples_per_dim^n_dims, n_dims) containing grid points. + """ + n_dims = bounds.shape[0] + + # Create linspace for each dimension + grids = [] + for i in range(n_dims): + grid = torch.linspace( + bounds[i, 0].item(), + bounds[i, 1].item(), + samples_per_dim, + device=self.device, + ) + grids.append(grid) + + # Create meshgrid and flatten + mesh = torch.meshgrid(*grids, indexing="ij") + samples = torch.stack([m.flatten() for m in mesh], dim=-1) + + return samples + + # Note: Constraint-based sampling methods temporarily disabled + + def get_strategy_name(self) -> str: + """Get the name of the sampling strategy. + + Returns: + String identifier for the sampling strategy. + """ + return SamplingStrategy.UNIFORM.value + + def __repr__(self) -> str: + """String representation of the sampler.""" + return ( + f"{self.__class__.__name__}(" + f"strategy={self.get_strategy_name()}, " + f"samples_per_dim={self.samples_per_dim}, " + f"seed={self.seed})" + ) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/lhs_sampler.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/lhs_sampler.py new file mode 100644 index 0000000..89175b0 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/lhs_sampler.py @@ -0,0 +1,226 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Union + +try: + from scipy.stats import qmc + + SCIPY_AVAILABLE = True +except ImportError: + SCIPY_AVAILABLE = False + +__all__ = ["LatinHypercubeSampler"] + +from embodichain.lab.sim.utility.workspace_analyzer.configs.sampling_config import ( + SamplingStrategy, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.base_sampler import ( + BaseSampler, +) + +from embodichain.utils import logger + + +class LatinHypercubeSampler(BaseSampler): + """Latin Hypercube Sampler (LHS) for stratified sampling. + + Latin Hypercube Sampling ensures that each dimension is divided into + equally probable intervals, with exactly one sample in each interval. + This provides better coverage with fewer samples compared to random sampling. + + Advantages: + - Excellent coverage with small sample sizes + - Each dimension is uniformly sampled + - No sample clustering in any dimension + - Works well in high dimensions + - Popular in experimental design and sensitivity analysis + + Disadvantages: + - Samples may align in projections (can be mitigated with optimization) + - Not deterministic across different sample sizes + - May have correlation between dimensions (unless optimized) + + Attributes: + strength: Strength of the LHS (1 or 2). Higher strength reduces correlation. + optimization: Optimization method ('random-cd', 'lloyd', None). + 'random-cd': Random coordinate descent (fast, good quality) + 'lloyd': Lloyd's algorithm (slower, better quality) + None: No optimization (fastest, may have correlation) + """ + + def __init__( + self, + seed: int = 42, + device: torch.device | None = None, + strength: int = 1, + optimization: str | None = "random-cd", + ): + """Initialize the Latin Hypercube sampler. + + Args: + seed: Random seed for reproducibility. Defaults to 42. + device: PyTorch device (cpu/cuda). Defaults to cpu. + strength: Strength of the LHS (1 or 2). Defaults to 1. + Strength 1: Standard LHS + Strength 2: Improved spacing (requires more samples) + optimization: Optimization method to reduce correlation. + 'random-cd': Fast, good quality (recommended) + 'lloyd': Better quality, slower + None: No optimization (fastest) + Defaults to 'random-cd'. + constraint: Optional geometric constraint for sampling (e.g., SphereConstraint). + """ + super().__init__(seed, device) + self.strength = strength + self.optimization = optimization + + if not SCIPY_AVAILABLE and optimization is not None: + logger.log_warning( + "scipy is not available. LHS optimization will be disabled. " + "Install scipy for optimized sampling: pip install scipy" + ) + self.optimization = None + + def sample( + self, bounds: Union[torch.Tensor, np.ndarray], num_samples: int + ) -> torch.Tensor: + """Generate Latin Hypercube samples within the given bounds. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing the sampled points. + + Raises: + ValueError: If bounds are invalid or num_samples is non-positive. + + Examples: + >>> sampler = LatinHypercubeSampler(seed=42, optimization='random-cd') + >>> bounds = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]], dtype=torch.float32) + >>> samples = sampler.sample(bounds, num_samples=50) + >>> samples.shape + torch.Size([50, 2]) + """ + bounds = self._validate_bounds(bounds) + + if num_samples <= 0: + raise ValueError(f"num_samples must be positive, got {num_samples}") + + n_dims = bounds.shape[0] + + # Check strength constraints + if self.strength == 2: + min_samples = n_dims + 1 + if num_samples < min_samples: + logger.log_warning( + f"Strength 2 LHS requires at least {min_samples} samples for {n_dims} dimensions. " + f"Got {num_samples}. Falling back to strength 1." + ) + strength = 1 + else: + strength = self.strength + else: + strength = self.strength + + # Generate Latin Hypercube samples + if SCIPY_AVAILABLE and self.optimization is not None: + samples_unit = self._generate_lhs_scipy(n_dims, num_samples, strength) + else: + samples_unit = self._generate_lhs_basic(n_dims, num_samples) + + # Convert to tensor and scale to bounds + samples_unit_tensor = self._to_tensor(samples_unit) + samples = self._scale_samples(samples_unit_tensor, bounds) + + # Validate samples + self._validate_samples(samples, bounds) + + return samples + + def _generate_lhs_scipy( + self, n_dims: int, num_samples: int, strength: int + ) -> np.ndarray: + """Generate optimized LHS using scipy. + + Args: + n_dims: Number of dimensions. + num_samples: Number of samples to generate. + strength: Strength of the LHS (1 or 2). + + Returns: + Array of shape (num_samples, n_dims) with values in [0, 1]. + """ + # Create LHS engine + lhs_engine = qmc.LatinHypercube( + d=n_dims, strength=strength, optimization=self.optimization, seed=self.seed + ) + + # Generate samples + samples = lhs_engine.random(n=num_samples) + + return samples.astype(np.float32) + + def _generate_lhs_basic(self, n_dims: int, num_samples: int) -> np.ndarray: + """Generate basic LHS without optimization. + + This is a simple implementation that doesn't optimize for correlation + but still provides the stratification property of LHS. + + Args: + n_dims: Number of dimensions. + num_samples: Number of samples to generate. + + Returns: + Array of shape (num_samples, n_dims) with values in [0, 1]. + """ + samples = np.zeros((num_samples, n_dims), dtype=np.float32) + + for dim in range(n_dims): + # Divide [0, 1] into num_samples equally-sized intervals + intervals = np.arange(num_samples, dtype=np.float32) / num_samples + + # Sample uniformly within each interval + samples[:, dim] = ( + intervals + self.rng.rand(num_samples).astype(np.float32) / num_samples + ) + + # Randomly permute to decorrelate dimensions + self.rng.shuffle(samples[:, dim]) + + return samples + + def get_strategy_name(self) -> str: + """Get the name of the sampling strategy. + + Returns: + String identifier for the sampling strategy. + """ + return SamplingStrategy.LATIN_HYPERCUBE.value + + def __repr__(self) -> str: + """String representation of the sampler.""" + return ( + f"{self.__class__.__name__}(" + f"strategy={self.get_strategy_name()}, " + f"strength={self.strength}, " + f"optimization={self.optimization}, " + f"seed={self.seed})" + ) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/random_sampler.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/random_sampler.py new file mode 100644 index 0000000..94122de --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/random_sampler.py @@ -0,0 +1,112 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Union, TYPE_CHECKING + +from embodichain.lab.sim.utility.workspace_analyzer.configs.sampling_config import ( + SamplingStrategy, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.base_sampler import ( + BaseSampler, +) + +__all__ = ["RandomSampler"] + + +class RandomSampler(BaseSampler): + """Random sampler using uniform distribution. + + This sampler generates samples uniformly at random within the specified bounds. + It's simple and fast, but doesn't guarantee uniform coverage. Samples may + cluster in some regions while leaving other regions sparse. + + Advantages: + - Fast and simple + - No assumptions about the space + - Works well for any number of dimensions + + Disadvantages: + - Uneven coverage (clusters and gaps) + - Slower convergence than quasi-random methods + - May miss important regions + """ + + def __init__( + self, + seed: int = 42, + device: torch.device | None = None, + ): + """Initialize the random sampler. + + Args: + seed: Random seed for reproducibility. Defaults to 42. + device: PyTorch device for tensor operations. + """ + super().__init__(seed, device) + + def _sample_from_bounds( + self, bounds: Union[torch.Tensor, np.ndarray], num_samples: int + ) -> torch.Tensor: + """Generate random samples within the given bounds. + + Samples are drawn from a uniform distribution within each dimension's bounds. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds for each dimension. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing the sampled points. + + Raises: + ValueError: If bounds are invalid or num_samples is non-positive. + + Examples: + >>> sampler = RandomSampler(seed=42) + >>> bounds = torch.tensor([[-1, 1], [-1, 1]], dtype=torch.float32) + >>> samples = sampler.sample(bounds, num_samples=100) + >>> samples.shape + torch.Size([100, 2]) + >>> torch.all(samples >= -1) and torch.all(samples <= 1) + True + """ + bounds = self._validate_bounds(bounds) + + if num_samples <= 0: + raise ValueError(f"num_samples must be positive, got {num_samples}") + + n_dims = bounds.shape[0] + + # Generate random samples in [0, 1]^n_dims using torch + samples_unit = torch.rand(num_samples, n_dims, device=self.device) + + # Scale to the actual bounds + samples = self._scale_samples(samples_unit, bounds) + + # Validate samples + self._validate_samples(samples, bounds) + + return samples + + def get_strategy_name(self) -> str: + """Get the name of the sampling strategy. + + Returns: + String identifier for the sampling strategy. + """ + return SamplingStrategy.RANDOM.value diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/sampler_factory.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/sampler_factory.py new file mode 100644 index 0000000..50001d4 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/sampler_factory.py @@ -0,0 +1,276 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations +from typing import Dict, Type, Any, Union +from threading import Lock + +import torch +import numpy as np + +from embodichain.lab.sim.utility.workspace_analyzer.configs.sampling_config import ( + SamplingStrategy, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.base_sampler import ( + BaseSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.iniform_sampler import ( + UniformSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.random_sampler import ( + RandomSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.halton_sampler import ( + HaltonSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.sobol_sampler import ( + SobolSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.lhs_sampler import ( + LatinHypercubeSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.importance_sampler import ( + ImportanceSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.gaussian_sampler import ( + GaussianSampler, +) + +# Note: Constraint imports temporarily disabled + +from embodichain.utils import logger + + +class SamplerFactory: + """Factory class for creating samplers (Singleton pattern). + + This factory allows registration and creation of samplers based on + the sampling strategy. It uses the singleton pattern to ensure only + one instance exists throughout the application. + + The factory comes pre-registered with built-in samplers: + - UNIFORM: UniformSampler + - RANDOM: RandomSampler + + Additional samplers can be registered using register_sampler(). + + Examples: + >>> factory = SamplerFactory() + >>> sampler = factory.create_sampler(SamplingStrategy.UNIFORM, seed=42) + >>> isinstance(sampler, UniformSampler) + True + + >>> # Register custom sampler + >>> factory.register_sampler("custom", CustomSampler) + >>> custom_sampler = factory.create_sampler("custom", seed=42) + """ + + _instance: SamplerFactory | None = None + _lock: Lock = Lock() + + def __new__(cls): + """Create or return the singleton instance. + + Returns: + The singleton SamplerFactory instance. + """ + if cls._instance is None: + with cls._lock: + # Double-checked locking + if cls._instance is None: + cls._instance = super(SamplerFactory, cls).__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self): + """Initialize the factory with built-in samplers. + + This method only runs once due to the singleton pattern. + """ + # Prevent re-initialization + if self._initialized: + return + + self._samplers: Dict[str, Type[BaseSampler]] = {} + self._register_builtin_samplers() + self._initialized = True + + def _register_builtin_samplers(self) -> None: + """Register the built-in samplers.""" + self._samplers[SamplingStrategy.UNIFORM.value] = UniformSampler + self._samplers[SamplingStrategy.RANDOM.value] = RandomSampler + self._samplers[SamplingStrategy.HALTON.value] = HaltonSampler + self._samplers[SamplingStrategy.SOBOL.value] = SobolSampler + self._samplers[SamplingStrategy.LATIN_HYPERCUBE.value] = LatinHypercubeSampler + self._samplers[SamplingStrategy.IMPORTANCE.value] = ImportanceSampler + self._samplers[SamplingStrategy.GAUSSIAN.value] = GaussianSampler + self._samplers[SamplingStrategy.SPHERE.value] = UniformSampler + + logger.log_debug(f"Registered built-in samplers: {list(self._samplers.keys())}") + + def register_sampler(self, name: str, sampler_class: Type[BaseSampler]) -> None: + """Register a new sampler class. + + Args: + name: String identifier for the sampler strategy. + sampler_class: The sampler class to register. Must inherit from BaseSampler. + + Raises: + TypeError: If sampler_class is not a subclass of BaseSampler. + ValueError: If name already exists and overwrite=False. + + Examples: + >>> factory = SamplerFactory() + >>> factory.register_sampler("my_sampler", MySamplerClass) + """ + if not issubclass(sampler_class, BaseSampler): + raise TypeError( + f"sampler_class must be a subclass of BaseSampler, " + f"got {sampler_class}" + ) + + if name in self._samplers: + logger.log_warning( + f"Sampler '{name}' already registered. Overwriting with {sampler_class.__name__}." + ) + + self._samplers[name] = sampler_class + logger.log_info(f"Registered sampler '{name}': {sampler_class.__name__}") + + def create_sampler( + self, strategy: SamplingStrategy | str | None = None, **kwargs: Any + ) -> BaseSampler: + """Create a sampler instance based on the strategy. + + Args: + strategy: The sampling strategy to use. Can be a SamplingStrategy enum + or a string identifier. If None, defaults to RANDOM. + **kwargs: Additional keyword arguments to pass to the sampler constructor. + Common arguments include: + - seed: Random seed for reproducibility + - samples_per_dim: For UniformSampler + - device: PyTorch device for tensor operations + Note: constraint parameter is temporarily disabled + + Returns: + An instance of the requested sampler. + + Raises: + ValueError: If the strategy is not registered. + + Examples: + >>> factory = SamplerFactory() + >>> # Bounds-based usage + >>> sampler = factory.create_sampler(SamplingStrategy.UNIFORM, seed=42) + >>> sampler = factory.create_sampler("random", seed=123) + + Note: Constraint-based sampling examples are temporarily disabled + """ + # Default to RANDOM if no strategy specified + if strategy is None: + strategy = SamplingStrategy.RANDOM + + # Convert enum to string if necessary + if isinstance(strategy, SamplingStrategy): + strategy_name = strategy.value + else: + strategy_name = strategy + + # Check if strategy is registered + if strategy_name not in self._samplers: + available = list(self._samplers.keys()) + raise ValueError( + f"Unknown sampling strategy: '{strategy_name}'. " + f"Available strategies: {available}. " + f"You can register a custom sampler using register_sampler()." + ) + + # Create and return sampler instance + sampler_class = self._samplers[strategy_name] + sampler = sampler_class(**kwargs) + + logger.log_info( + f"Created sampler: {sampler_class.__name__} with strategy '{strategy_name}'" + ) + + return sampler + + def list_available_strategies(self) -> list[str]: + """List all registered sampling strategies. + + Returns: + List of registered strategy names. + """ + return list(self._samplers.keys()) + + def is_registered(self, strategy: SamplingStrategy | str) -> bool: + """Check if a strategy is registered. + + Args: + strategy: The sampling strategy to check. + + Returns: + True if the strategy is registered, False otherwise. + """ + if isinstance(strategy, SamplingStrategy): + strategy_name = strategy.value + else: + strategy_name = strategy + + return strategy_name in self._samplers + + @classmethod + def reset_instance(cls) -> None: + """Reset the singleton instance (mainly for testing). + + Warning: + This should only be used in testing scenarios. + """ + with cls._lock: + cls._instance = None + + def __repr__(self) -> str: + """String representation of the factory.""" + strategies = self.list_available_strategies() + return f"SamplerFactory(strategies={strategies})" + + +# Convenience function for creating samplers +def create_sampler( + strategy: SamplingStrategy | str | None = None, **kwargs: Any +) -> BaseSampler: + """Convenience function to create a sampler. + + This is a shorthand for SamplerFactory().create_sampler(). + + Args: + strategy: The sampling strategy to use. + **kwargs: Additional keyword arguments to pass to the sampler constructor. + + Returns: + An instance of the requested sampler. + + Examples: + >>> sampler = create_sampler(SamplingStrategy.UNIFORM, seed=42) + >>> sampler = create_sampler("random", seed=123) + + Note: Constraint-based sampling is temporarily disabled + """ + factory = SamplerFactory() + return factory.create_sampler(strategy, **kwargs) + + +# Note: Constraint-based convenience functions have been removed diff --git a/embodichain/lab/sim/utility/workspace_analyzer/samplers/sobol_sampler.py b/embodichain/lab/sim/utility/workspace_analyzer/samplers/sobol_sampler.py new file mode 100644 index 0000000..de5f7e3 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/samplers/sobol_sampler.py @@ -0,0 +1,216 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Union + +try: + from scipy.stats import qmc + + SCIPY_AVAILABLE = True +except ImportError: + SCIPY_AVAILABLE = False + +from embodichain.lab.sim.utility.workspace_analyzer.configs.sampling_config import ( + SamplingStrategy, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers.base_sampler import ( + BaseSampler, +) + +from embodichain.utils import logger + + +class SobolSampler(BaseSampler): + """Sobol sequence sampler using quasi-random low-discrepancy sequences. + + The Sobol sequence is a low-discrepancy sequence that provides excellent + uniformity in high-dimensional spaces. It's widely used in finance, + engineering, and scientific computing for Monte Carlo simulations. + + Advantages: + - Excellent uniformity in high dimensions (up to ~40 dimensions) + - Industry standard for quasi-Monte Carlo methods + - Better convergence than random sampling (O(1/n) vs O(1/√n)) + - Well-suited for integration and optimization + + Disadvantages: + - Requires scipy library + - Sequential generation (but can be scrambled for randomization) + - Initial points may not be well-distributed (use skip parameter) + + Attributes: + scramble: Whether to scramble the sequence for better randomization. + skip: Number of initial samples to skip. + + Notes: + This implementation uses scipy.stats.qmc.Sobol for efficient generation. + Falls back to a basic implementation if scipy is not available. + """ + + def __init__( + self, + seed: int = 42, + device: torch.device | None = None, + scramble: bool = True, + skip: int = 0, + ): + """Initialize the Sobol sampler. + + Args: + seed: Random seed for scrambling. Defaults to 42. + device: PyTorch device (cpu/cuda). Defaults to cpu. + scramble: Whether to scramble the sequence. Defaults to True. + Scrambling improves randomization while maintaining low discrepancy. + skip: Number of initial samples to skip. Defaults to 0. + Recommended: 0 for scrambled, >0 (e.g., 100) for unscrambled. + constraint: Optional geometric constraint for sampling (e.g., SphereConstraint). + """ + super().__init__(seed, device) + self.scramble = scramble + self.skip = skip + + if not SCIPY_AVAILABLE: + logger.log_warning( + "scipy is not available. Sobol sampler will use a basic fallback implementation. " + "For optimal performance, install scipy: pip install scipy" + ) + + def sample( + self, bounds: Union[torch.Tensor, np.ndarray], num_samples: int + ) -> torch.Tensor: + """Generate Sobol sequence samples within the given bounds. + + Args: + bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds. + num_samples: Number of samples to generate. + + Returns: + Tensor of shape (num_samples, n_dims) containing the sampled points. + + Raises: + ValueError: If bounds are invalid or num_samples is non-positive. + + Examples: + >>> sampler = SobolSampler(scramble=True, seed=42) + >>> bounds = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]], dtype=torch.float32) + >>> samples = sampler.sample(bounds, num_samples=100) + >>> samples.shape + torch.Size([100, 2]) + """ + bounds = self._validate_bounds(bounds) + + if num_samples <= 0: + raise ValueError(f"num_samples must be positive, got {num_samples}") + + n_dims = bounds.shape[0] + + if n_dims > 21201: # Maximum dimension for scipy's Sobol + raise ValueError( + f"Sobol sequence supports up to 21201 dimensions, got {n_dims}" + ) + + # Generate Sobol sequence + if SCIPY_AVAILABLE: + samples_unit = self._generate_sobol_scipy(n_dims, num_samples) + else: + samples_unit = self._generate_sobol_fallback(n_dims, num_samples) + + # Convert to tensor and scale to bounds + samples_unit_tensor = self._to_tensor(samples_unit) + samples = self._scale_samples(samples_unit_tensor, bounds) + + # Validate samples + self._validate_samples(samples, bounds) + + return samples + + def _generate_sobol_scipy(self, n_dims: int, num_samples: int) -> np.ndarray: + """Generate Sobol sequence using scipy. + + Args: + n_dims: Number of dimensions. + num_samples: Number of samples to generate. + + Returns: + Array of shape (num_samples, n_dims) with values in [0, 1]. + """ + # Create Sobol engine + sobol_engine = qmc.Sobol(d=n_dims, scramble=self.scramble, seed=self.seed) + + # Skip initial samples if requested + if self.skip > 0: + sobol_engine.fast_forward(self.skip) + + # Generate samples + samples = sobol_engine.random(n=num_samples) + + return samples.astype(np.float32) + + def _generate_sobol_fallback(self, n_dims: int, num_samples: int) -> np.ndarray: + """Fallback Sobol generator when scipy is not available. + + This is a basic implementation and may not match scipy's quality. + For production use, install scipy. + + Args: + n_dims: Number of dimensions. + num_samples: Number of samples to generate. + + Returns: + Array of shape (num_samples, n_dims) with values in [0, 1]. + """ + logger.log_warning( + "Using fallback Sobol generator. Results may not match scipy quality." + ) + + # Simple fallback: use stratified random sampling + # This is NOT a true Sobol sequence but provides reasonable coverage + samples = np.zeros((num_samples, n_dims), dtype=np.float32) + + for dim in range(n_dims): + # Divide [0, 1] into num_samples intervals + intervals = np.linspace(0, 1, num_samples + 1) + # Sample randomly within each interval + samples[:, dim] = intervals[:-1] + self.rng.rand(num_samples) / num_samples + + # Shuffle to reduce correlation + for dim in range(n_dims): + self.rng.shuffle(samples[:, dim]) + + return samples + + def get_strategy_name(self) -> str: + """Get the name of the sampling strategy. + + Returns: + String identifier for the sampling strategy. + """ + return SamplingStrategy.SOBOL.value + + def __repr__(self) -> str: + """String representation of the sampler.""" + return ( + f"{self.__class__.__name__}(" + f"strategy={self.get_strategy_name()}, " + f"scramble={self.scramble}, " + f"skip={self.skip}, " + f"seed={self.seed})" + ) + + +__all__ = ["SobolSampler"] diff --git a/embodichain/lab/sim/utility/workspace_analyzer/visualizers/__init__.py b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/__init__.py new file mode 100644 index 0000000..dec80b6 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/__init__.py @@ -0,0 +1,56 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.base_visualizer import ( + BaseVisualizer, + IVisualizer, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + VisualizationType, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.point_cloud_visualizer import ( + PointCloudVisualizer, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.voxel_visualizer import ( + VoxelVisualizer, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.sphere_visualizer import ( + SphereVisualizer, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.axis_visualizer import ( + AxisVisualizer, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.visualizer_factory import ( + VisualizerFactory, + create_visualizer, +) + +__all__ = [ + "BaseVisualizer", + "IVisualizer", + "VisualizationType", + "PointCloudVisualizer", + "VoxelVisualizer", + "SphereVisualizer", + "AxisVisualizer", + "VisualizerFactory", + "create_visualizer", +] diff --git a/embodichain/lab/sim/utility/workspace_analyzer/visualizers/axis_visualizer.py b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/axis_visualizer.py new file mode 100644 index 0000000..3b4c188 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/axis_visualizer.py @@ -0,0 +1,478 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Union, Dict, Any +from pathlib import Path + +try: + import open3d as o3d + + OPEN3D_AVAILABLE = True +except ImportError: + OPEN3D_AVAILABLE = False + +try: + import matplotlib.pyplot as plt + from mpl_toolkits.mplot3d import Axes3D + + MATPLOTLIB_AVAILABLE = True +except ImportError: + MATPLOTLIB_AVAILABLE = False + +from embodichain.utils import logger +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.base_visualizer import ( + BaseVisualizer, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.visualization_config import ( + VisualizationType, +) + +__all__ = ["AxisVisualizer"] + + +class AxisVisualizer(BaseVisualizer): + """Visualizer for coordinate axes/frames at specified poses. + + This visualizer creates coordinate axes (X, Y, Z) at given transformation matrices, + useful for visualizing robot end-effector poses, workspace reference frames, etc. + + Supports multiple backends: + - 'sim_manager': Uses SimulationManager.draw_marker() with MarkerCfg + - 'open3d': Creates coordinate frames using Open3D + - 'matplotlib': Draws axis lines in 3D matplotlib plot + - 'data': Returns axis data without visualization + """ + + def __init__( + self, + backend: str = "sim_manager", + axis_length: float = 0.15, + axis_size: float = 0.005, + config: Dict[str, Any] | None = None, + sim_manager: Any | None = None, + control_part_name: str | None = None, + reference_pose: Union[np.ndarray, torch.Tensor] | None = None, + arena_index: int = 0, + ): + """Initialize the axis visualizer. + + Args: + backend: Visualization backend ('sim_manager', 'open3d', 'matplotlib', or 'data'). + Defaults to 'sim_manager'. + axis_length: Length of each axis. Defaults to 0.15. + axis_size: Thickness/size of axes. Defaults to 0.005. + config: Optional configuration dictionary. Defaults to None. + sim_manager: SimulationManager instance for 'sim_manager' backend. Defaults to None. + control_part_name: Control part name for naming (compatibility). Defaults to None. + reference_pose: Reference pose (4x4 matrix) for orientation. Defaults to None. + arena_index: Arena index for sim_manager markers. Defaults to 0. + """ + super().__init__(backend, config) + self.axis_length = axis_length + self.axis_size = axis_size + self.sim_manager = sim_manager + self.control_part_name = control_part_name + self.reference_pose = reference_pose + self.arena_index = arena_index + + def visualize( + self, + poses: Union[torch.Tensor, np.ndarray], + colors: Union[torch.Tensor, np.ndarray] | None = None, + **kwargs: Any, + ) -> Any: + """Visualize coordinate axes at specified poses or points. + + Args: + poses: Either transformation matrices of shape (4, 4) or (N, 4, 4), + or point coordinates of shape (N, 3) which will be converted to poses. + colors: Optional colors (not used for axes but kept for interface compatibility). + **kwargs: Additional visualization parameters: + - axis_length: Override default axis length + - axis_size: Override default axis size + - name_prefix: Prefix for axis names (default: "axis") + - arena_index: Arena index for sim_manager (default: 0) + + Returns: + Backend-specific axis representation. + + Examples: + >>> import numpy as np + >>> visualizer = AxisVisualizer(backend='sim_manager', sim_manager=sim) + >>> # Using transformation matrix + >>> pose = np.eye(4) + >>> pose[:3, 3] = [1.0, 0.5, 1.2] # Set position + >>> result = visualizer.visualize(pose) + >>> + >>> # Using point coordinates + >>> points = np.array([[1.0, 0.5, 1.2], [2.0, 1.0, 0.8]]) + >>> result = visualizer.visualize(points) + >>> visualizer.show() + """ + # Convert to numpy + poses = self._to_numpy(poses) + + # Convert points to poses if needed + poses = self._convert_points_to_poses(poses) + + # Validate pose dimensions + self._validate_poses(poses) + + # Get parameters + axis_length = kwargs.get("axis_length", self.axis_length) + axis_size = kwargs.get("axis_size", self.axis_size) + name_prefix = kwargs.get("name_prefix", "axis") + arena_index = kwargs.get("arena_index", self.arena_index) + + # Dispatch to backend implementation + if self.backend == "sim_manager": + return self._visualize_sim_manager( + poses, axis_length, axis_size, name_prefix, arena_index, **kwargs + ) + elif self.backend == "open3d": + return self._visualize_open3d( + poses, axis_length, axis_size, name_prefix, **kwargs + ) + elif self.backend == "matplotlib": + return self._visualize_matplotlib( + poses, axis_length, axis_size, name_prefix, **kwargs + ) + elif self.backend == "data": + return self._visualize_data( + poses, axis_length, axis_size, name_prefix, **kwargs + ) + else: + raise ValueError(f"Unsupported backend: {self.backend}") + + def _visualize_sim_manager( + self, + poses: np.ndarray, + axis_length: float, + axis_size: float, + name_prefix: str, + arena_index: int, + **kwargs: Any, + ) -> Any: + """Visualize axes using sim_manager backend.""" + if self.sim_manager is None: + raise ValueError("sim_manager is required for 'sim_manager' backend") + + # Import here to avoid circular imports + from embodichain.lab.sim.cfg import MarkerCfg + + axis_markers = [] + + # Handle single pose (4,4) or multiple poses (N,4,4) + if poses.ndim == 2: + poses = poses[np.newaxis, :, :] + + for i, pose_matrix in enumerate(poses): + marker_name = f"{name_prefix}_{i}" + + try: + # Create axis marker using MarkerCfg + marker_cfg = MarkerCfg( + name=marker_name, + marker_type="axis", + axis_xpos=pose_matrix, # 4x4 transformation matrix + axis_size=axis_size, + axis_len=axis_length, + arena_index=arena_index, + ) + + # Draw the marker + self.sim_manager.draw_marker(cfg=marker_cfg) + axis_markers.append(marker_name) + + except Exception as e: + logger.log_warning(f"Failed to draw axis {marker_name}: {e}") + + logger.log_info(f"Created {len(axis_markers)} coordinate axes with sim_manager") + self._last_visualization = axis_markers + return axis_markers + + def _visualize_open3d( + self, + poses: np.ndarray, + axis_length: float, + axis_size: float, + name_prefix: str, + **kwargs: Any, + ) -> Any: + """Visualize axes using Open3D backend.""" + if not OPEN3D_AVAILABLE: + raise RuntimeError( + "Open3D is not available. Install with: pip install open3d" + ) + + # Create coordinate frame geometries + frame_geometries = [] + + if poses.ndim == 2: + poses = poses[np.newaxis, :, :] + + for i, pose_matrix in enumerate(poses): + frame = o3d.geometry.TriangleMesh.create_coordinate_frame( + size=axis_length, origin=[0, 0, 0] + ) + frame.transform(pose_matrix) + frame_geometries.append(frame) + + logger.log_info( + f"Created {len(frame_geometries)} coordinate frames with Open3D" + ) + self._last_visualization = frame_geometries + return frame_geometries + + def _visualize_matplotlib( + self, + poses: np.ndarray, + axis_length: float, + axis_size: float, + name_prefix: str, + **kwargs: Any, + ) -> Any: + """Visualize axes using matplotlib backend.""" + if not MATPLOTLIB_AVAILABLE: + raise RuntimeError( + "Matplotlib is not available. Install with: pip install matplotlib" + ) + + # Create figure and 3D axis + fig = plt.figure(figsize=(10, 8)) + ax = fig.add_subplot(111, projection="3d") + + if poses.ndim == 2: + poses = poses[np.newaxis, :, :] + + colors = ["red", "green", "blue"] # X, Y, Z axes + labels = ["X", "Y", "Z"] + + for i, pose_matrix in enumerate(poses): + origin = pose_matrix[:3, 3] + + for j, (color, label) in enumerate(zip(colors, labels)): + axis_direction = pose_matrix[:3, j] + axis_end = origin + axis_direction * axis_length + + ax.plot3D( + [origin[0], axis_end[0]], + [origin[1], axis_end[1]], + [origin[2], axis_end[2]], + color=color, + linewidth=axis_size * 1000, + label=f"{label}_axis_{i}" if i == 0 else "", + ) + + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.legend() + ax.set_title(f"Coordinate Axes ({len(poses)} frames)") + + logger.log_info(f"Created matplotlib plot with {len(poses)} coordinate frames") + self._last_visualization = fig + return fig + + def _visualize_data( + self, + poses: np.ndarray, + axis_length: float, + axis_size: float, + name_prefix: str, + **kwargs: Any, + ) -> Any: + """Return axis data without visualization.""" + axis_data = { + "poses": poses, + "axis_length": axis_length, + "axis_size": axis_size, + "name_prefix": name_prefix, + "type": "coordinate_axes", + "num_frames": len(poses) if poses.ndim == 3 else 1, + } + + logger.log_info( + f"Generated axis data for {axis_data['num_frames']} coordinate frames" + ) + self._last_visualization = axis_data + return axis_data + + def _convert_points_to_poses(self, data: np.ndarray) -> np.ndarray: + """Convert point coordinates to transformation matrices if needed. + + Args: + data: Either points (N, 3) or poses (4, 4) or (N, 4, 4) + + Returns: + Transformation matrices (4, 4) or (N, 4, 4) + """ + # Check if input is points (N, 3) + if data.ndim == 2 and data.shape[1] == 3: + # Convert points to poses + num_points = data.shape[0] + + # Use reference pose if available, otherwise identity rotation + if self.reference_pose is not None: + # Convert reference pose to numpy if needed + if isinstance(self.reference_pose, torch.Tensor): + ref_pose = self.reference_pose.cpu().numpy() + else: + ref_pose = self.reference_pose + + # Ensure 4x4 matrix + if ref_pose.ndim == 3: + ref_pose = ref_pose[0] # Take first pose if batch + + # Create poses with reference orientation and point positions + poses = np.tile(ref_pose, (num_points, 1, 1)) + poses[:, :3, 3] = data # Override translation with point positions + + logger.log_debug( + f"Using reference pose orientation for {num_points} coordinate axes" + ) + else: + # Fallback to identity rotation + poses = np.tile(np.eye(4), (num_points, 1, 1)) + poses[:, :3, 3] = data # Set translation + + logger.log_debug( + f"Using identity orientation for {num_points} coordinate axes" + ) + + return poses + + elif data.ndim == 1 and data.shape[0] == 3: + # Single point case + if self.reference_pose is not None: + # Convert reference pose to numpy if needed + if isinstance(self.reference_pose, torch.Tensor): + pose = self.reference_pose.cpu().numpy() + else: + pose = self.reference_pose.copy() + + # Ensure 2D matrix + if pose.ndim == 3: + pose = pose[0] + + # Override translation + pose[:3, 3] = data + else: + # Fallback to identity + pose = np.eye(4) + pose[:3, 3] = data + + return pose + else: + # Already poses, return as is + return data + + def _validate_poses(self, poses: np.ndarray) -> None: + """Validate pose array dimensions and values. + + Args: + poses: Array of transformation matrices + + Raises: + ValueError: If poses have invalid shape or values + """ + if poses.ndim not in [2, 3]: + raise ValueError( + f"Poses must be 2D (4,4) or 3D (N,4,4) array, got shape {poses.shape}" + ) + + if poses.ndim == 2: + if poses.shape != (4, 4): + raise ValueError(f"Single pose must be (4,4), got {poses.shape}") + else: # poses.ndim == 3 + if poses.shape[1:] != (4, 4): + raise ValueError(f"Multiple poses must be (N,4,4), got {poses.shape}") + + # Check if last row is [0, 0, 0, 1] for valid transformation matrices + if poses.ndim == 2: + poses_to_check = [poses] + else: + poses_to_check = poses + + for i, pose in enumerate(poses_to_check): + expected_bottom_row = np.array([0, 0, 0, 1]) + if not np.allclose(pose[3, :], expected_bottom_row, atol=1e-6): + logger.log_warning( + f"Pose {i} bottom row {pose[3, :]} != [0,0,0,1]. May not be a valid transformation matrix." + ) + + def get_type_name(self) -> str: + """Return the type name for this visualizer.""" + return VisualizationType.AXIS.value + + def _save_impl(self, filepath: Path, **kwargs: Any) -> None: + """Save the visualization to file.""" + if self._last_visualization is None: + raise RuntimeError("No visualization to save. Call visualize() first.") + + if self.backend == "open3d": + if filepath.suffix.lower() in [".ply", ".obj", ".stl"]: + # Save combined mesh + combined_mesh = o3d.geometry.TriangleMesh() + for frame in self._last_visualization: + combined_mesh += frame + o3d.io.write_triangle_mesh(str(filepath), combined_mesh) + logger.log_info(f"Saved Open3D coordinate frames to {filepath}") + else: + logger.log_warning( + f"Unsupported file format {filepath.suffix} for Open3D backend" + ) + + elif self.backend == "matplotlib": + self._last_visualization.savefig(filepath, **kwargs) + logger.log_info(f"Saved matplotlib plot to {filepath}") + + elif self.backend in ["sim_manager", "data"]: + # Save as numpy file + if self.backend == "sim_manager": + data_to_save = { + "axis_names": self._last_visualization, + "type": "sim_manager_axes", + } + else: + data_to_save = self._last_visualization + + np.save(filepath.with_suffix(".npy"), data_to_save, allow_pickle=True) + logger.log_info(f"Saved axis data to {filepath.with_suffix('.npy')}") + + def _show_impl(self, **kwargs: Any) -> None: + """Display the visualization.""" + if self._last_visualization is None: + logger.log_warning("No visualization to show. Call visualize() first.") + return + + if self.backend == "open3d": + o3d.visualization.draw_geometries(self._last_visualization) + + elif self.backend == "matplotlib": + plt.show() + + elif self.backend == "sim_manager": + logger.log_info( + f"Axes are displayed in simulation. Marker names: {self._last_visualization}" + ) + + elif self.backend == "data": + logger.log_info( + f"Data backend - no visual display. Use .save() to export data." + ) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/visualizers/base_visualizer.py b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/base_visualizer.py new file mode 100644 index 0000000..d0f8e48 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/base_visualizer.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from abc import ABC, abstractmethod +from typing import Protocol, Union, Dict, Any, Tuple +from pathlib import Path + +try: + import open3d as o3d + + OPEN3D_AVAILABLE = True +except ImportError: + OPEN3D_AVAILABLE = False + +try: + import matplotlib.pyplot as plt + from mpl_toolkits.mplot3d import Axes3D + + MATPLOTLIB_AVAILABLE = True +except ImportError: + MATPLOTLIB_AVAILABLE = False + +from embodichain.utils import logger +from embodichain.lab.sim.utility.workspace_analyzer.configs.visualization_config import ( + VisualizationConfig, +) + + +__all__ = [ + "IVisualizer", + "BaseVisualizer", +] + + +class IVisualizer(Protocol): + """Interface for all visualizers. + + This protocol defines the contract that all visualizers must follow. + """ + + def visualize( + self, + points: Union[torch.Tensor, np.ndarray], + colors: Union[torch.Tensor, np.ndarray] | None = None, + **kwargs: Any, + ) -> Any: + """Visualize the workspace data. + + Args: + points: Array of shape (N, 3) containing point positions. + colors: Optional array of shape (N, 3) or (N, 4) containing colors. + **kwargs: Additional visualization parameters. + + Returns: + Visualization object (e.g., Open3D geometry, matplotlib figure). + """ + ... + + def save(self, filepath: Union[str, Path], **kwargs: Any) -> None: + """Save visualization to file. + + Args: + filepath: Path to save the visualization. + **kwargs: Additional save parameters. + """ + ... + + def get_type_name(self) -> str: + """Get the name of the visualization type. + + Returns: + String identifier for the visualization type. + """ + ... + + +class BaseVisualizer(ABC): + """Abstract base class for all visualizers. + + This class provides common functionality and enforces the implementation + of the visualization method in all derived classes. + + Attributes: + backend: Visualization backend ('open3d', 'matplotlib', etc.). + config: Configuration dictionary for visualization parameters. + """ + + def __init__(self, backend: str = "open3d", config: Dict[str, Any] | None = None): + """Initialize the base visualizer. + + Args: + backend: Visualization backend to use. Defaults to "open3d". + config: Optional configuration dictionary. Defaults to None. + """ + self.backend = backend + self.config = config or {} + self._validate_backend() + + # Store last visualization for reuse + self._last_visualization = None + + @classmethod + def from_config(cls, config: VisualizationConfig, backend: str = "open3d"): + """Create a visualizer instance from a VisualizationConfig. + + Args: + config: VisualizationConfig instance with visualization settings. + backend: Visualization backend to use. + + Returns: + Configured visualizer instance. + """ + config_dict = { + "enabled": config.enabled, + "voxel_size": config.voxel_size, + "nb_neighbors": config.nb_neighbors, + "std_ratio": config.std_ratio, + "is_voxel_down": config.is_voxel_down, + "color_by_distance": config.color_by_distance, + } + return cls(backend=backend, config=config_dict) + + def _validate_backend(self) -> None: + """Validate the visualization backend is available. + + Raises: + ImportError: If the required backend is not available. + """ + if self.backend == "data": + # Data backend doesn't require any external dependencies + return + elif self.backend == "open3d" and not OPEN3D_AVAILABLE: + raise ImportError( + "Open3D is not available. Install it with: pip install open3d" + ) + elif self.backend == "matplotlib" and not MATPLOTLIB_AVAILABLE: + raise ImportError( + "Matplotlib is not available. Install it with: pip install matplotlib" + ) + + @abstractmethod + def visualize( + self, + points: Union[torch.Tensor, np.ndarray], + colors: Union[torch.Tensor, np.ndarray] | None = None, + **kwargs: Any, + ) -> Any: + """Visualize the workspace data. + + This method must be implemented by all derived classes. + + Args: + points: Array of shape (N, 3) containing point positions. + colors: Optional array of shape (N, 3) or (N, 4) containing colors. + **kwargs: Additional visualization parameters. + + Returns: + Visualization object. + + Raises: + NotImplementedError: If the method is not implemented in the derived class. + """ + raise NotImplementedError("Subclasses must implement the visualize method") + + @abstractmethod + def get_type_name(self) -> str: + """Get the name of the visualization type. + + Returns: + String identifier for the visualization type. + """ + raise NotImplementedError("Subclasses must implement get_type_name method") + + def save(self, filepath: Union[str, Path], **kwargs: Any) -> None: + """Save the last visualization to file. + + Args: + filepath: Path to save the visualization. + **kwargs: Additional save parameters. + + Raises: + RuntimeError: If no visualization has been created yet. + """ + if self._last_visualization is None: + raise RuntimeError("No visualization available. Call visualize() first.") + + filepath = Path(filepath) + filepath.parent.mkdir(parents=True, exist_ok=True) + + self._save_impl(filepath, **kwargs) + logger.log_info(f"Visualization saved to {filepath}") + + @abstractmethod + def _save_impl(self, filepath: Path, **kwargs: Any) -> None: + """Implementation of save functionality. + + Args: + filepath: Path to save the visualization. + **kwargs: Additional save parameters. + """ + raise NotImplementedError("Subclasses must implement _save_impl method") + + def _to_numpy(self, data: Union[torch.Tensor, np.ndarray]) -> np.ndarray: + """Convert data to numpy array. + + Args: + data: Input data (torch tensor or numpy array). + + Returns: + NumPy array. + """ + if isinstance(data, torch.Tensor): + return data.detach().cpu().numpy() + return np.asarray(data) + + def _validate_points(self, points: np.ndarray) -> None: + """Validate the points array. + + Args: + points: Array of shape (N, 3) containing point positions. + + Raises: + ValueError: If points array is invalid. + """ + if points.ndim != 2 or points.shape[1] != 3: + raise ValueError(f"Points must have shape (N, 3), got {points.shape}") + + if len(points) == 0: + raise ValueError("Points array is empty") + + def _validate_colors( + self, colors: Union[torch.Tensor, np.ndarray] | None, num_points: int + ) -> Union[np.ndarray, None]: + """Validate and normalize the colors array. + + Args: + colors: Optional array of shape (N, 3) or (N, 4) containing colors. + num_points: Number of points to validate against. + + Returns: + Validated colors array or None. + + Raises: + ValueError: If colors array is invalid. + """ + if colors is None: + return None + + colors = self._to_numpy(colors) + + if colors.ndim != 2 or colors.shape[0] != num_points: + raise ValueError( + f"Colors must have shape ({num_points}, 3) or ({num_points}, 4), " + f"got {colors.shape}" + ) + + if colors.shape[1] not in [3, 4]: + raise ValueError( + f"Colors must have 3 (RGB) or 4 (RGBA) channels, " + f"got {colors.shape[1]}" + ) + + # Normalize colors to [0, 1] if needed + if colors.max() > 1.0: + colors = colors / 255.0 + + return colors + + def _get_default_colors( + self, num_points: int, color: Tuple[float, float, float] = (0.0, 1.0, 0.0) + ) -> np.ndarray: + """Generate default colors for points. + + Args: + num_points: Number of points. + color: RGB color tuple. Defaults to green (0.0, 1.0, 0.0). + + Returns: + Array of shape (num_points, 3) with default colors. + """ + return np.tile(np.array(color), (num_points, 1)) + + def show(self, **kwargs: Any) -> None: + """Display the visualization interactively. + + Args: + **kwargs: Backend-specific display parameters. + + Raises: + RuntimeError: If no visualization has been created yet. + """ + if self._last_visualization is None: + raise RuntimeError("No visualization available. Call visualize() first.") + + self._show_impl(**kwargs) + + @abstractmethod + def _show_impl(self, **kwargs: Any) -> None: + """Implementation of show functionality. + + Args: + **kwargs: Backend-specific display parameters. + """ + raise NotImplementedError("Subclasses must implement _show_impl method") + + def clear(self) -> None: + """Clear the current visualization.""" + self._last_visualization = None + + def __repr__(self) -> str: + """String representation of the visualizer.""" + return ( + f"{self.__class__.__name__}(" + f"type={self.get_type_name()}, " + f"backend={self.backend})" + ) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/visualizers/point_cloud_visualizer.py b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/point_cloud_visualizer.py new file mode 100644 index 0000000..6494ac7 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/point_cloud_visualizer.py @@ -0,0 +1,265 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Union, Any, Dict +from pathlib import Path + +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + VisualizationType, +) +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.base_visualizer import ( + BaseVisualizer, + OPEN3D_AVAILABLE, +) + +if OPEN3D_AVAILABLE: + import open3d as o3d + +from embodichain.utils import logger + +__all__ = ["PointCloudVisualizer"] + + +class PointCloudVisualizer(BaseVisualizer): + """Point cloud visualizer using Open3D or matplotlib. + + Attributes: + point_size: Size of points in visualization. + """ + + def __init__( + self, + backend: str = "sim_manager", + point_size: float = 2.0, + config: Dict[str, Any] | None = None, + sim_manager: Any | None = None, + control_part_name: str | None = None, + ): + """Initialize the point cloud visualizer. + + Args: + backend: Visualization backend ('sim_manager', 'open3d', 'matplotlib', or 'data'). + Defaults to 'sim_manager'. 'data' backend returns raw data without visualization. + 'sim_manager' backend uses simulation environment for visualization. + point_size: Size of points in visualization. Defaults to 2.0. + config: Optional configuration dictionary. Defaults to None. + sim_manager: SimulationManager instance for 'sim_manager' backend. Defaults to None. + control_part_name: Control part name for naming the point cloud. Defaults to None. + """ + super().__init__(backend, config) + self.point_size = point_size + self.sim_manager = sim_manager + self.control_part_name = control_part_name + + def visualize( + self, + points: Union[torch.Tensor, np.ndarray], + colors: Union[torch.Tensor, np.ndarray] | None = None, + **kwargs: Any, + ) -> Any: + """Visualize points as a point cloud. + + Args: + points: Array of shape (N, 3) containing point positions. + colors: Optional array of shape (N, 3) or (N, 4) containing colors. + **kwargs: Additional visualization parameters: + - point_size: Override default point size + + Returns: + Open3D PointCloud geometry or matplotlib figure. + + Examples: + >>> visualizer = PointCloudVisualizer() + >>> points = np.random.rand(1000, 3) + >>> colors = np.random.rand(1000, 3) + >>> pcd = visualizer.visualize(points, colors) + >>> visualizer.show() + """ + # Convert to numpy + points = self._to_numpy(points) + self._validate_points(points) + + # Get visualization parameters + point_size = kwargs.get("point_size", self.point_size) + + # Validate and prepare colors + colors = self._validate_colors(colors, len(points)) + if colors is None: + colors = self._get_default_colors(len(points)) + + # Convert to RGB if RGBA + if colors.shape[1] == 4: + colors = colors[:, :3] + + if self.backend == "data": + # Return raw data for user to handle + data = { + "points": points, + "colors": colors, + "point_size": point_size, + "type": "point_cloud", + } + self._last_visualization = {"data": data} + return data + elif self.backend == "sim_manager": + pcd_handle = self._create_sim_manager_point_cloud( + points, colors, point_size + ) + self._last_visualization = { + "point_cloud_handle": pcd_handle, + "point_size": point_size, + } + return pcd_handle + elif self.backend == "open3d": + pcd = self._create_open3d_point_cloud(points, colors) + self._last_visualization = { + "point_cloud": pcd, + "point_size": point_size, + } + return pcd + elif self.backend == "matplotlib": + fig = self._create_matplotlib_point_cloud(points, colors, point_size) + self._last_visualization = {"figure": fig} + return fig + else: + raise ValueError(f"Unsupported backend: {self.backend}") + + def _create_sim_manager_point_cloud( + self, points: np.ndarray, colors: np.ndarray, point_size: float + ) -> Any: + if self.sim_manager is None: + raise ValueError("sim_manager is required for 'sim_manager' backend") + + # Get simulation environment + env = self.sim_manager.get_env() + if env is None: + raise RuntimeError("Simulation manager has no active simulation") + + # Create point cloud name + pcd_name = f"workspace_pcd_{self.control_part_name or 'default'}" + + # Create point cloud in simulation + pcd_handle = env.create_point_cloud(name=pcd_name) + pcd_handle.add_points(points) + pcd_handle.set_colors(colors) + pcd_handle.set_point_size(point_size) + + logger.log_info( + f"Created point cloud '{pcd_name}' with {len(points)} points " + f"(point_size={point_size})" + ) + + return pcd_handle + + def _create_open3d_point_cloud( + self, points: np.ndarray, colors: np.ndarray + ) -> "o3d.geometry.PointCloud": + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points) + pcd.colors = o3d.utility.Vector3dVector(colors) + + logger.log_info(f"Created point cloud with {len(points)} points") + + return pcd + + def _create_matplotlib_point_cloud( + self, points: np.ndarray, colors: np.ndarray, point_size: float + ): + import matplotlib.pyplot as plt + from mpl_toolkits.mplot3d import Axes3D + + fig = plt.figure(figsize=(10, 8)) + ax = fig.add_subplot(111, projection="3d") + + ax.scatter( + points[:, 0], points[:, 1], points[:, 2], c=colors, s=point_size, alpha=0.6 + ) + + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.set_title("Workspace Point Cloud") + + return fig + + def _save_impl(self, filepath: Path, **kwargs: Any) -> None: + if self.backend == "data": + # Save data as numpy file + data = self._last_visualization["data"] + np.savez(filepath, **data) + elif self.backend == "open3d": + pcd = self._last_visualization["point_cloud"] + + # Determine file format from extension + suffix = filepath.suffix.lower() + if suffix in [".pcd", ".ply", ".xyz", ".xyzrgb", ".pts"]: + o3d.io.write_point_cloud(str(filepath), pcd) + elif suffix in [".png", ".jpg", ".jpeg"]: + # Render to image + vis = o3d.visualization.Visualizer() + vis.create_window(visible=False) + vis.add_geometry(pcd) + + # Coordinate frame removed - implement separately if needed + + vis.update_geometry(pcd) + vis.poll_events() + vis.update_renderer() + vis.capture_screen_image(str(filepath)) + vis.destroy_window() + else: + raise ValueError( + f"Unsupported file format: {suffix}. " + f"Use .pcd, .ply, .xyz, .xyzrgb, .pts, .png, .jpg" + ) + + elif self.backend == "matplotlib": + fig = self._last_visualization["figure"] + fig.savefig(filepath, dpi=300, bbox_inches="tight") + + def _show_impl(self, **kwargs: Any) -> None: + if self.backend == "data": + logger.log_warning( + "Cannot display visualization with 'data' backend. " + "Use 'open3d' or 'matplotlib' backend for interactive display." + ) + return + elif self.backend == "open3d": + geometries = [self._last_visualization["point_cloud"]] + + # Coordinate frame removed - implement separately if needed + + # Set point size in visualization + vis = o3d.visualization.Visualizer() + vis.create_window() + for geom in geometries: + vis.add_geometry(geom) + + render_option = vis.get_render_option() + render_option.point_size = self._last_visualization.get("point_size", 2.0) + + vis.run() + vis.destroy_window() + + elif self.backend == "matplotlib": + import matplotlib.pyplot as plt + + plt.show() + + def get_type_name(self) -> str: + return VisualizationType.POINT_CLOUD.value diff --git a/embodichain/lab/sim/utility/workspace_analyzer/visualizers/sphere_visualizer.py b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/sphere_visualizer.py new file mode 100644 index 0000000..0338533 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/sphere_visualizer.py @@ -0,0 +1,310 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Union, Any, Dict +from pathlib import Path + +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + VisualizationType, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.base_visualizer import ( + BaseVisualizer, + OPEN3D_AVAILABLE, +) + +if OPEN3D_AVAILABLE: + import open3d as o3d + +from embodichain.utils import logger + + +__all__ = ["SphereVisualizer"] + + +class SphereVisualizer(BaseVisualizer): + """Sphere-based visualizer using Open3D or matplotlib. + + Attributes: + sphere_radius: Radius of each sphere. + sphere_resolution: Resolution of sphere mesh (higher = smoother). + """ + + def __init__( + self, + backend: str = "sim_manager", + sphere_radius: float = 0.005, + sphere_resolution: int = 10, + config: Dict[str, Any] | None = None, + sim_manager: Any | None = None, + control_part_name: str | None = None, + ): + """Initialize the sphere visualizer. + + Args: + backend: Visualization backend ('sim_manager', 'open3d', 'matplotlib', or 'data'). + Defaults to 'sim_manager'. 'data' backend returns sphere data without visualization. + sphere_radius: Radius of each sphere. Defaults to 0.005. + sphere_resolution: Sphere mesh resolution. Defaults to 10. + config: Optional configuration dictionary. Defaults to None. + sim_manager: SimulationManager instance for 'sim_manager' backend. Defaults to None. + control_part_name: Control part name for naming. Defaults to None. + """ + super().__init__(backend, config) + self.sphere_radius = sphere_radius + self.sphere_resolution = sphere_resolution + self.sim_manager = sim_manager + self.control_part_name = control_part_name + + def visualize( + self, + points: Union[torch.Tensor, np.ndarray], + colors: Union[torch.Tensor, np.ndarray] | None = None, + **kwargs: Any, + ) -> Any: + """Visualize points as spheres. + + Args: + points: Array of shape (N, 3) containing point positions. + colors: Optional array of shape (N, 3) or (N, 4) containing colors. + **kwargs: Additional visualization parameters: + - sphere_radius: Override default sphere radius + - sphere_resolution: Override sphere resolution + - max_spheres: Maximum number of spheres to render (for performance) + + Returns: + Open3D TriangleMesh or matplotlib figure. + + Examples: + >>> visualizer = SphereVisualizer(sphere_radius=0.01) + >>> points = np.random.rand(100, 3) + >>> colors = np.random.rand(100, 3) + >>> mesh = visualizer.visualize(points, colors) + >>> visualizer.show() + """ + # Convert to numpy + points = self._to_numpy(points) + self._validate_points(points) + + # Get visualization parameters + sphere_radius = kwargs.get("sphere_radius", self.sphere_radius) + sphere_resolution = kwargs.get("sphere_resolution", self.sphere_resolution) + max_spheres = kwargs.get("max_spheres", None) + + # Limit number of spheres for performance + if max_spheres is not None and len(points) > max_spheres: + logger.log_warning( + f"Limiting visualization to {max_spheres} spheres " + f"(total points: {len(points)})" + ) + indices = np.random.choice(len(points), max_spheres, replace=False) + points = points[indices] + if colors is not None: + colors = ( + colors[indices] + if isinstance(colors, np.ndarray) + else self._to_numpy(colors)[indices] + ) + + # Validate and prepare colors + colors = self._validate_colors(colors, len(points)) + if colors is None: + colors = self._get_default_colors(len(points)) + + # Convert to RGB if RGBA + if colors.shape[1] == 4: + colors = colors[:, :3] + + if self.backend == "data": + # Return sphere data + data = { + "centers": points, + "colors": colors, + "radius": sphere_radius, + "resolution": sphere_resolution, + "num_spheres": len(points), + "type": "spheres", + } + self._last_visualization = {"data": data} + logger.log_info( + f"Created sphere data with {len(points)} spheres (radius={sphere_radius})" + ) + return data + elif self.backend == "sim_manager": + spheres_handle = self._create_sim_manager_spheres( + points, colors, sphere_radius + ) + self._last_visualization = { + "spheres_handle": spheres_handle, + "radius": sphere_radius, + } + return spheres_handle + elif self.backend == "open3d": + mesh = self._create_open3d_spheres( + points, colors, sphere_radius, sphere_resolution + ) + self._last_visualization = {"mesh": mesh} + return mesh + elif self.backend == "matplotlib": + fig = self._create_matplotlib_spheres(points, colors, sphere_radius) + self._last_visualization = {"figure": fig} + return fig + else: + raise ValueError(f"Unsupported backend: {self.backend}") + + def _create_sim_manager_spheres( + self, points: np.ndarray, colors: np.ndarray, sphere_radius: float + ) -> Any: + if self.sim_manager is None: + raise ValueError("sim_manager is required for 'sim_manager' backend") + + # Get simulation env + env = self.sim_manager.get_env() + if env is None: + raise RuntimeError("Simulation manager has no active env") + + sphere_handles = [] + for i, point in enumerate(points): + sphere_handle = env.create_sphere(radius=sphere_radius, resolution=10) + # Unpack array to individual x, y, z coordinates + sphere_handle.set_location( + float(point[0]), float(point[1]), float(point[2]) + ) + # TODO: Unsupported in current sim_manager API + # sphere_handle.set_color(colors[i].tolist()) + sphere_handle.set_name(f"workspace_sphere_{i}") + sphere_handles.append(sphere_handle) + + logger.log_info(f"Created {len(points)} spheres with radius={sphere_radius}") + + return sphere_handles + + def _create_open3d_spheres( + self, + points: np.ndarray, + colors: np.ndarray, + sphere_radius: float, + sphere_resolution: int, + ) -> "o3d.geometry.TriangleMesh": + # Create a template sphere + sphere_template = o3d.geometry.TriangleMesh.create_sphere( + radius=sphere_radius, resolution=sphere_resolution + ) + + # Combine all spheres into one mesh + combined_mesh = o3d.geometry.TriangleMesh() + + for point, color in zip(points, colors): + # Copy and translate sphere + sphere = o3d.geometry.TriangleMesh(sphere_template) + sphere.translate(point) + sphere.paint_uniform_color(color) + + # Merge into combined mesh + combined_mesh += sphere + + # Compute normals for proper lighting + combined_mesh.compute_vertex_normals() + + logger.log_info(f"Created {len(points)} spheres with radius={sphere_radius}") + + return combined_mesh + + def _create_matplotlib_spheres( + self, points: np.ndarray, colors: np.ndarray, sphere_radius: float + ): + import matplotlib.pyplot as plt + from mpl_toolkits.mplot3d import Axes3D + + fig = plt.figure(figsize=(10, 8)) + ax = fig.add_subplot(111, projection="3d") + + # Scale marker size based on sphere radius + marker_size = (sphere_radius * 1000) ** 2 + + ax.scatter( + points[:, 0], + points[:, 1], + points[:, 2], + c=colors, + s=marker_size, + alpha=0.8, + marker="o", + ) + + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.set_title(f"Workspace Spheres (radius={sphere_radius:.4f})") + + return fig + + def _save_impl(self, filepath: Path, **kwargs: Any) -> None: + if self.backend == "data": + # Save sphere data + data = self._last_visualization["data"] + np.savez(filepath, **data) + elif self.backend == "open3d": + mesh = self._last_visualization["mesh"] + + # Determine file format from extension + suffix = filepath.suffix.lower() + if suffix in [".ply", ".obj", ".stl", ".gltf", ".glb"]: + o3d.io.write_triangle_mesh(str(filepath), mesh) + elif suffix in [".png", ".jpg", ".jpeg"]: + # Render to image + vis = o3d.visualization.Visualizer() + vis.create_window(visible=False) + vis.add_geometry(mesh) + + vis.update_geometry(mesh) + vis.poll_events() + vis.update_renderer() + vis.capture_screen_image(str(filepath)) + vis.destroy_window() + else: + raise ValueError( + f"Unsupported file format: {suffix}. " + f"Use .ply, .obj, .stl, .gltf, .glb, .png, .jpg" + ) + + elif self.backend == "matplotlib": + fig = self._last_visualization["figure"] + fig.savefig(filepath, dpi=300, bbox_inches="tight") + + def _show_impl(self, **kwargs: Any) -> None: + if self.backend == "data": + logger.log_warning( + "Cannot display visualization with 'data' backend. " + "Use 'open3d' or 'matplotlib' backend for interactive display." + ) + return + elif self.backend == "open3d": + geometries = [self._last_visualization["mesh"]] + + # Coordinate frame removed - implement separately if needed + + o3d.visualization.draw_geometries(geometries) + + elif self.backend == "matplotlib": + import matplotlib.pyplot as plt + + plt.show() + + def get_type_name(self) -> str: + return VisualizationType.SPHERE.value diff --git a/embodichain/lab/sim/utility/workspace_analyzer/visualizers/visualizer_factory.py b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/visualizer_factory.py new file mode 100644 index 0000000..f8638a9 --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/visualizer_factory.py @@ -0,0 +1,274 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations +from typing import Dict, Type, Any +from threading import Lock + +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + VisualizationType, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.base_visualizer import ( + BaseVisualizer, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.point_cloud_visualizer import ( + PointCloudVisualizer, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.voxel_visualizer import ( + VoxelVisualizer, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.sphere_visualizer import ( + SphereVisualizer, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.axis_visualizer import ( + AxisVisualizer, +) + +from embodichain.utils import logger + +__all__ = [ + "VisualizerFactory", + "create_visualizer", +] + + +class VisualizerFactory: + """Factory class for creating visualizers (Singleton pattern). + + This factory allows registration and creation of visualizers based on + the visualization type. It uses the singleton pattern to ensure only + one instance exists throughout the application. + + The factory comes pre-registered with built-in visualizers: + - POINT_CLOUD: PointCloudVisualizer + - VOXEL: VoxelVisualizer + - SPHERE: SphereVisualizer + + Additional visualizers can be registered using register_visualizer(). + + Examples: + >>> factory = VisualizerFactory() + >>> visualizer = factory.create_visualizer( + ... VisualizationType.POINT_CLOUD, + ... backend='open3d' + ... ) + >>> isinstance(visualizer, PointCloudVisualizer) + True + + >>> # Register custom visualizer + >>> factory.register_visualizer("custom", CustomVisualizer) + >>> custom_viz = factory.create_visualizer("custom") + """ + + _instance: VisualizerFactory | None = None + _lock: Lock = Lock() + + def __new__(cls): + """Create or return the singleton instance. + + Returns: + The singleton VisualizerFactory instance. + """ + if cls._instance is None: + with cls._lock: + # Double-checked locking + if cls._instance is None: + cls._instance = super(VisualizerFactory, cls).__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self): + """Initialize the factory with built-in visualizers. + + This method only runs once due to the singleton pattern. + """ + # Prevent re-initialization + if self._initialized: + return + + self._visualizers: Dict[str, Type[BaseVisualizer]] = {} + self._register_builtin_visualizers() + self._initialized = True + + def _register_builtin_visualizers(self) -> None: + """Register the built-in visualizers.""" + self._visualizers[VisualizationType.POINT_CLOUD.value] = PointCloudVisualizer + self._visualizers[VisualizationType.VOXEL.value] = VoxelVisualizer + self._visualizers[VisualizationType.SPHERE.value] = SphereVisualizer + self._visualizers[VisualizationType.AXIS.value] = AxisVisualizer + + logger.log_debug( + f"Registered built-in visualizers: {list(self._visualizers.keys())}" + ) + + def register_visualizer( + self, name: str, visualizer_class: Type[BaseVisualizer] + ) -> None: + """Register a new visualizer class. + + Args: + name: String identifier for the visualizer type. + visualizer_class: The visualizer class to register. + Must inherit from BaseVisualizer. + + Raises: + TypeError: If visualizer_class is not a subclass of BaseVisualizer. + + Examples: + >>> factory = VisualizerFactory() + >>> factory.register_visualizer("my_viz", MyVisualizerClass) + """ + if not issubclass(visualizer_class, BaseVisualizer): + raise TypeError( + f"visualizer_class must be a subclass of BaseVisualizer, " + f"got {visualizer_class}" + ) + + if name in self._visualizers: + logger.log_warning( + f"Visualizer '{name}' already registered. " + f"Overwriting with {visualizer_class.__name__}." + ) + + self._visualizers[name] = visualizer_class + logger.log_info(f"Registered visualizer '{name}': {visualizer_class.__name__}") + + def create_visualizer( + self, viz_type: VisualizationType | str | None = None, **kwargs: Any + ) -> BaseVisualizer: + """Create a visualizer instance based on the type. + + Args: + viz_type: The visualization type to use. Can be a VisualizationType enum + or a string identifier. If None, defaults to POINT_CLOUD. + **kwargs: Additional keyword arguments to pass to the visualizer constructor. + Common arguments include: + - backend: Visualization backend ('open3d', 'matplotlib', 'data') + 'data' backend returns processed data without visualization + - voxel_size: For VoxelVisualizer + - sphere_radius: For SphereVisualizer + - point_size: For PointCloudVisualizer + + Returns: + An instance of the requested visualizer. + + Raises: + ValueError: If the visualization type is not registered. + + Examples: + >>> factory = VisualizerFactory() + >>> viz = factory.create_visualizer( + ... VisualizationType.POINT_CLOUD, + ... backend='open3d' + ... ) + >>> viz = factory.create_visualizer("voxel", voxel_size=0.02) + >>> viz = factory.create_visualizer() # Uses default (POINT_CLOUD) + """ + # Default to POINT_CLOUD if no type specified + if viz_type is None: + viz_type = VisualizationType.POINT_CLOUD + + # Convert enum to string if necessary + if isinstance(viz_type, VisualizationType): + type_name = viz_type.value + else: + type_name = viz_type + + # Check if type is registered + if type_name not in self._visualizers: + available = list(self._visualizers.keys()) + raise ValueError( + f"Unknown visualization type: '{type_name}'. " + f"Available types: {available}. " + f"You can register a custom visualizer using register_visualizer()." + ) + + # Create and return visualizer instance + visualizer_class = self._visualizers[type_name] + visualizer = visualizer_class(**kwargs) + + logger.log_info( + f"Created visualizer: {visualizer_class.__name__} with kwargs: {kwargs}" + ) + + return visualizer + + def list_available_types(self) -> list[str]: + """List all registered visualization types. + + Returns: + List of registered type names. + """ + return list(self._visualizers.keys()) + + def is_registered(self, viz_type: VisualizationType | str) -> bool: + """Check if a visualization type is registered. + + Args: + viz_type: The visualization type to check. + + Returns: + True if the type is registered, False otherwise. + """ + if isinstance(viz_type, VisualizationType): + type_name = viz_type.value + else: + type_name = viz_type + + return type_name in self._visualizers + + @classmethod + def reset_instance(cls) -> None: + """Reset the singleton instance (mainly for testing). + + Warning: + This should only be used in testing scenarios. + """ + with cls._lock: + cls._instance = None + + def __repr__(self) -> str: + """String representation of the factory.""" + types = self.list_available_types() + return f"VisualizerFactory(types={types})" + + +# Convenience function for creating visualizers +def create_visualizer( + viz_type: VisualizationType | str | None = None, **kwargs: Any +) -> BaseVisualizer: + """Convenience function to create a visualizer. + + This is a shorthand for VisualizerFactory().create_visualizer(). + + Args: + viz_type: The visualization type to use. + **kwargs: Additional keyword arguments to pass to the visualizer constructor. + + Returns: + An instance of the requested visualizer. + + Examples: + >>> viz = create_visualizer(VisualizationType.POINT_CLOUD, backend='open3d') + >>> viz = create_visualizer("voxel", voxel_size=0.02) + """ + factory = VisualizerFactory() + return factory.create_visualizer(viz_type, **kwargs) diff --git a/embodichain/lab/sim/utility/workspace_analyzer/visualizers/voxel_visualizer.py b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/voxel_visualizer.py new file mode 100644 index 0000000..0324a2f --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/visualizers/voxel_visualizer.py @@ -0,0 +1,342 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import numpy as np +import torch +from typing import Union, Any, Dict +from pathlib import Path + +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + VisualizationType, +) + +from embodichain.lab.sim.utility.workspace_analyzer.visualizers.base_visualizer import ( + BaseVisualizer, + OPEN3D_AVAILABLE, +) + +if OPEN3D_AVAILABLE: + import open3d as o3d + +from embodichain.utils import logger + + +__all__ = ["VoxelVisualizer"] + + +class VoxelVisualizer(BaseVisualizer): + """Voxel grid visualizer using Open3D or matplotlib. + + Attributes: + voxel_size: Size of each voxel cube. + """ + + def __init__( + self, + backend: str = "sim_manager", + voxel_size: float = 0.01, + config: Dict[str, Any] | None = None, + sim_manager: Any | None = None, + control_part_name: str | None = None, + ): + """Initialize the voxel visualizer. + + Args: + backend: Visualization backend ('sim_manager', 'open3d', 'matplotlib', or 'data'). + Defaults to 'sim_manager'. 'data' backend returns voxelized data without visualization. + voxel_size: Size of each voxel. Defaults to 0.01. + config: Optional configuration dictionary. Defaults to None. + sim_manager: SimulationManager instance for 'sim_manager' backend. Defaults to None. + control_part_name: Control part name for naming. Defaults to None. + """ + super().__init__(backend, config) + self.voxel_size = voxel_size + self.sim_manager = sim_manager + self.control_part_name = control_part_name + + def visualize( + self, + points: Union[torch.Tensor, np.ndarray], + colors: Union[torch.Tensor, np.ndarray] | None = None, + **kwargs: Any, + ) -> Any: + """Visualize points as a voxel grid. + + Args: + points: Array of shape (N, 3) containing point positions. + colors: Optional array of shape (N, 3) or (N, 4) containing colors. + **kwargs: Additional visualization parameters: + - voxel_size: Override default voxel size + Returns: + Open3D VoxelGrid geometry or matplotlib figure. + + Examples: + >>> visualizer = VoxelVisualizer(voxel_size=0.02) + >>> points = np.random.rand(1000, 3) + >>> voxel_grid = visualizer.visualize(points) + >>> visualizer.show() + """ + # Convert to numpy + points = self._to_numpy(points) + self._validate_points(points) + + # Get visualization parameters + voxel_size = kwargs.get("voxel_size", self.voxel_size) + + # Validate and prepare colors + colors = self._validate_colors(colors, len(points)) + if colors is None: + colors = self._get_default_colors(len(points)) + + # Convert to RGB if RGBA + if colors.shape[1] == 4: + colors = colors[:, :3] + + if self.backend == "data": + # Return voxelized data + voxel_data = self._create_voxel_data(points, colors, voxel_size) + self._last_visualization = {"data": voxel_data} + return voxel_data + elif self.backend == "sim_manager": + voxels_handle = self._create_sim_manager_voxels(points, colors, voxel_size) + self._last_visualization = { + "voxels_handle": voxels_handle, + "voxel_size": voxel_size, + } + return voxels_handle + elif self.backend == "open3d": + voxel_grid = self._create_open3d_voxel_grid(points, colors, voxel_size) + self._last_visualization = { + "voxel_grid": voxel_grid, + "voxel_size": voxel_size, + } + return voxel_grid + elif self.backend == "matplotlib": + fig = self._create_matplotlib_voxel_grid(points, colors, voxel_size) + self._last_visualization = {"figure": fig} + return fig + else: + raise ValueError(f"Unsupported backend: {self.backend}") + + def _create_sim_manager_voxels( + self, points: np.ndarray, colors: np.ndarray, voxel_size: float + ) -> Any: + if self.sim_manager is None: + raise ValueError("sim_manager is required for 'sim_manager' backend") + + # Get simulation env + env = self.sim_manager.get_env() + if env is None: + raise RuntimeError("Simulation manager has no active env") + + cube_handles = [] + for i, point in enumerate(points): + cube_handle = env.create_cube(l=voxel_size, w=voxel_size, h=voxel_size) + cube_handle.set_location(float(point[0]), float(point[1]), float(point[2])) + # TODO: Unsupported in current sim_manager API + # cube_handle.set_color(colors[i].tolist()) + cube_handle.set_name(f"workspace_cube_{i}") + cube_handles.append(cube_handle) + + logger.log_info(f"Created {len(points)} cubes with size={voxel_size}") + + return cube_handles + + def _create_open3d_voxel_grid( + self, points: np.ndarray, colors: np.ndarray, voxel_size: float + ) -> "o3d.geometry.VoxelGrid": + # First create a point cloud + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points) + pcd.colors = o3d.utility.Vector3dVector(colors) + + # Convert to voxel grid + voxel_grid = o3d.geometry.VoxelGrid.create_from_point_cloud( + pcd, voxel_size=voxel_size + ) + + num_voxels = len(voxel_grid.get_voxels()) + logger.log_info( + f"Created voxel grid with {num_voxels} voxels " f"(voxel_size={voxel_size})" + ) + + return voxel_grid + + def _create_voxel_data( + self, points: np.ndarray, colors: np.ndarray, voxel_size: float + ) -> Dict[str, Any]: + # Discretize points to voxel grid + min_bounds = points.min(axis=0) + max_bounds = points.max(axis=0) + + # Create voxel indices + voxel_indices = np.floor((points - min_bounds) / voxel_size).astype(int) + + # Create unique voxels with averaged colors + unique_voxels = {} + for idx, color in zip(voxel_indices, colors): + key = tuple(idx) + if key not in unique_voxels: + unique_voxels[key] = [] + unique_voxels[key].append(color) + + # Average colors for each voxel + voxel_positions = [] + voxel_colors = [] + voxel_indices_list = [] + for idx, color_list in unique_voxels.items(): + voxel_positions.append(np.array(idx) * voxel_size + min_bounds) + voxel_colors.append(np.mean(color_list, axis=0)) + voxel_indices_list.append(idx) + + data = { + "voxel_positions": np.array(voxel_positions), + "voxel_colors": np.array(voxel_colors), + "voxel_indices": np.array(voxel_indices_list), + "voxel_size": voxel_size, + "min_bounds": min_bounds, + "max_bounds": max_bounds, + "num_voxels": len(unique_voxels), + "type": "voxel_grid", + } + + logger.log_info( + f"Created voxel data with {data['num_voxels']} voxels " + f"(voxel_size={voxel_size})" + ) + + return data + + def _create_matplotlib_voxel_grid( + self, points: np.ndarray, colors: np.ndarray, voxel_size: float + ): + import matplotlib.pyplot as plt + from mpl_toolkits.mplot3d import Axes3D + + # Discretize points to voxel grid + min_bounds = points.min(axis=0) + + # Create voxel indices + voxel_indices = np.floor((points - min_bounds) / voxel_size).astype(int) + + # Create unique voxels with averaged colors + unique_voxels = {} + for idx, color in zip(voxel_indices, colors): + key = tuple(idx) + if key not in unique_voxels: + unique_voxels[key] = [] + unique_voxels[key].append(color) + + # Average colors for each voxel + voxel_positions = [] + voxel_colors = [] + for idx, color_list in unique_voxels.items(): + voxel_positions.append(np.array(idx) * voxel_size + min_bounds) + voxel_colors.append(np.mean(color_list, axis=0)) + + voxel_positions = np.array(voxel_positions) + voxel_colors = np.array(voxel_colors) + + # Create figure + fig = plt.figure(figsize=(10, 8)) + ax = fig.add_subplot(111, projection="3d") + + # Plot voxels as cubes + ax.scatter( + voxel_positions[:, 0], + voxel_positions[:, 1], + voxel_positions[:, 2], + c=voxel_colors, + s=100, + marker="s", + alpha=0.8, + ) + + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.set_title(f"Workspace Voxel Grid (size={voxel_size:.4f})") + + return fig + + def _save_impl(self, filepath: Path, **kwargs: Any) -> None: + if self.backend == "data": + # Save voxel data + data = self._last_visualization["data"] + np.savez(filepath, **data) + elif self.backend == "open3d": + voxel_grid = self._last_visualization["voxel_grid"] + + # Determine file format from extension + suffix = filepath.suffix.lower() + if suffix in [".ply"]: + # Extract voxel centers and save as point cloud + voxels = voxel_grid.get_voxels() + points = np.array( + [ + voxel_grid.origin + voxel.grid_index * voxel_grid.voxel_size + for voxel in voxels + ] + ) + colors = np.array([voxel.color for voxel in voxels]) + + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points) + pcd.colors = o3d.utility.Vector3dVector(colors) + o3d.io.write_point_cloud(str(filepath), pcd) + elif suffix in [".png", ".jpg", ".jpeg"]: + # Render to image + vis = o3d.visualization.Visualizer() + vis.create_window(visible=False) + vis.add_geometry(voxel_grid) + + # Coordinate frame removed - implement separately if needed + + vis.update_geometry(voxel_grid) + vis.poll_events() + vis.update_renderer() + vis.capture_screen_image(str(filepath)) + vis.destroy_window() + else: + raise ValueError( + f"Unsupported file format: {suffix}. " f"Use .ply, .png, .jpg" + ) + + elif self.backend == "matplotlib": + fig = self._last_visualization["figure"] + fig.savefig(filepath, dpi=300, bbox_inches="tight") + + def _show_impl(self, **kwargs: Any) -> None: + if self.backend == "data": + logger.log_warning( + "Cannot display visualization with 'data' backend. " + "Use 'open3d' or 'matplotlib' backend for interactive display." + ) + return + elif self.backend == "open3d": + geometries = [self._last_visualization["voxel_grid"]] + + # Coordinate frame removed - implement separately if needed + + o3d.visualization.draw_geometries(geometries) + + elif self.backend == "matplotlib": + import matplotlib.pyplot as plt + + plt.show() + + def get_type_name(self) -> str: + return VisualizationType.VOXEL.value diff --git a/embodichain/lab/sim/utility/workspace_analyzer/workspace_analyzer.py b/embodichain/lab/sim/utility/workspace_analyzer/workspace_analyzer.py new file mode 100644 index 0000000..41fdc1f --- /dev/null +++ b/embodichain/lab/sim/utility/workspace_analyzer/workspace_analyzer.py @@ -0,0 +1,2151 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import time +import torch +import numpy as np + +from tqdm import tqdm +from enum import Enum +from pathlib import Path +from dataclasses import dataclass +from contextlib import contextmanager +from typing import List, Tuple, Dict, Any +import os +import sys + +try: + import psutil +except ImportError: + psutil = None + +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.objects.robot import Robot + +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + CacheConfig, + DimensionConstraint, + SamplingConfig, + VisualizationType, + VisualizationConfig, + MetricConfig, +) +from embodichain.lab.sim.utility.workspace_analyzer.samplers import ( + SamplerFactory, + BaseSampler, +) +from embodichain.lab.sim.utility.workspace_analyzer.caches import CacheManager +from embodichain.lab.sim.utility.workspace_analyzer.constraints import ( + WorkspaceConstraintChecker, +) + +from embodichain.utils import logger + +__all__ = [ + "AnalysisMode", + "WorkspaceAnalyzerConfig", + "WorkspaceAnalyzer", +] + + +class AnalysisMode(Enum): + """Workspace analysis mode.""" + + JOINT_SPACE = "joint_space" + """Sample in joint space, compute FK to get workspace points.""" + + CARTESIAN_SPACE = "cartesian_space" + """Sample in Cartesian space, compute IK to verify reachability.""" + + PLANE_SAMPLING = "plane_sampling" + """Sample on a specific plane within Cartesian space.""" + + +@dataclass +class WorkspaceAnalyzerConfig: + """Complete configuration for workspace analyzer.""" + + mode: AnalysisMode = AnalysisMode.JOINT_SPACE + """Analysis mode: joint space or Cartesian space sampling.""" + + sampling: SamplingConfig = None + """Sampling configuration.""" + cache: CacheConfig = None + """Cache configuration.""" + constraint: DimensionConstraint = None + """Dimension constraint configuration.""" + visualization: VisualizationConfig = None + """Visualization configuration.""" + metric: MetricConfig = None + """Metric configuration.""" + + ik_samples_per_point: int = 1 + """For Cartesian mode: number of random joint seeds to try for each Cartesian point.""" + reference_pose: Any | None = None + """Optional reference pose (4x4 matrix) for IK target orientation. If None, uses current robot pose.""" + control_part_name: str | None = None + """Name of the control part (e.g., 'left_arm', 'right_arm'). If None, uses the default solver or first available control part.""" + + # Plane sampling parameters + enable_plane_sampling: bool = False + """Whether to enable plane sampling functionality (uses existing samplers directly)""" + + plane_normal: torch.Tensor | None = None + """Normal vector of the plane for plane sampling [nx, ny, nz]""" + + plane_point: torch.Tensor | None = None + """A point on the plane for plane sampling [x, y, z]""" + + plane_bounds: torch.Tensor | None = None + """Bounds for 2D plane coordinates [[u_min, u_max], [v_min, v_max]]""" + + # Geometric constraint parameters for sampling + constraint_type: str | None = None + """Type of geometric constraint: 'box', 'sphere', None. If None, no constraint applied.""" + + constraint_bounds: torch.Tensor | None = None + """Bounds for constraint: For box: [[x_min, x_max], [y_min, y_max], ...]. For sphere: used to auto-calculate radius if sphere_radius is None.""" + + sphere_center: torch.Tensor | None = None + """Center point for sphere constraint [x, y, z, ...]. If None and constraint_type='sphere', calculated from constraint_bounds.""" + + sphere_radius: float | None = None + """Radius for sphere constraint. If None and constraint_type='sphere', auto-calculated from constraint_bounds.""" + + sphere_radius_mode: str = "inscribed" + """Mode for auto-calculating sphere radius from bounds: 'inscribed' or 'circumscribed'. Only used if sphere_radius is None.""" + + def __post_init__(self): + """Initialize sub-configs with defaults if not provided.""" + if self.sampling is None: + self.sampling = SamplingConfig() + if self.cache is None: + self.cache = CacheConfig() + if self.constraint is None: + self.constraint = DimensionConstraint() + if self.visualization is None: + self.visualization = VisualizationConfig() + if self.metric is None: + self.metric = MetricConfig() + + +class WorkspaceAnalyzer: + """Main workspace analyzer class for robotic manipulation. + + Analyzes the reachable workspace of a robot by sampling joint configurations, + computing forward kinematics, and generating metrics and visualizations. + + Note: + Currently designed for single environment operation (num_envs=1). When multiple + environments are present, the analyzer will use the first environment (index 0) + and log appropriate warnings. Multi-environment support will be added in future versions. + """ + + # Default priority order for control parts selection + DEFAULT_CONTROL_PART_PRIORITY = ["left_arm", "right_arm"] + + def __init__( + self, + robot: Robot, + config: WorkspaceAnalyzerConfig | None = None, + sim_manager: SimulationManager | None = None, + ): + """Initialize the workspace analyzer. + + Args: + robot: Robot instance to analyze. + config: Configuration object. If None, uses defaults. + sim_manager: SimulationManager instance. Defaults None. + """ + self.robot = robot + self.config = config or WorkspaceAnalyzerConfig() + self.sim_manager = sim_manager + + # Check multi-environment compatibility and add protection + self._check_num_envs_compatibility() + + # Use sim_manager's device if available, otherwise default to CPU + self.device = ( + sim_manager.device if sim_manager is not None else torch.device("cpu") + ) + + # Determine control part name from config + self.control_part_name = self._determine_control_part( + self.config.control_part_name + ) + + # Extract joint limits from robot + self._setup_joint_limits() + + # Initialize components + self.sampler = self._create_sampler() + self.cache = self._create_cache() + self.constraint_checker = self._create_constraint_checker() + + # Storage for analysis results + self.workspace_points: torch.Tensor | None = None + self.joint_configurations: torch.Tensor | None = None + self.metrics_results: Dict[str, Any] = {} + self.current_mode: AnalysisMode | None = None + self.success_rates: torch.Tensor | None = None + + def _determine_control_part(self, control_part_name: str | None) -> str | None: + """Determine the control part name to use. + + Args: + control_part_name: User-specified control part name, or None. + + Returns: + The control part name to use, or None for default solver. + """ + if control_part_name is not None: + # User explicitly specified a control part + logger.log_info(f"Using user-specified control part: {control_part_name}") + return control_part_name + + # Try to find a suitable default control part + if hasattr(self.robot, "cfg") and hasattr(self.robot.cfg, "control_parts"): + control_parts = self.robot.cfg.control_parts + if control_parts: + # Try priority parts first + for part in self.DEFAULT_CONTROL_PART_PRIORITY: + if part in control_parts: + logger.log_info(f"Auto-selected control part: {part}") + return part + + # If no priority part found, use the first available + first_part = next(iter(control_parts.keys())) + logger.log_info( + f"Auto-selected first available control part: {first_part}" + ) + return first_part + + # Fall back to None (will use default solver) + logger.log_info("No specific control part specified, using default solver") + return None + + def _check_num_envs_compatibility(self) -> None: + """Check multi-environment compatibility and provide appropriate warnings. + + WorkspaceAnalyzer is currently designed for num_envs=1. If multiple environments + are present, it will use the first environment and log appropriate warnings. + """ + if self.sim_manager is None: + # No sim_manager provided, cannot check num_envs + logger.log_debug( + "No SimulationManager provided, assuming single environment setup" + ) + return + + num_envs = self.sim_manager.num_envs + + if num_envs == 1: + logger.log_debug( + "WorkspaceAnalyzer initialized with single environment (num_envs=1)" + ) + else: + logger.log_warning( + f"WorkspaceAnalyzer is currently designed for single environment operation (num_envs=1), " + f"but {num_envs} environments detected. Will use the first environment (index 0) for analysis. " + f"Multi-environment support will be added in future versions." + ) + logger.log_info("Using environment index 0 for workspace analysis") + + def _setup_joint_limits(self) -> None: + """Extract and setup joint limits from the robot.""" + # Get all joint limits from robot (using first entity/environment) + all_joint_limits = self.robot.body_data.qpos_limits[0] + + # If control_part_name is specified, get only the joints for that part + if self.control_part_name is not None: + joint_ids = self.robot.get_joint_ids(self.control_part_name) + self.qpos_limits = all_joint_limits[joint_ids] + logger.log_info( + f"Using {len(joint_ids)} joints from control part '{self.control_part_name}'" + ) + else: + # Use all joints + self.qpos_limits = all_joint_limits + logger.log_info("Using all robot joints (no control part specified)") + + # Apply scaling factor if specified + if self.config.constraint.joint_limits_scale != 1.0: + scale = self.config.constraint.joint_limits_scale + center = (self.qpos_limits[:, 0] + self.qpos_limits[:, 1]) / 2 + range_half = (self.qpos_limits[:, 1] - self.qpos_limits[:, 0]) / 2 * scale + + self.qpos_limits[:, 0] = center - range_half + self.qpos_limits[:, 1] = center + range_half + + logger.log_info(f"Joint limits scaled by factor: {scale}") + + self.num_joints = len(self.qpos_limits) + logger.log_debug(f"Number of joints: {self.num_joints}") + + def _create_sampler(self) -> BaseSampler: + """Create sampler based on configuration.""" + factory = SamplerFactory() + + return factory.create_sampler( + strategy=self.config.sampling.strategy, + seed=self.config.sampling.seed, + ) + + # Note: Geometric constraint creation methods temporarily removed + + # Note: Explicit constraint creation temporarily removed + + def _compute_dynamic_workspace_bounds(self) -> torch.Tensor: + """Compute workspace bounds dynamically from joint space FK. + + Returns: + Tensor of shape (3, 2) representing [[x_min, x_max], [y_min, y_max], [z_min, z_max]] + """ + logger.log_info("Computing workspace bounds dynamically from joint space FK...") + + # Create a temporary sampler without constraints for initial FK computation + from embodichain.lab.sim.utility.workspace_analyzer.samplers import ( + RandomSampler, + ) + + temp_sampler = RandomSampler(seed=self.config.sampling.seed) + + # Sample joint space to compute FK bounds + joint_samples = temp_sampler.sample(num_samples=1000, bounds=self.qpos_limits) + + # Compute FK for all samples with progress tracking + workspace_pts_list = [] + + pbar = self._create_optimized_tqdm( + range(len(joint_samples)), + desc="Computing Workspace Bounds (FK)", + unit="cfg", + color="cyan", + emoji="📏", + ) + + successful_fk = 0 + for i in pbar: + qpos = joint_samples[i : i + 1] # Keep batch dimension + try: + pose = self.robot.compute_fk( + qpos=qpos, + name=self.control_part_name, + to_matrix=True, + ) + position = pose[:, :3, 3] # Extract position + workspace_pts_list.append(position) + successful_fk += 1 + except Exception: + continue + + # Update progress bar with success rate + self._update_progress_with_stats( + pbar, i, successful_fk, metric_name="FK success", show_rate=True + ) + + if workspace_pts_list: + workspace_pts = torch.cat(workspace_pts_list, dim=0) + # Compute min/max bounds for each dimension + min_bounds = workspace_pts.min(dim=0).values + max_bounds = workspace_pts.max(dim=0).values + + # Add margin (10%) + margin = (max_bounds - min_bounds) * 0.1 + min_bounds = min_bounds - margin + max_bounds = max_bounds + margin + + # Create bounds tensor: [[x_min, x_max], [y_min, y_max], [z_min, z_max]] + bounds = torch.stack([min_bounds, max_bounds], dim=1) + + logger.log_info( + f"Computed workspace bounds from {len(workspace_pts)} FK samples:\n" + f"\t X: [{min_bounds[0]:.3f}, {max_bounds[0]:.3f}] m\n" + f"\t Y: [{min_bounds[1]:.3f}, {max_bounds[1]:.3f}] m\n" + f"\t Z: [{min_bounds[2]:.3f}, {max_bounds[2]:.3f}] m" + ) + return bounds + else: + # Fallback to default bounds if FK computation fails + logger.log_warning("FK computation failed, using fallback bounds") + return torch.tensor( + [[-1.0, 1.0], [-1.0, 1.0], [0.0, 2.0]], device=self.device + ) + + def _create_cache(self): + """Create cache manager based on configuration.""" + if not self.config.cache.enabled: + return None + return CacheManager.create_cache_from_config(self.config.cache) + + def _create_constraint_checker(self) -> WorkspaceConstraintChecker: + """Create constraint checker based on configuration.""" + return WorkspaceConstraintChecker.from_config( + self.config.constraint, device=self.device + ) + + def _create_optimized_tqdm( + self, iterable, desc: str, unit: str, color: str = "blue", emoji: str = "⚡" + ): + """Create an optimized tqdm progress bar with adaptive updates and smart formatting. + + Args: + iterable: The iterable to track progress for + desc: Description text + unit: Unit name (e.g., 'cfg', 'pt') + color: Progress bar color + emoji: Emoji for the description + + Returns: + Configured tqdm instance + """ + total = len(iterable) if hasattr(iterable, "__len__") else None + + # Adaptive parameters based on total count + if total: + if total < 100: + mininterval, maxinterval = 0.1, 1.0 + smoothing = 0.1 + elif total < 1000: + mininterval, maxinterval = 0.2, 2.0 + smoothing = 0.05 + else: + mininterval, maxinterval = 0.5, 5.0 + smoothing = 0.02 + else: + mininterval, maxinterval = 0.5, 5.0 + smoothing = 0.05 + + # Terminal width detection + try: + terminal_width = os.get_terminal_size().columns + ncols = min(120, max(80, terminal_width - 10)) + except OSError: + # Terminal size unavailable (e.g., non-terminal environment) + ncols = 100 + + # Color codes for different states + color_codes = { + "blue": "\033[34m", + "cyan": "\033[36m", + "magenta": "\033[35m", + "green": "\033[32m", + "yellow": "\033[33m", + "red": "\033[31m", + } + + # Enhanced bar format with better spacing + bar_format = ( + f"{color_codes.get(color, '')}{{desc}}\033[0m: " + f"{{percentage:3.0f}}%|{{bar}}| {{n_fmt}}/{{total_fmt}} " + f"[{{elapsed}}<{{remaining}}, {{rate_fmt}}{{postfix}}]" + ) + + # Performance-aware tqdm configuration + pbar = tqdm( + iterable, + desc=f"{emoji} {desc}", + unit=unit, + unit_scale=True, + smoothing=smoothing, + mininterval=mininterval, + maxinterval=maxinterval, + bar_format=bar_format, + ncols=ncols, + dynamic_ncols=True, + colour=color, + leave=True, + ascii=False if sys.stdout.encoding == "utf-8" else True, + # Advanced features + position=0, # Top position for multiple bars + file=sys.stdout, + disable=False, + ) + + # Add performance tracking attributes + pbar._start_time = time.time() + pbar._last_update = 0 + pbar._update_count = 0 + + return pbar + + def _update_progress_with_stats( + self, + pbar, + current_idx: int, + success_count: int, + metric_name: str = "success", + show_rate: bool = True, + ): + """Update progress bar with intelligent statistics and color coding. + + Args: + pbar: tqdm progress bar instance + current_idx: Current iteration index + success_count: Number of successful operations + metric_name: Name of the metric being tracked + show_rate: Whether to show the success rate + """ + if not show_rate: + return + + total_processed = current_idx + 1 + rate = (success_count / total_processed) * 100 if total_processed > 0 else 0 + + # Intelligent color coding with thresholds + if rate >= 85: + color, icon = "\033[92m", "🟢" # Bright green, excellent + elif rate >= 70: + color, icon = "\033[32m", "✅" # Green, good + elif rate >= 50: + color, icon = "\033[93m", "🟡" # Bright yellow, moderate + elif rate >= 30: + color, icon = "\033[33m", "🟠" # Yellow, low + else: + color, icon = "\033[91m", "🔴" # Bright red, poor + + # Adaptive display with performance metrics + current_time = time.time() + + # Smart update throttling based on performance + if hasattr(pbar, "_last_update") and hasattr(pbar, "_update_count"): + time_since_last = current_time - pbar._last_update + pbar._update_count += 1 + + # Adaptive update frequency based on processing speed + if pbar._update_count > 100: # After first 100 updates + avg_time_per_update = time_since_last / max( + 1, + ( + pbar._update_count - pbar._last_update_count + if hasattr(pbar, "_last_update_count") + else 1 + ), + ) + if avg_time_per_update < 0.01: # Very fast processing + update_threshold = 0.5 # Update every 0.5s + elif avg_time_per_update < 0.1: # Medium speed + update_threshold = 0.3 # Update every 0.3s + else: # Slow processing + update_threshold = 0.1 # Update every 0.1s + + if time_since_last < update_threshold: + return # Skip update to reduce overhead + + pbar._last_update = current_time + pbar._last_update_count = pbar._update_count + + # Enhanced display with ETA and throughput + if total_processed < 10: + # Show individual counts for small numbers + stats = f" {icon} {success_count}/{total_processed}" + else: + # Show percentage and throughput for larger numbers + if hasattr(pbar, "_start_time"): + elapsed = current_time - pbar._start_time + throughput = total_processed / elapsed if elapsed > 0 else 0 + if throughput > 10: + stats = f" {icon} {color}{rate:.1f}%\033[0m {metric_name} ({throughput:.0f}/s)" + else: + stats = f" {icon} {color}{rate:.1f}%\033[0m {metric_name} ({throughput:.1f}/s)" + else: + stats = f" {icon} {color}{rate:.1f}%\033[0m {metric_name}" + + pbar.set_postfix_str(stats, refresh=False) + + def sample_joint_space(self, num_samples: int | None = None) -> torch.Tensor: + """Sample joint configurations within joint limits. + + Args: + num_samples: Number of samples to generate. If None, uses config value. + + Returns: + Tensor of shape (num_samples, num_joints) containing joint configurations. + """ + num_samples = num_samples or self.config.sampling.num_samples + + # Performance-aware sampling with progress indication + start_time = time.time() + + # Sample from joint space + joint_samples = self.sampler.sample( + bounds=self.qpos_limits, num_samples=num_samples + ) + + sampling_time = time.time() - start_time + samples_per_sec = ( + num_samples / sampling_time if sampling_time > 0 else float("inf") + ) + + logger.log_info( + f"Generated {num_samples} joint space samples " + f"({samples_per_sec:.0f} samples/s)" + ) + return joint_samples + + def sample_cartesian_space(self, num_samples: int | None = None) -> torch.Tensor: + """Sample Cartesian positions within workspace bounds. + + Args: + num_samples: Number of samples to generate. If None, uses config value. + + Returns: + Tensor of shape (num_samples, 3) containing Cartesian positions. + """ + num_samples = num_samples or self.config.sampling.num_samples + + # Determine Cartesian bounds + if ( + self.config.constraint.min_bounds is not None + and self.config.constraint.max_bounds is not None + ): + cartesian_bounds = torch.stack( + [ + torch.tensor(self.config.constraint.min_bounds, device=self.device), + torch.tensor(self.config.constraint.max_bounds, device=self.device), + ], + dim=1, + ) + else: + # Compute bounds from joint space FK using dedicated method + logger.log_info( + "No Cartesian bounds specified, computing from joint space..." + ) + cartesian_bounds = self._compute_dynamic_workspace_bounds() + + # Sample from Cartesian space using bounds + cartesian_samples = self.sampler.sample( + bounds=cartesian_bounds, num_samples=num_samples + ) + + # Check how many samples pass workspace constraints + valid_bounds = self.constraint_checker.check_bounds(cartesian_samples) + valid_collision = self.constraint_checker.check_collision(cartesian_samples) + valid_constraints = valid_bounds & valid_collision + + constraint_pass_rate = valid_constraints.sum().item() / num_samples * 100 + exclude_zones_count = self.constraint_checker.get_num_exclude_zones() + + logger.log_info( + f"Generated {num_samples} Cartesian space samples. " + f"Constraint check: {valid_constraints.sum()}/{num_samples} " + f"({constraint_pass_rate:.1f}%) pass bounds+collision constraints " + f"({exclude_zones_count} exclude zones configured)" + ) + return cartesian_samples + + def sample_plane( + self, + num_samples: int | None = None, + plane_normal: torch.Tensor | None = None, + plane_point: torch.Tensor | None = None, + plane_bounds: torch.Tensor | None = None, + ) -> torch.Tensor: + """Sample points on a specified plane using existing samplers (ultra-simplified version). + + Args: + num_samples: Number of samples to generate. If None, uses config value. + plane_normal: Plane normal vector [nx, ny, nz]. Defaults to [0,0,1] (XY plane). + plane_point: A point on the plane [x, y, z]. Defaults to [0,0,0]. + plane_bounds: 2D bounds [[u_min, u_max], [v_min, v_max]]. Defaults to [[-1,1], [-1,1]]. + + Returns: + Tensor of shape (num_samples, 3) containing 3D points on the plane. + """ + num_samples = num_samples or self.config.sampling.num_samples + + # Set default values + if plane_normal is None: + plane_normal = torch.tensor([0.0, 0.0, 1.0], device=self.device) # XY plane + else: + plane_normal = plane_normal.to(self.device) / torch.norm( + plane_normal.to(self.device) + ) + + if plane_point is None: + plane_point = torch.tensor([0.0, 0.0, 0.0], device=self.device) + else: + plane_point = plane_point.to(self.device) + + if plane_bounds is None: + # Compute dynamic workspace bounds from joint space FK + dynamic_bounds_3d = self._compute_dynamic_workspace_bounds() + + # Project 3D bounds to 2D plane coordinate system + plane_bounds = self._compute_plane_bounds_from_3d( + dynamic_bounds_3d, plane_normal, plane_point + ) + + logger.log_info( + f"Using dynamic plane bounds computed from FK: " + f"U: [{plane_bounds[0, 0]:.3f}, {plane_bounds[0, 1]:.3f}], " + f"V: [{plane_bounds[1, 0]:.3f}, {plane_bounds[1, 1]:.3f}] " + f"(projected to plane with normal {plane_normal.cpu().numpy()})" + ) + else: + plane_bounds = plane_bounds.to(self.device) + + # Generate 2D samples and convert to 3D + plane_samples_2d = self.sampler.sample(num_samples, bounds=plane_bounds) + plane_samples_3d = self._plane_to_world_optimized( + plane_samples_2d, plane_normal, plane_point + ) + + logger.log_info( + f"Generated {num_samples} plane samples using {self.sampler.get_strategy_name()}" + ) + + return plane_samples_3d + + def _project_to_plane( + self, + points_3d: torch.Tensor, + plane_normal: torch.Tensor, + plane_point: torch.Tensor, + ) -> torch.Tensor: + """Project 3D points onto a specified plane. + + Args: + points_3d: 3D points to project, shape (num_samples, 3) + plane_normal: Normal vector of the plane [nx, ny, nz] + plane_point: A point on the plane [x, y, z] + + Returns: + Projected 3D points on the plane, shape (num_samples, 3) + """ + # Normalize the plane normal + plane_normal = plane_normal / torch.norm(plane_normal) + + # Vector from plane_point to each 3D point + vectors_to_points = points_3d - plane_point.unsqueeze(0) + + # Project vectors onto plane normal (signed distance from plane) + distances = torch.sum(vectors_to_points * plane_normal.unsqueeze(0), dim=1) + + # Project points onto plane by subtracting the normal component + projected_points = points_3d - distances.unsqueeze(1) * plane_normal.unsqueeze( + 0 + ) + + return projected_points + + def _plane_to_world_optimized( + self, + plane_coords: torch.Tensor, + plane_normal: torch.Tensor, + plane_point: torch.Tensor, + ) -> torch.Tensor: + """Convert 2D plane coordinates to 3D world coordinates with optimized basis generation. + + This method uses a more numerically stable approach to generate orthogonal basis vectors + and supports orientation optimization for better workspace coverage. + + Args: + plane_coords: 2D coordinates on the plane, shape (num_samples, 2) + plane_normal: Normal vector of the plane [nx, ny, nz] + plane_point: A point on the plane [x, y, z] + + Returns: + 3D world coordinates, shape (num_samples, 3) + """ + num_samples = plane_coords.shape[0] + + # Generate orthogonal basis vectors using improved method + u, v = self._generate_orthogonal_basis(plane_normal) + + # Convert 2D plane coordinates to 3D with vectorized operations + world_coords = ( + plane_point.unsqueeze(0) + + plane_coords[:, 0:1] # Base point broadcast to all samples + * u.unsqueeze(0) + + plane_coords[:, 1:2] # First plane direction + * v.unsqueeze(0) # Second plane direction + ) + + return world_coords + + def _generate_orthogonal_basis( + self, plane_normal: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Generate orthogonal basis vectors for a plane with improved numerical stability. + + This method uses the most stable approach based on the plane normal direction + and optionally optimizes orientation for workspace coverage. + + Args: + plane_normal: Normal vector of the plane [nx, ny, nz] + + Returns: + Tuple of two orthogonal unit vectors (u, v) that span the plane + """ + # Find the coordinate with smallest absolute value for numerical stability + abs_normal = torch.abs(plane_normal) + min_idx = torch.argmin(abs_normal) + + # Create an arbitrary vector with 1 in the most stable coordinate + arbitrary = torch.zeros_like(plane_normal) + arbitrary[min_idx] = 1.0 + + # Generate first tangent vector using Gram-Schmidt + u = arbitrary - torch.dot(arbitrary, plane_normal) * plane_normal + u = u / torch.norm(u) + + # Generate second tangent vector via cross product + v = torch.linalg.cross(plane_normal, u) + v = v / torch.norm(v) + + return u, v + + def _compute_plane_bounds_from_3d( + self, + bounds_3d: torch.Tensor, + plane_normal: torch.Tensor, + plane_point: torch.Tensor, + ) -> torch.Tensor: + """Project 3D workspace bounds to 2D plane coordinate system. + + Args: + bounds_3d: 3D workspace bounds, shape (3, 2) representing + [[x_min, x_max], [y_min, y_max], [z_min, z_max]] + plane_normal: Normal vector of the plane [nx, ny, nz] + plane_point: A point on the plane [x, y, z] + + Returns: + 2D plane bounds, shape (2, 2) representing [[u_min, u_max], [v_min, v_max]] + """ + # Generate orthogonal basis vectors for the plane + u, v = self._generate_orthogonal_basis(plane_normal) + + # Create all 8 corners of the 3D bounding box + corners_3d = [] + for x in [bounds_3d[0, 0], bounds_3d[0, 1]]: # x_min, x_max + for y in [bounds_3d[1, 0], bounds_3d[1, 1]]: # y_min, y_max + for z in [bounds_3d[2, 0], bounds_3d[2, 1]]: # z_min, z_max + corners_3d.append(torch.tensor([x, y, z], device=self.device)) + + corners_3d = torch.stack(corners_3d) # Shape: (8, 3) + + # Project all corners to the plane coordinate system + # Transform from world coordinates to plane coordinates + corners_relative = corners_3d - plane_point.unsqueeze( + 0 + ) # Relative to plane point + + # Project onto plane basis vectors + u_coords = torch.sum(corners_relative * u.unsqueeze(0), dim=1) # Shape: (8,) + v_coords = torch.sum(corners_relative * v.unsqueeze(0), dim=1) # Shape: (8,) + + # Find min/max in both plane directions + u_min, u_max = u_coords.min(), u_coords.max() + v_min, v_max = v_coords.min(), v_coords.max() + + # Create 2D plane bounds + plane_bounds = torch.tensor( + [[u_min, u_max], [v_min, v_max]], device=self.device + ) + + return plane_bounds + + def _get_robot_base_position(self) -> torch.Tensor: + """Get the robot base position as default plane point. + + Returns: + Robot base position as a 3D tensor + """ + try: + # Try to get current robot pose (using first environment) + current_pose = self.robot.compute_fk( + qpos=self.robot.get_qpos()[None, :], # Add batch dimension + name=self.control_part_name, + to_matrix=True, + ) + # Use current end-effector position projected to a reasonable height + base_pos = current_pose[0, :3, 3].clone() + base_pos[2] = 0.0 # Project to ground plane + return base_pos + except Exception: + # Fallback to origin + return torch.tensor([0.0, 0.0, 0.0], device=self.device) + + def _generate_plane_samples( + self, + plane_bounds: torch.Tensor, + num_samples: int, + ) -> torch.Tensor: + """Generate 2D plane samples using existing base sampler directly.""" + return self.sampler.sample(num_samples, bounds=plane_bounds) + + def compute_workspace_points( + self, joint_configs: torch.Tensor, batch_size: int | None = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute end-effector positions for given joint configurations. + + Args: + joint_configs: Joint configurations, shape (num_samples, num_joints). + batch_size: Batch size for FK computation. If None, uses config value. + + Returns: + Tuple of: + - workspace_points: End-effector positions, shape (num_valid, 3) + - valid_configs: Valid joint configurations, shape (num_valid, num_joints) + """ + num_samples = len(joint_configs) + + workspace_points_list = [] + valid_configs_list = [] + + logger.log_info(f"Computing FK for {num_samples} samples...") + + # Track valid points for progress bar + total_valid = 0 + + # Robot expects one configuration at a time (batch_size from robot environments, not samples) + # Process each configuration individually + pbar = self._create_optimized_tqdm( + range(num_samples), + desc="Forward Kinematics", + unit="cfg", + color="cyan", + emoji="🤖", + ) + for i in pbar: + qpos = joint_configs[i : i + 1] # Keep batch dimension + + try: + # Compute forward kinematics + pose = self.robot.compute_fk( + qpos=qpos, + name=self.control_part_name, + to_matrix=True, + ) + + # Extract position (x, y, z) + position = pose[:, :3, 3] # Shape: (1, 3) + + # Filter by constraints (bounds + collision check) + valid_bounds = self.constraint_checker.check_bounds(position) + valid_collision = self.constraint_checker.check_collision(position) + valid_mask = valid_bounds & valid_collision + + # Store valid results + if valid_mask.any(): + workspace_points_list.append(position[valid_mask]) + valid_configs_list.append(qpos[valid_mask]) + total_valid += 1 + + # Update progress bar with intelligent statistics + self._update_progress_with_stats( + pbar, i, total_valid, metric_name="valid", show_rate=True + ) + + except Exception as e: + logger.log_warning(f"FK computation failed for sample {i}: {e}") + continue + + # Concatenate all results + if workspace_points_list: + workspace_points = torch.cat(workspace_points_list, dim=0) + valid_configs = torch.cat(valid_configs_list, dim=0) + else: + workspace_points = torch.empty((0, 3), device=self.device) + valid_configs = torch.empty((0, self.num_joints), device=self.device) + + # Performance summary for FK computation + pbar.close() # Ensure progress bar is closed + success_rate = len(workspace_points) / num_samples * 100 + + # Performance indicator based on success rate + if success_rate >= 90: + perf_icon = "🏆" # Trophy for excellent performance + elif success_rate >= 75: + perf_icon = "✅" # Check mark for good performance + elif success_rate >= 50: + perf_icon = "🟡" # Yellow circle for moderate performance + else: + perf_icon = "⚠️" # Warning for low performance + + logger.log_info( + f"{perf_icon} FK Results: {len(workspace_points)}/{num_samples} valid points " + f"({success_rate:.1f}% success rate)" + ) + + return workspace_points, valid_configs + + def compute_reachability( + self, cartesian_points: torch.Tensor, batch_size: int | None = None + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute reachability for Cartesian points using IK. + + Args: + cartesian_points: Cartesian positions, shape (num_samples, 3). + batch_size: Batch size for IK computation. If None, uses config value. + + Returns: + Tuple of: + - all_points: All Cartesian positions, shape (num_samples, 3) + - reachable_points: Reachable positions, shape (num_reachable, 3) + - success_rates: IK success rate for each point, shape (num_samples,) + - reachability_mask: Boolean mask indicating reachable points, shape (num_samples,) + - best_configs: Best joint configurations, shape (num_reachable, num_joints) + """ + num_samples = len(cartesian_points) + ik_samples_per_point = self.config.ik_samples_per_point + + # Pre-filter Cartesian points by workspace constraints + # This eliminates points that are outside bounds or in collision zones + valid_cartesian_mask = self.constraint_checker.check_bounds( + cartesian_points + ) & self.constraint_checker.check_collision(cartesian_points) + + logger.log_info( + f"Pre-filtered Cartesian points: {valid_cartesian_mask.sum()}/{num_samples} " + f"points pass workspace constraints ({(valid_cartesian_mask.sum()/num_samples*100):.1f}%)" + ) + + # Store results for all points (including invalid ones for consistent indexing) + all_success_rates = torch.zeros(num_samples, device=self.device) + reachable_points_list = [] + best_configs_list = [] + + logger.log_info( + f"Computing IK for {num_samples} Cartesian samples " + f"({ik_samples_per_point} seeds per point)..." + ) + + # Create a random sampler for generating IK seeds (avoid UniformSampler issues) + from embodichain.lab.sim.utility.workspace_analyzer.samplers import ( + RandomSampler, + ) + + random_sampler = RandomSampler(seed=self.config.sampling.seed) + + # Get reference end-effector pose for IK target orientation + # Priority: use reference_pose if provided, otherwise compute from current joint configuration + if ( + hasattr(self.config, "reference_pose") + and self.config.reference_pose is not None + ): + # Use provided reference pose (should be 4x4 transformation matrix) + reference_pose = self.config.reference_pose + if isinstance(reference_pose, np.ndarray): + reference_pose = torch.from_numpy(reference_pose).to(self.device) + if reference_pose.dim() == 2: # Shape: (4, 4) -> (1, 4, 4) + reference_pose = reference_pose.unsqueeze(0) + current_ee_pose = reference_pose # Shape: (1, 4, 4) + logger.log_info("Using provided reference pose for IK target orientation") + else: + # Fallback: compute current end-effector pose from joint configuration + try: + # Using first environment (index 0) for qpos retrieval + current_qpos = self.robot.get_qpos()[0][ + self.robot.get_joint_ids(self.control_part_name) + ] + current_ee_pose = self.robot.compute_fk( + name=self.control_part_name, + qpos=current_qpos.unsqueeze(0), + to_matrix=True, + ) # Shape: (1, 4, 4) + logger.log_info( + "Computing reference pose from current robot configuration" + ) + except Exception as e: + logger.log_warning(f"Failed to compute current robot pose: {e}") + # Create identity pose as fallback + current_ee_pose = torch.eye(4, device=self.device).unsqueeze(0) + current_ee_pose[0, :3, 3] = torch.tensor( + [0.5, 0.0, 1.0], device=self.device + ) # Default position + logger.log_info("Using default identity pose as fallback") + + # Print current joint configuration and computed pose + pose_np = current_ee_pose[0].cpu().numpy() + position = pose_np[:3, 3] + rotation_matrix = pose_np[:3, :3] + + # Convert rotation matrix to Euler angles + import scipy.spatial.transform as spt + + euler_angles = spt.Rotation.from_matrix(rotation_matrix).as_euler( + "xyz", degrees=True + ) + + # Print detailed reference pose information + pose_np = current_ee_pose[0].cpu().numpy() + position = pose_np[:3, 3] + rotation_matrix = pose_np[:3, :3] + + # Convert rotation matrix to Euler angles (ZYX convention) + import scipy.spatial.transform as spt + + euler_angles = spt.Rotation.from_matrix(rotation_matrix).as_euler( + "xyz", degrees=True + ) + + # Format matrix with proper indentation + matrix_lines = np.array2string(pose_np, precision=4, suppress_small=True).split( + "\n" + ) + matrix_str = "\n".join(f"\t {line}" for line in matrix_lines) + logger.log_info( + f"🎯 Using provided reference pose for IK target orientation:\n" + f"\t Position: [{position[0]:.4f}, {position[1]:.4f}, {position[2]:.4f}] m\n" + f"\t Rotation (XYZ Euler): [{euler_angles[0]:.2f}°, {euler_angles[1]:.2f}°, {euler_angles[2]:.2f}°]\n" + f"\t Matrix:\n{matrix_str}" + ) + + # Track statistics for progress bar + total_reachable = 0 + + # Process each point individually (robot expects batch_size from environments, not samples) + pbar = self._create_optimized_tqdm( + range(num_samples), + desc="Inverse Kinematics", + unit="pt", + color="magenta", + emoji="🎯", + ) + + for i in pbar: + position = cartesian_points[i] # Shape: (3,) + + # Skip points that don't satisfy workspace constraints + if not valid_cartesian_mask[i]: + # Mark as unreachable due to constraint violation + all_success_rates[i] = 0.0 + # Update progress bar + reachability_rate = total_reachable / (i + 1) * 100 + if reachability_rate >= 70: + reach_color = "\033[32m" # Green for high reachability + elif reachability_rate >= 40: + reach_color = "\033[33m" # Yellow for medium reachability + else: + reach_color = "\033[31m" # Red for low reachability + pbar.set_postfix_str( + f"🎯 Reachable: {total_reachable}/{i+1} | {reach_color}{reachability_rate:.1f}%\033[0m rate (❌ constraint)" + ) + continue + + # Create target pose: use current orientation, replace position with sampled position + pose = current_ee_pose.clone() + pose[0, :3, 3] = position + + # Try multiple random seeds for this point + success_count = 0 + best_qpos = None + + logger.set_log_level("ERROR") # Suppress warnings during IK attempts + for seed_idx in range(ik_samples_per_point): + # Generate random joint seed using RandomSampler + random_seed = random_sampler.sample( + bounds=self.qpos_limits, num_samples=1 + ) # Shape: (1, num_joints) + + try: + # Compute IK + ret, qpos = self.robot.compute_ik( + pose=pose, + joint_seed=random_seed, + name=self.control_part_name, + ) + + # Count successes + if ret is not None and ret[0]: + success_count += 1 + # Store first successful configuration + if best_qpos is None: + best_qpos = qpos[0] # Extract from batch dimension + + except Exception as e: + logger.log_warning( + f"IK computation failed for sample {i}, seed {seed_idx}: {e}" + ) + continue + logger.set_log_level("INFO") # Restore log level + + # Calculate success rate for this point + success_rate = success_count / ik_samples_per_point + all_success_rates[i] = success_rate + + # Filter by success threshold for reachable points + if success_rate and best_qpos is not None: + reachable_points_list.append(position.unsqueeze(0)) # Add batch dim + best_configs_list.append(best_qpos.unsqueeze(0)) # Add batch dim + total_reachable += 1 + + # Update progress bar with reachability statistics + reachability_rate = total_reachable / (i + 1) * 100 + # Use color coding for the reachability rate + if reachability_rate >= 70: + reach_color = "\033[32m" # Green for high reachability + elif reachability_rate >= 40: + reach_color = "\033[33m" # Yellow for medium reachability + else: + reach_color = "\033[31m" # Red for low reachability + + # Add success rate indicator for this specific point + if success_rate: + point_status = "✅ IK" + elif success_rate > 0: + point_status = f"🟡 IK({success_rate:.1f})" + else: + point_status = "❌ IK" + + pbar.set_postfix_str( + f"🎯 Reachable: {total_reachable}/{i+1} | {reach_color}{reachability_rate:.1f}%\033[0m rate | {point_status}" + ) + + # Concatenate reachable results + if reachable_points_list: + reachable_points = torch.cat(reachable_points_list, dim=0) + best_configs = torch.cat(best_configs_list, dim=0) + else: + reachable_points = torch.empty((0, 3), device=self.device) + best_configs = torch.empty((0, self.num_joints), device=self.device) + + # Create reachability mask + reachability_mask = all_success_rates > 0 + + # Performance summary for IK computation + pbar.close() # Ensure progress bar is closed + reachability = len(reachable_points) / num_samples * 100 + + # Reachability performance indicator + if reachability >= 80: + reach_icon = "🏆" # Trophy for high reachability + elif reachability >= 60: + reach_icon = "🚀" # Rocket for good reachability + elif reachability >= 40: + reach_icon = "🟡" # Yellow for moderate reachability + elif reachability >= 20: + reach_icon = "🟠" # Orange for low reachability + else: + reach_icon = "⚠️" # Warning for very low reachability + + logger.log_info( + f"{reach_icon} IK Results: {len(reachable_points)}/{num_samples} reachable points " + f"({reachability:.1f}% reachability)" + ) + + return ( + cartesian_points, + reachable_points, + all_success_rates, + reachability_mask, + best_configs, + ) + + def analyze( + self, + num_samples: int | None = None, + force_recompute: bool = False, + visualize: bool = False, + ) -> Dict[str, Any]: + """Perform complete workspace analysis. + + Args: + num_samples: Number of samples to generate. If None, uses config value. + force_recompute: If True, recomputes even if cached results exist. + visualize: If True, visualizes the workspace points. Prefers sim_manager visualization + if available, otherwise falls back to visualizers module. + + Returns: + Dictionary containing analysis results. + """ + logger.log_info("Starting Workspace Analysis...") + + start_time = time.time() + + # Check cache + if not force_recompute and self.cache is not None: + cached_results = self._load_from_cache() + if cached_results is not None: + logger.log_info("Loaded results from cache") + return cached_results + + # Choose analysis mode + if self.config.mode == AnalysisMode.JOINT_SPACE: + # Joint space mode: Sample joints → FK → Workspace points + logger.log_info(f"Mode: {AnalysisMode.JOINT_SPACE.value}") + + # Step 1: Sample joint space + logger.log_info("[1/3] Sampling joint space...") + joint_configs = self.sample_joint_space(num_samples) + + # Step 2: Compute workspace points + logger.log_info("[2/3] Computing workspace points via computing FK...") + workspace_points, valid_configs = self.compute_workspace_points( + joint_configs + ) + + # Store results + self.workspace_points = workspace_points + self.joint_configurations = valid_configs + self.current_mode = AnalysisMode.JOINT_SPACE + self.success_rates = None # All points are reachable in joint space mode + + # Add constraint check statistics + constraint_stats = self._compute_constraint_statistics(workspace_points) + + results = { + "mode": AnalysisMode.JOINT_SPACE.value, + "workspace_points": workspace_points, + "joint_configurations": valid_configs, + "num_samples": num_samples or self.config.sampling.num_samples, + "num_valid": len(workspace_points), + "constraint_statistics": constraint_stats, + } + + elif self.config.mode == AnalysisMode.CARTESIAN_SPACE: + # Cartesian space mode: Sample Cartesian → IK → Verify reachability + logger.log_info(f"Mode: {AnalysisMode.CARTESIAN_SPACE.value}") + + # Step 1: Sample Cartesian space + logger.log_info("[1/3] Sampling Cartesian space...") + cartesian_samples = self.sample_cartesian_space(num_samples) + + # Step 2: Compute reachability via IK + logger.log_info("[2/3] Computing reachability via computing IK...") + ( + all_points, + reachable_points, + success_rates, + reachability_mask, + best_configs, + ) = self.compute_reachability(cartesian_samples) + + # Store results - now storing all points for visualization + self.workspace_points = all_points # Store all sampled points + self.reachable_points = reachable_points # Store only reachable points + self.joint_configurations = best_configs + self.current_mode = AnalysisMode.CARTESIAN_SPACE + self.success_rates = success_rates # Store success rates for all points + self.reachability_mask = reachability_mask # Store reachability mask + + # Add constraint check statistics for both all_points and reachable_points + constraint_stats_all = self._compute_constraint_statistics(all_points) + constraint_stats_reachable = ( + self._compute_constraint_statistics(reachable_points) + if len(reachable_points) > 0 + else {} + ) + + results = { + "mode": AnalysisMode.CARTESIAN_SPACE.value, + "all_points": all_points, # All sampled Cartesian points + "workspace_points": all_points, # For compatibility + "reachable_points": reachable_points, # Only reachable points + "joint_configurations": best_configs, + "success_rates": success_rates, + "reachability_mask": reachability_mask, + "num_samples": num_samples or self.config.sampling.num_samples, + "num_reachable": len(reachable_points), + "constraint_statistics": { + "all_points": constraint_stats_all, + "reachable_points": constraint_stats_reachable, + }, + } + + elif self.config.mode == AnalysisMode.PLANE_SAMPLING: + # Plane sampling mode: Sample on plane → IK → Verify reachability + logger.log_info(f"Mode: {AnalysisMode.PLANE_SAMPLING.value}") + + # Step 1: Sample on specified plane + logger.log_info("[1/3] Sampling on specified plane...") + cartesian_samples = self.sample_plane( + num_samples=num_samples, + plane_normal=self.config.plane_normal, + plane_point=self.config.plane_point, + plane_bounds=self.config.plane_bounds, + ) + + # Step 2: Compute reachability via IK + logger.log_info("[2/3] Computing reachability via computing IK...") + ( + all_points, + reachable_points, + success_rates, + reachability_mask, + best_configs, + ) = self.compute_reachability(cartesian_samples) + + # Store results + self.workspace_points = all_points + self.reachable_points = reachable_points + self.joint_configurations = best_configs + self.current_mode = AnalysisMode.PLANE_SAMPLING + self.success_rates = success_rates + self.reachability_mask = reachability_mask + + # Add constraint check statistics + constraint_stats_all = self._compute_constraint_statistics(all_points) + constraint_stats_reachable = ( + self._compute_constraint_statistics(reachable_points) + if len(reachable_points) > 0 + else {} + ) + + results = { + "mode": AnalysisMode.PLANE_SAMPLING.value, + "all_points": all_points, # All sampled plane points + "workspace_points": all_points, # For compatibility + "reachable_points": reachable_points, # Only reachable points + "joint_configurations": best_configs, + "success_rates": success_rates, + "reachability_mask": reachability_mask, + "num_samples": len(cartesian_samples), + "num_reachable": len(reachable_points), + "constraint_statistics": { + "all_points": constraint_stats_all, + "reachable_points": constraint_stats_reachable, + }, + "plane_sampling_config": { + "plane_normal": self.config.plane_normal, + "plane_point": self.config.plane_point, + "plane_bounds": self.config.plane_bounds, + }, + } + + else: + raise ValueError(f"Unknown analysis mode: {self.config.mode}") + + # Step 3: Compute metrics (common for both modes) + logger.log_info("[3/3] Computing metrics...") + metrics = self._compute_metrics() + results["metrics"] = metrics + results["config"] = self.config + results["analysis_time"] = time.time() - start_time + + # Cache results + if self.cache is not None: + self._save_to_cache(results) + + # Enhanced completion summary with performance insights + self._log_analysis_summary(results) + + # Visualize if requested + if visualize: + self._visualize_workspace() + + return results + + def _log_analysis_summary(self, results: Dict[str, Any]) -> None: + """Log enhanced analysis summary with performance insights.""" + analysis_time = results["analysis_time"] + mode = results["mode"] + + # Time-based performance indicators + if analysis_time < 30: + time_icon, time_color = "⚡", "\033[92m" # Lightning, bright green + elif analysis_time < 120: + time_icon, time_color = "🚀", "\033[32m" # Rocket, green + elif analysis_time < 300: + time_icon, time_color = "⏱️", "\033[33m" # Clock, yellow + else: + time_icon, time_color = "🐌", "\033[31m" # Snail, red + + logger.log_info( + f"{time_icon} Analysis completed in {time_color}{analysis_time:.2f}s\033[0m" + ) + + if mode == "joint_space": + success_rate = results["num_valid"] / results["num_samples"] * 100 + logger.log_info( + f"📊 Joint Space Results: {results['num_valid']}/{results['num_samples']} " + f"valid points ({success_rate:.1f}% success)" + ) + + # Show constraint statistics + if "constraint_statistics" in results: + stats = results["constraint_statistics"] + logger.log_info( + f"🔒 Constraint Check: Bounds: {stats['bounds_pass_rate']:.1f}% | " + f"Collision: {stats['collision_pass_rate']:.1f}% | " + f"Overall: {stats['overall_pass_rate']:.1f}% " + f"({stats['exclude_zones_count']} exclude zones)" + ) + + elif mode in ["cartesian_space", "plane_sampling"]: + reachability = results["num_reachable"] / results["num_samples"] * 100 + mode_name = ( + "Plane Sampling" if mode == "plane_sampling" else "Cartesian Space" + ) + logger.log_info( + f"📊 {mode_name} Results: {results['num_reachable']}/{results['num_samples']} " + f"reachable points ({reachability:.1f}% reachability)" + ) + + # Show plane sampling specific info + if mode == "plane_sampling" and "plane_sampling_config" in results: + plane_config = results["plane_sampling_config"] + if plane_config: + logger.log_info( + f"🎯 Plane Configuration: Normal: {plane_config['plane_normal']}, " + f"Point: {plane_config['plane_point']}" + ) + + # Show constraint statistics for all points and reachable points + if "constraint_statistics" in results: + all_stats = results["constraint_statistics"]["all_points"] + logger.log_info( + f"🔒 All Points Constraint Check: Bounds: {all_stats['bounds_pass_rate']:.1f}% | " + f"Collision: {all_stats['collision_pass_rate']:.1f}% | " + f"Overall: {all_stats['overall_pass_rate']:.1f}% " + f"({all_stats['exclude_zones_count']} exclude zones)" + ) + + if ( + "reachable_points" in results["constraint_statistics"] + and results["constraint_statistics"]["reachable_points"] + ): + reach_stats = results["constraint_statistics"]["reachable_points"] + logger.log_info( + f"✅ Reachable Points Constraint Check: Bounds: {reach_stats['bounds_pass_rate']:.1f}% | " + f"Collision: {reach_stats['collision_pass_rate']:.1f}% | " + f"Overall: {reach_stats['overall_pass_rate']:.1f}%" + ) + + def _visualize_workspace(self) -> None: + """Visualize the workspace using configured visualization type and backend. + + Uses the vis_type specified in configuration (default: POINT_CLOUD). + Tries multiple backends in order: sim_manager → open3d → matplotlib. + """ + # Early return checks + if self.workspace_points is None or len(self.workspace_points) == 0: + logger.log_warning("No workspace points available for visualization") + return + + if not self.config.visualization.enabled: + logger.log_warning("Visualization is disabled in configuration") + return + + # Define backend priority order + backends = self._get_backend_priority_list() + + # Try each backend in order until one succeeds + for i, backend in enumerate(backends): + try: + logger.log_info(f"Attempting visualization with '{backend}' backend") + self.visualize( + vis_type=self.config.visualization.vis_type, + show=True, + backend=backend, + ) + logger.log_info(f"Successfully visualized with '{backend}' backend") + return + + except Exception as e: + logger.log_warning(f"Failed to visualize with '{backend}' backend: {e}") + + # If this is not the last backend, try the next one + if i < len(backends) - 1: + continue + else: + logger.log_error( + f"All visualization backends failed. " + f"Tried: {', '.join(backends)}" + ) + break + + def _compute_constraint_statistics(self, points: torch.Tensor) -> Dict[str, Any]: + """Compute constraint check statistics for a set of points. + + Args: + points: Tensor of shape (N, 3) containing workspace points. + + Returns: + Dictionary containing constraint statistics. + """ + if len(points) == 0: + return { + "num_points": 0, + "bounds_pass_count": 0, + "bounds_pass_rate": 0.0, + "collision_pass_count": 0, + "collision_pass_rate": 0.0, + "overall_pass_count": 0, + "overall_pass_rate": 0.0, + "exclude_zones_count": self.constraint_checker.get_num_exclude_zones(), + } + + num_points = len(points) + + # Check bounds constraints + bounds_pass = self.constraint_checker.check_bounds(points) + bounds_pass_count = bounds_pass.sum().item() + bounds_pass_rate = bounds_pass_count / num_points * 100 + + # Check collision constraints (exclude zones) + collision_pass = self.constraint_checker.check_collision(points) + collision_pass_count = collision_pass.sum().item() + collision_pass_rate = collision_pass_count / num_points * 100 + + # Overall constraint pass (both bounds and collision) + overall_pass = bounds_pass & collision_pass + overall_pass_count = overall_pass.sum().item() + overall_pass_rate = overall_pass_count / num_points * 100 + + return { + "num_points": num_points, + "bounds_pass_count": bounds_pass_count, + "bounds_pass_rate": bounds_pass_rate, + "collision_pass_count": collision_pass_count, + "collision_pass_rate": collision_pass_rate, + "overall_pass_count": overall_pass_count, + "overall_pass_rate": overall_pass_rate, + "exclude_zones_count": self.constraint_checker.get_num_exclude_zones(), + } + + def _get_backend_priority_list(self) -> List[str]: + """Get the priority-ordered list of visualization backends to try. + + Returns: + List of backend names in order of preference. + """ + backends = [] + + # Prefer sim_manager if available + if self.sim_manager is not None and hasattr(self.sim_manager, "get_env"): + backends.append("sim_manager") + + # Always include open3d and matplotlib as fallbacks + backends.extend(["open3d", "matplotlib"]) + + return backends + + def _create_visualizer_with_config(self, factory, vis_type, backend): + """Create a visualizer with appropriate configuration parameters. + + Args: + factory: VisualizerFactory instance. + vis_type: VisualizationType enum. + backend: Backend string. + + Returns: + Configured visualizer instance. + """ + # Prepare common arguments for all visualizers + common_kwargs = { + "backend": backend, + "sim_manager": self.sim_manager, + "control_part_name": self.control_part_name, + } + + # Add visualization-type specific arguments + if vis_type == VisualizationType.POINT_CLOUD: + common_kwargs["point_size"] = getattr( + self.config.visualization, "point_size", 8.0 + ) + elif vis_type == VisualizationType.VOXEL: + common_kwargs["voxel_size"] = getattr( + self.config.visualization, "voxel_size", 0.05 + ) + elif vis_type == VisualizationType.SPHERE: + common_kwargs["sphere_radius"] = getattr( + self.config.visualization, "sphere_radius", 0.005 + ) + common_kwargs["sphere_resolution"] = getattr( + self.config.visualization, "sphere_resolution", 10 + ) + elif vis_type == VisualizationType.AXIS: + common_kwargs["axis_length"] = getattr( + self.config.visualization, "axis_length", 0.05 + ) + common_kwargs["axis_size"] = getattr( + self.config.visualization, "axis_size", 0.003 + ) + # Pass reference pose if available + if ( + hasattr(self.config, "reference_pose") + and self.config.reference_pose is not None + ): + common_kwargs["reference_pose"] = self.config.reference_pose + # For other visualization types (MESH, HEATMAP), use only common arguments + + return factory.create_visualizer(viz_type=vis_type, **common_kwargs) + + def _generate_point_colors_and_sizes( + self, points: np.ndarray, filtered_to_reachable: bool = False + ) -> Tuple[np.ndarray, np.ndarray]: + """Generate colors and sizes for workspace points based on reachability. + + Args: + points: Workspace points, shape (N, 3). + filtered_to_reachable: Whether points have been pre-filtered to only include reachable ones. + + Returns: + Tuple of: + - Colors array, shape (N, 3) with RGB values in [0, 1]. + - Sizes array, shape (N,) with point sizes. + """ + num_points = len(points) + colors = np.zeros((num_points, 3)) + sizes = ( + np.ones(num_points) * self.config.visualization.point_size + ) # Default size + + # Check if we have current_mode attribute (set during analyze) + if not hasattr(self, "current_mode"): + # Fallback: color based on available reachability information + if ( + hasattr(self, "reachability_mask") + and self.reachability_mask is not None + ): + reachability_mask_np = self.reachability_mask.cpu().numpy() + if len(reachability_mask_np) == num_points: + colors[reachability_mask_np, 1] = 1.0 # Green for reachable + colors[~reachability_mask_np, 0] = 1.0 # Red for unreachable + logger.log_debug("Using available reachability mask for coloring") + else: + colors[:, 1] = 1.0 # Green fallback + logger.log_debug( + "Reachability mask size mismatch, using green fallback" + ) + else: + colors[:, 1] = 1.0 # Green fallback + logger.log_debug( + "No reachability information available, using green fallback" + ) + return colors, sizes + + if self.current_mode == AnalysisMode.JOINT_SPACE: + # Joint space mode: all points are reachable (green, same size) + colors[:, 1] = 1.0 # Green channel = 1.0 + logger.log_debug(f"Coloring {num_points} points as reachable (green)") + + elif self.current_mode in [ + AnalysisMode.CARTESIAN_SPACE, + AnalysisMode.PLANE_SAMPLING, + ]: + # Cartesian/Plane space mode: different colors and sizes based on reachability + mode_name = ( + "Cartesian" + if self.current_mode == AnalysisMode.CARTESIAN_SPACE + else "Plane sampling" + ) + if self.success_rates is not None and hasattr(self, "reachability_mask"): + if filtered_to_reachable: + # Points have been pre-filtered, but we still need to check IK reachability + # Only color as green if we have verified IK solutions + if ( + hasattr(self, "reachability_mask") + and self.reachability_mask is not None + ): + # Use the actual IK reachability results + reachability_mask_np = self.reachability_mask.cpu().numpy() + if len(reachability_mask_np) == num_points: + # Apply the actual reachability coloring + reachable_indices = reachability_mask_np + colors[reachable_indices, 1] = 1.0 # Green for IK-reachable + sizes[reachable_indices] = ( + self.config.visualization.point_size * 1.5 + ) + + unreachable_indices = ~reachability_mask_np + colors[unreachable_indices, 0] = ( + 1.0 # Red for IK-unreachable + ) + sizes[unreachable_indices] = ( + self.config.visualization.point_size * 0.7 + ) + + num_reachable = np.sum(reachable_indices) + num_unreachable = np.sum(unreachable_indices) + logger.log_debug( + f"Coloring {num_reachable} IK-reachable points (green) and " + f"{num_unreachable} IK-unreachable points (red) in {mode_name} mode" + ) + else: + # Fallback: color based on geometric constraint only + colors[:, 1] = 1.0 # All green (geometrically valid) + sizes[:] = self.config.visualization.point_size * 1.5 + logger.log_warning( + f"IK reachability mask length mismatch. " + f"Coloring {num_points} geometrically valid points (green) in {mode_name} mode" + ) + else: + # No IK verification available, assume geometric validity only + colors[:, 1] = 1.0 # All green (geometrically valid) + sizes[:] = self.config.visualization.point_size * 1.5 + logger.log_debug( + f"No IK verification available. " + f"Coloring {num_points} geometrically valid points (green) in {mode_name} mode" + ) + else: + # Original logic for showing both reachable and unreachable points + reachability_mask_np = self.reachability_mask.cpu().numpy() + + # Check if mask length matches points length + if len(reachability_mask_np) != num_points: + logger.log_warning( + f"Reachability mask length ({len(reachability_mask_np)}) doesn't match " + f"points length ({num_points}). Defaulting to all green." + ) + colors[:, 1] = 1.0 # All green as fallback + sizes[:] = self.config.visualization.point_size * 1.5 + else: + # Reachable points: green color, larger size + reachable_indices = reachability_mask_np + colors[reachable_indices, 1] = 1.0 # Pure green + sizes[reachable_indices] = ( + self.config.visualization.point_size * 1.5 + ) # Larger size + + # Unreachable points: red color, smaller size + unreachable_indices = ~reachability_mask_np + colors[unreachable_indices, 0] = 1.0 # Pure red + sizes[unreachable_indices] = ( + self.config.visualization.point_size * 0.7 + ) # Smaller size + + num_reachable = np.sum(reachable_indices) + num_unreachable = np.sum(unreachable_indices) + logger.log_debug( + f"Coloring {num_reachable} reachable points (green, large) and " + f"{num_unreachable} unreachable points (red, small) in {mode_name} mode" + ) + else: + # No success rates available, assume all reachable + colors[:, 1] = 1.0 # Green + logger.log_warning( + f"No success rates available in {mode_name} mode, " + "defaulting to green (reachable)" + ) + + return colors, sizes + + def _generate_point_colors(self, points: np.ndarray) -> np.ndarray: + """Generate colors for workspace points based on reachability (backward compatibility). + + Args: + points: Workspace points, shape (N, 3). + + Returns: + Colors array, shape (N, 3) with RGB values in [0, 1]. + """ + colors, _ = self._generate_point_colors_and_sizes( + points, filtered_to_reachable=False + ) + return colors + + def _compute_metrics(self) -> Dict[str, Any]: + """Compute workspace metrics based on configuration.""" + if self.workspace_points is None or len(self.workspace_points) == 0: + logger.log_warning("No workspace points available for metrics computation") + return {} + + metrics = {} + + # TODO: Implement metric computation using metrics module + # For now, compute basic statistics + points_np = self.workspace_points.cpu().numpy() + + metrics["bounding_box"] = { + "min": points_np.min(axis=0).tolist(), + "max": points_np.max(axis=0).tolist(), + } + + metrics["centroid"] = points_np.mean(axis=0).tolist() + + dimensions = points_np.max(axis=0) - points_np.min(axis=0) + metrics["dimensions"] = dimensions.tolist() + + # Approximate volume (bounding box) + metrics["bounding_box_volume"] = float(np.prod(dimensions)) + + logger.log_info(f"Computed {len(metrics)} metrics") + + return metrics + + def visualize( + self, + vis_type: VisualizationType | str | None = None, + show: bool = True, + save_path: str | None = None, + backend: str | None = None, + ) -> Any: + """Visualize the workspace. + + Args: + vis_type: Type of visualization to create. Can be VisualizationType enum or string. + If None, uses the vis_type from configuration (default: POINT_CLOUD). + Supported types: 'point_cloud', 'voxel', 'sphere'. + show: Whether to display the visualization. + save_path: Optional path to save the visualization. + backend: Backend to use ('sim_manager', 'open3d', 'matplotlib', 'data'). + If None, automatically selects based on availability. + + Returns: + Visualization object. + """ + if self.workspace_points is None or len(self.workspace_points) == 0: + logger.log_error("No workspace points available for visualization") + return None + + if not self.config.visualization.enabled: + logger.log_warning("Visualization is disabled in configuration") + return None + + # Use configured vis_type if not specified + if vis_type is None: + vis_type = self.config.visualization.vis_type + + # Handle string vis_type by converting to enum + if isinstance(vis_type, str): + try: + vis_type = VisualizationType(vis_type) + except ValueError: + logger.log_warning( + f"Unknown visualization type '{vis_type}', falling back to POINT_CLOUD" + ) + vis_type = VisualizationType.POINT_CLOUD + + # Convert points to numpy first + points_np = self.workspace_points.cpu().numpy() + filtered_points = False # Track if points were filtered + + # Enhanced visualization logging with point count info + vis_start_time = time.time() + logger.log_info( + f"Creating {vis_type.value} visualization for {len(points_np)} points..." + ) + + # Auto-select backend if not specified + if backend is None: + if self.sim_manager is not None and hasattr(self.sim_manager, "sim"): + backend = "sim_manager" + else: + backend = "open3d" + + # Filter points if configured to hide unreachable ones in Cartesian/Plane space mode + if ( + self.current_mode + in [AnalysisMode.CARTESIAN_SPACE, AnalysisMode.PLANE_SAMPLING] + and not self.config.visualization.show_unreachable_points + and hasattr(self, "reachability_mask") + ): + # Only show reachable points + reachable_mask = self.reachability_mask.cpu().numpy() + + # Check if mask length matches points length before filtering + if len(reachable_mask) != len(points_np): + logger.log_warning( + f"Cannot filter points: reachability mask length ({len(reachable_mask)}) " + f"doesn't match points length ({len(points_np)}). Showing all points." + ) + else: + points_np = points_np[reachable_mask] + filtered_points = True + logger.log_info( + f"Filtering to show only {len(points_np)} reachable points" + ) + + # Generate colors and sizes based on reachability + colors, sizes = self._generate_point_colors_and_sizes( + points_np, filtered_points + ) + + # Create visualizer using factory pattern + from embodichain.lab.sim.utility.workspace_analyzer.visualizers import ( + VisualizerFactory, + ) + + factory = VisualizerFactory() + visualizer = self._create_visualizer_with_config(factory, vis_type, backend) + + # Create visualization with sizes if supported + try: + # Try to pass sizes to visualizer (some backends may support it) + vis_obj = visualizer.visualize(points_np, colors=colors, sizes=sizes) + except TypeError: + # Fallback to colors-only visualization if sizes not supported + vis_obj = visualizer.visualize(points_np, colors=colors) + + # Performance tracking for visualization + vis_time = time.time() - vis_start_time + logger.log_info(f"✨ Visualization created in {vis_time:.2f}s") + + # Save if requested + if save_path: + save_start = time.time() + visualizer.save(save_path) + save_time = time.time() - save_start + logger.log_info(f"💾 Saved visualization to {save_path} ({save_time:.2f}s)") + + # Show if requested (only for non-sim_manager backends) + if show and backend != "sim_manager": + try: + visualizer.show() + except Exception as e: + logger.log_warning(f"Failed to show visualization: {e}") + + return vis_obj + + def _load_from_cache(self) -> dict[str, Any] | None: + """Load analysis results from cache.""" + if self.cache is None: + return None + + # TODO: Implement cache loading logic + return None + + def _save_to_cache(self, results: Dict[str, Any]) -> None: + """Save analysis results to cache.""" + if self.cache is None: + return + + # TODO: Implement cache saving logic + pass + + def get_workspace_bounds(self) -> Dict[str, np.ndarray]: + """Get the bounding box of the analyzed workspace. + + Returns: + Dictionary with 'min' and 'max' bounds. + """ + if self.workspace_points is None or len(self.workspace_points) == 0: + logger.log_warning("No workspace points available") + return {"min": None, "max": None} + + points_np = self.workspace_points.cpu().numpy() + return {"min": points_np.min(axis=0), "max": points_np.max(axis=0)} + + def export_results(self, output_path: str, format: str = "npz") -> None: + """Export analysis results to file. + + Args: + output_path: Path to save the results. + format: Output format ('npz', 'pkl', 'json'). + """ + if self.workspace_points is None: + logger.log_error("No analysis results to export") + return + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + if format == "npz": + np.savez( + output_path, + workspace_points=self.workspace_points.cpu().numpy(), + joint_configurations=self.joint_configurations.cpu().numpy(), + metrics=self.metrics_results, + ) + elif format == "pkl": + import pickle + + with open(output_path, "wb") as f: + pickle.dump( + { + "workspace_points": self.workspace_points.cpu().numpy(), + "joint_configurations": self.joint_configurations.cpu().numpy(), + "metrics": self.metrics_results, + }, + f, + ) + elif format == "json": + import json + + # Convert tensors to lists for JSON serialization + with open(output_path, "w") as f: + json.dump( + { + "workspace_points": self.workspace_points.cpu() + .numpy() + .tolist(), + "metrics": self.metrics_results, + }, + f, + indent=2, + ) + else: + logger.log_error(f"Unsupported format: {format}") + return + + # File size information for export + try: + file_size = output_path.stat().st_size + if file_size > 1024 * 1024: # > 1MB + size_str = f"{file_size / (1024*1024):.1f} MB" + elif file_size > 1024: # > 1KB + size_str = f"{file_size / 1024:.1f} KB" + else: + size_str = f"{file_size} bytes" + logger.log_info(f"💾 Exported results to {output_path} ({size_str})") + except OSError: + # File size unavailable + logger.log_info(f"💾 Exported results to {output_path}") + + @contextmanager + def profiling(self): + """Enhanced context manager for profiling workspace analysis with detailed metrics.""" + logger.log_info("🔍 Starting profiled analysis...") + start_time = time.time() + start_mem = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0 + + # CPU memory tracking (if psutil is available) + start_cpu_mem = 0 + process = None + if psutil is not None: + try: + process = psutil.Process() + start_cpu_mem = process.memory_info().rss + except (psutil.Error, OSError): + # Process info unavailable + process = None + + yield + + end_time = time.time() + end_mem = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0 + end_cpu_mem = 0 + if process is not None: + try: + end_cpu_mem = process.memory_info().rss + except (psutil.Error, OSError): + # Process info unavailable, use start value + end_cpu_mem = start_cpu_mem + + # Detailed performance summary + analysis_time = end_time - start_time + if analysis_time < 60: + time_str = f"{analysis_time:.2f}s" + else: + minutes = int(analysis_time // 60) + seconds = analysis_time % 60 + time_str = f"{minutes}m {seconds:.1f}s" + + logger.log_info(f"⏱️ Analysis time: {time_str}") + + # Memory usage summary + if torch.cuda.is_available(): + gpu_mem_used = (end_mem - start_mem) / 1024**2 + logger.log_info(f"💾 GPU memory used: {gpu_mem_used:.2f} MB") + + # CPU memory tracking (if available) + if process is not None and end_cpu_mem > start_cpu_mem: + cpu_mem_used = (end_cpu_mem - start_cpu_mem) / 1024**2 + logger.log_info(f"💻 CPU memory used: {cpu_mem_used:.2f} MB") + + # Performance rating + if analysis_time < 30 and (not torch.cuda.is_available() or gpu_mem_used < 100): + logger.log_info("🚀 Performance: Excellent!") + elif analysis_time < 120: + logger.log_info("✅ Performance: Good") + elif analysis_time < 300: + logger.log_info("🟡 Performance: Moderate") + else: + logger.log_info("🐌 Performance: Needs optimization") diff --git a/embodichain/toolkits/urdf_assembly/sensor.py b/embodichain/toolkits/urdf_assembly/sensor.py index 87fcd5e..7e072bd 100644 --- a/embodichain/toolkits/urdf_assembly/sensor.py +++ b/embodichain/toolkits/urdf_assembly/sensor.py @@ -103,10 +103,10 @@ def attach_sensor( sensor_source: Sensor definition source (multiple formats supported). parent_component (str): Target component name for sensor attachment. parent_link (str): Specific link within parent component for attachment. - transform (Optional[np.ndarray]): Optional 4x4 homogeneous transformation matrix. - sensor_type (Optional[str]): Sensor type classification. - extract_links (Optional[List[str]]): Specific link names to extract from URDF. - extract_joints (Optional[List[str]]): Specific joint names to extract from URDF. + transform (np.ndarray | None): Optional 4x4 homogeneous transformation matrix. + sensor_type (str | None): Sensor type classification. + extract_links (list[str] | None): Specific link names to extract from URDF. + extract_joints (list[str] | None): Specific joint names to extract from URDF. Returns: bool: True if sensor attachment successful, False on failure. diff --git a/embodichain/utils/math.py b/embodichain/utils/math.py index 238767a..c2ac591 100644 --- a/embodichain/utils/math.py +++ b/embodichain/utils/math.py @@ -1570,6 +1570,7 @@ def sample_uniform( lower: Union[torch.Tensor, float], upper: Union[torch.Tensor, float], size: Union[int, tuple[int, ...]], + device: torch.device = torch.device("cpu"), ) -> torch.Tensor: """Sample uniformly within a range. @@ -1586,7 +1587,7 @@ def sample_uniform( if isinstance(size, int): size = (size,) # return tensor - return torch.rand(*size, device=lower.device) * (upper - lower) + lower + return torch.rand(*size, device=device) * (upper - lower) + lower def sample_log_uniform( diff --git a/examples/sim/demo/grasp_cup_to_caffe.py b/examples/sim/demo/grasp_cup_to_caffe.py index 77f31fb..d7ca271 100644 --- a/examples/sim/demo/grasp_cup_to_caffe.py +++ b/examples/sim/demo/grasp_cup_to_caffe.py @@ -83,11 +83,11 @@ def initialize_simulation(args) -> SimulationManager: sim_device=args.device, enable_rt=args.enable_rt, physics_dt=1.0 / 100.0, + num_envs=args.num_envs, + arena_space=2.5, ) sim = SimulationManager(config) - sim.build_multiple_arenas(args.num_envs, space=2.5) - if args.enable_rt: light = sim.add_light( cfg=LightCfg( diff --git a/examples/sim/gizmo/gizmo_camera.py b/examples/sim/gizmo/gizmo_camera.py index 2b782a0..382c4f2 100644 --- a/examples/sim/gizmo/gizmo_camera.py +++ b/examples/sim/gizmo/gizmo_camera.py @@ -40,9 +40,6 @@ def main(): parser = argparse.ArgumentParser( description="Create and simulate a camera with gizmo in SimulationManager" ) - parser.add_argument( - "--num_envs", type=int, default=1, help="Number of environments to simulate" - ) parser.add_argument( "--device", type=str, @@ -72,10 +69,6 @@ def main(): sim = SimulationManager(sim_cfg) sim.set_manual_update(False) - # Build multiple arenas if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=3.0) - # Add some objects to the scene for camera to observe for i in range(5): cube_cfg = RigidObjectCfg( diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index 6b664ef..1860ea1 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -43,9 +43,6 @@ def main(): default=False, help="Run simulation in headless mode", ) - parser.add_argument( - "--num_envs", type=int, default=1, help="Number of parallel environments" - ) parser.add_argument( "--device", type=str, default="cpu", help="Simulation device (cuda or cpu)" ) @@ -70,10 +67,6 @@ def main(): # Create the simulation instance sim = SimulationManager(sim_cfg) - # Build multiple arenas if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=3.0) - # Add two cubes to the scene cube1: RigidObject = sim.add_rigid_object( cfg=RigidObjectCfg( diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index fae7996..81fc4b2 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -41,9 +41,6 @@ def main(): parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - parser.add_argument( - "--num_envs", type=int, default=1, help="Number of parallel environments" - ) parser.add_argument( "--device", type=str, default="cpu", help="Simulation device (cuda or cpu)" ) @@ -67,10 +64,6 @@ def main(): sim = SimulationManager(sim_cfg) sim.set_manual_update(False) - # Build multiple arenas if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=3.0) - # Get UR10 URDF path urdf_path = get_data_path("UniversalRobots/UR10/UR10.urdf") diff --git a/examples/sim/gizmo/gizmo_scene.py b/examples/sim/gizmo/gizmo_scene.py index 31c7d23..dd3540d 100644 --- a/examples/sim/gizmo/gizmo_scene.py +++ b/examples/sim/gizmo/gizmo_scene.py @@ -49,9 +49,6 @@ def main(): parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - parser.add_argument( - "--num_envs", type=int, default=1, help="Number of parallel environments" - ) parser.add_argument( "--device", type=str, default="cpu", help="Simulation device (cuda or cpu)" ) @@ -75,10 +72,6 @@ def main(): sim = SimulationManager(sim_cfg) sim.set_manual_update(False) - # Build multiple arenas if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=3.0) - # Get DexForce W1 URDF path urdf_path = get_data_path( "DexforceW1V021_INDUSTRIAL_DH_PGC_GRIPPER_M/DexforceW1V021.urdf" diff --git a/examples/sim/gizmo/gizmo_w1.py b/examples/sim/gizmo/gizmo_w1.py index 90d73ed..2cf84a9 100644 --- a/examples/sim/gizmo/gizmo_w1.py +++ b/examples/sim/gizmo/gizmo_w1.py @@ -41,9 +41,6 @@ def main(): parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) - parser.add_argument( - "--num_envs", type=int, default=1, help="Number of parallel environments" - ) parser.add_argument( "--device", type=str, default="cpu", help="Simulation device (cuda or cpu)" ) @@ -67,10 +64,6 @@ def main(): sim = SimulationManager(sim_cfg) sim.set_manual_update(False) - # Build multiple arenas if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=3.0) - # Get DexForce W1 URDF path urdf_path = get_data_path( "DexforceW1V021_INDUSTRIAL_DH_PGC_GRIPPER_M/DexforceW1V021.urdf" diff --git a/examples/sim/planners/motion_generator.py b/examples/sim/planners/motion_generator.py index d767aba..9e7b7c9 100644 --- a/examples/sim/planners/motion_generator.py +++ b/examples/sim/planners/motion_generator.py @@ -78,7 +78,6 @@ def main(interactive=False): # Initialize simulation sim = SimulationManager(SimulationManagerCfg(headless=False, sim_device="cpu")) - sim.build_multiple_arenas(1) sim.set_manual_update(False) # Robot configuration diff --git a/examples/sim/robot/dexforce_w1.py b/examples/sim/robot/dexforce_w1.py new file mode 100644 index 0000000..d36af47 --- /dev/null +++ b/examples/sim/robot/dexforce_w1.py @@ -0,0 +1,75 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +""" +This script shows how to customize the end-effectors of the DexForce W1 robot by +adding a pair of parallel grippers as the left and right hands. +""" + +import numpy as np + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.robots import DexforceW1Cfg + + +def main(): + np.set_printoptions(precision=5, suppress=True) + + config = SimulationManagerCfg() + sim = SimulationManager(config) + + cfg = DexforceW1Cfg.from_dict( + { + "uid": "dexforce_w1", + "version": "v021", + "arm_kind": "anthropomorphic", + "with_default_eef": False, + "control_parts": { + "left_eef": ["LEFT_FINGER1_JOINT", "LEFT_FINGER2_JOINT"], + "right_eef": ["RIGHT_FINGER1_JOINT", "RIGHT_FINGER2_JOINT"], + }, + "urdf_cfg": { + "components": [ + { + "component_type": "left_hand", + "urdf_path": "DH_PGC_140_50/DH_PGC_140_50.urdf", + }, + { + "component_type": "right_hand", + "urdf_path": "DH_PGC_140_50/DH_PGC_140_50.urdf", + }, + ] + }, + "drive_pros": { + "max_effort": { + "left_eef": 10.0, + "right_eef": 10.0, + } + }, + } + ) + + robot = sim.add_robot(cfg=cfg) + sim.update(step=1) + print("DexforceW1 with a user defined end-effector added to the simulation.") + + from IPython import embed + + embed() + + +if __name__ == "__main__": + main() diff --git a/examples/sim/scene/scene_demo.py b/examples/sim/scene/scene_demo.py index 20ff70a..ffab097 100644 --- a/examples/sim/scene/scene_demo.py +++ b/examples/sim/scene/scene_demo.py @@ -122,12 +122,11 @@ def main(): physics_dt=1.0 / 100.0, sim_device=args.device, enable_rt=not args.disable_rt, + num_envs=args.num_envs, + arena_space=10.0, ) sim = SimulationManager(sim_cfg) - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=10.0) - num_lights = 8 radius = 5 height = 8 diff --git a/examples/sim/sensors/batch_camera.py b/examples/sim/sensors/batch_camera.py index efbd17f..b4cd437 100644 --- a/examples/sim/sensors/batch_camera.py +++ b/examples/sim/sensors/batch_camera.py @@ -33,10 +33,13 @@ def main(args): config = SimulationManagerCfg( - headless=True, sim_device=args.device, arena_space=2, enable_rt=args.enable_rt + headless=True, + sim_device=args.device, + num_envs=args.num_envs, + arena_space=2, + enable_rt=args.enable_rt, ) sim = SimulationManager(config) - sim.build_multiple_arenas(args.num_envs) rigid_obj: RigidObject = sim.add_rigid_object( cfg=RigidObjectCfg( diff --git a/examples/sim/solvers/differential_solver.py b/examples/sim/solvers/differential_solver.py index 1f05657..27f00e7 100644 --- a/examples/sim/solvers/differential_solver.py +++ b/examples/sim/solvers/differential_solver.py @@ -34,10 +34,9 @@ def main(visualize: bool = True): sim_device = "cpu" num_envs = 9 # Number of parallel arenas/environments config = SimulationManagerCfg( - headless=False, sim_device=sim_device, arena_space=1.5 + headless=False, sim_device=sim_device, arena_space=1.5, num_envs=num_envs ) sim = SimulationManager(config) - sim.build_multiple_arenas(num_envs) sim.set_manual_update(False) # Load robot URDF file diff --git a/examples/sim/solvers/opw_solver.py b/examples/sim/solvers/opw_solver.py index aa47eae..e19c23f 100644 --- a/examples/sim/solvers/opw_solver.py +++ b/examples/sim/solvers/opw_solver.py @@ -34,7 +34,6 @@ def main(): sim_device = "cpu" config = SimulationManagerCfg(headless=False, sim_device=sim_device) sim = SimulationManager(config) - sim.build_multiple_arenas(1) sim.set_manual_update(False) # Robot configuration dictionary diff --git a/examples/sim/solvers/pink_solver.py b/examples/sim/solvers/pink_solver.py index aea419c..e63e326 100644 --- a/examples/sim/solvers/pink_solver.py +++ b/examples/sim/solvers/pink_solver.py @@ -34,7 +34,6 @@ def main(): sim_device = "cpu" config = SimulationManagerCfg(headless=False, sim_device=sim_device) sim = SimulationManager(config) - sim.build_multiple_arenas(1) sim.set_manual_update(False) # Load robot URDF file diff --git a/examples/sim/solvers/pinocchio_solver.py b/examples/sim/solvers/pinocchio_solver.py index bfd3730..af31ade 100644 --- a/examples/sim/solvers/pinocchio_solver.py +++ b/examples/sim/solvers/pinocchio_solver.py @@ -35,7 +35,6 @@ def main(): sim_device = "cpu" config = SimulationManagerCfg(headless=False, sim_device=sim_device) sim = SimulationManager(config) - sim.build_multiple_arenas(1) sim.set_manual_update(False) # Load robot URDF file diff --git a/examples/sim/solvers/pytorch_solver.py b/examples/sim/solvers/pytorch_solver.py index 822eac5..d587dc3 100644 --- a/examples/sim/solvers/pytorch_solver.py +++ b/examples/sim/solvers/pytorch_solver.py @@ -20,10 +20,9 @@ def main(): sim_device = "cpu" num_envs = 9 # Number of parallel environments config = SimulationManagerCfg( - headless=False, sim_device=sim_device, arena_space=2.0 + headless=False, sim_device=sim_device, arena_space=2.0, num_envs=num_envs ) sim = SimulationManager(config) - sim.build_multiple_arenas(num_envs) sim.set_manual_update(False) # Load robot URDF file diff --git a/examples/sim/solvers/srs_solver.py b/examples/sim/solvers/srs_solver.py index 5d9922b..314040e 100644 --- a/examples/sim/solvers/srs_solver.py +++ b/examples/sim/solvers/srs_solver.py @@ -38,7 +38,6 @@ def main(): ) ) - sim.build_multiple_arenas(1) sim.set_manual_update(False) robot: Robot = sim.add_robot(cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"})) diff --git a/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py b/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py new file mode 100644 index 0000000..f3b607b --- /dev/null +++ b/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py @@ -0,0 +1,98 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np +from IPython import embed + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.robots import DexforceW1Cfg +from embodichain.lab.sim.cfg import MarkerCfg +from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, + AnalysisMode, +) +from embodichain.lab.sim.utility.workspace_analyzer.configs.visualization_config import ( + VisualizationConfig, +) + +if __name__ == "__main__": + # Example usage + np.set_printoptions(precision=5, suppress=True) + torch.set_printoptions(precision=5, sci_mode=False) + + config = SimulationManagerCfg( + headless=False, sim_device="cpu", width=1080, height=1080 + ) + sim = SimulationManager(config) + sim.set_manual_update(False) + + cfg = DexforceW1Cfg.from_dict( + {"uid": "dexforce_w1", "version": "v021", "arm_kind": "industrial"} + ) + robot = sim.add_robot(cfg=cfg) + print("DexforceW1 robot added to the simulation.") + + # Set left arm joint positions (mirrored) + left_qpos = torch.tensor([0, -np.pi / 4, 0.0, -np.pi / 2, -np.pi / 4, 0.0, 0.0]) + right_qpos = -left_qpos + robot.set_qpos( + qpos=left_qpos, + joint_ids=robot.get_joint_ids("left_arm"), + ) + # Set right arm joint positions (mirrored) + robot.set_qpos( + qpos=right_qpos, + joint_ids=robot.get_joint_ids("right_arm"), + ) + + left_arm_pose = robot.compute_fk( + qpos=left_qpos, + name="left_arm", + to_matrix=True, + ) + + sim.draw_marker( + cfg=MarkerCfg( + name=f"left_arm_pose_axis", + marker_type="axis", + axis_xpos=left_arm_pose, + axis_size=0.005, + axis_len=0.15, + arena_index=0, + ) + ) + + cartesian_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.CARTESIAN_SPACE, + visualization=VisualizationConfig( + show_unreachable_points=False, point_size=8.0 + ), + control_part_name="left_arm", + ) + wa_cartesian = WorkspaceAnalyzer( + robot=robot, config=cartesian_config, sim_manager=sim + ) + results_cartesian = wa_cartesian.analyze(num_samples=1000, visualize=True) + print(f"\nCartesian Space Results:") + print( + f" Reachable points: {results_cartesian['num_reachable']} / {results_cartesian['num_samples']}" + ) + print(f" Analysis time: {results_cartesian['analysis_time']:.2f}s") + print(f" Metrics: {results_cartesian['metrics']}") + + embed(header="Workspace Analyzer Test Environment") diff --git a/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py b/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py new file mode 100644 index 0000000..d20db50 --- /dev/null +++ b/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py @@ -0,0 +1,55 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np +from IPython import embed + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.robots import DexforceW1Cfg + +from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + WorkspaceAnalyzer, +) + +if __name__ == "__main__": + # Example usage + np.set_printoptions(precision=5, suppress=True) + torch.set_printoptions(precision=5, sci_mode=False) + + config = SimulationManagerCfg(headless=False, sim_device="cpu") + sim_manager = SimulationManager(config) + sim_manager.set_manual_update(False) + + cfg = DexforceW1Cfg.from_dict( + {"uid": "dexforce_w1", "version": "v021", "arm_kind": "industrial"} + ) + robot = sim_manager.add_robot(cfg=cfg) + print("DexforceW1 robot added to the simulation.") + + print("Example: Joint Space Analysis") + + wa_joint = WorkspaceAnalyzer(robot=robot, sim_manager=sim_manager) + results_joint = wa_joint.analyze(num_samples=3000, visualize=True) + + print(f"\nJoint Space Results:") + print( + f" Valid points: {results_joint['num_valid']} / {results_joint['num_samples']}" + ) + print(f" Analysis time: {results_joint['analysis_time']:.2f}s") + print(f" Metrics: {results_joint['metrics']}") + + embed(header="End of Joint Space Analysis Example") diff --git a/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py b/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py new file mode 100644 index 0000000..b36c289 --- /dev/null +++ b/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py @@ -0,0 +1,105 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np +from IPython import embed + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.robots import DexforceW1Cfg +from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, + AnalysisMode, +) +from embodichain.lab.sim.cfg import MarkerCfg +from embodichain.lab.sim.utility.workspace_analyzer.configs.visualization_config import ( + VisualizationConfig, +) + + +if __name__ == "__main__": + # Example usage + np.set_printoptions(precision=5, suppress=True) + torch.set_printoptions(precision=5, sci_mode=False) + + config = SimulationManagerCfg( + headless=False, sim_device="cpu", width=1080, height=1080 + ) + sim = SimulationManager(config) + sim.set_manual_update(False) + + cfg = DexforceW1Cfg.from_dict( + {"uid": "dexforce_w1", "version": "v021", "arm_kind": "industrial"} + ) + robot = sim.add_robot(cfg=cfg) + print("DexforceW1 robot added to the simulation.") + + # Set left arm joint positions (mirrored) + left_qpos = torch.tensor([0, -np.pi / 4, 0.0, -np.pi / 2, -np.pi / 4, 0.0, 0.0]) + right_qpos = -left_qpos + robot.set_qpos( + qpos=left_qpos, + joint_ids=robot.get_joint_ids("left_arm"), + ) + # Set right arm joint positions (mirrored) + robot.set_qpos( + qpos=right_qpos, + joint_ids=robot.get_joint_ids("right_arm"), + ) + + left_arm_pose = robot.compute_fk( + qpos=left_qpos, + name="left_arm", + to_matrix=True, + ) + + sim.draw_marker( + cfg=MarkerCfg( + name=f"left_arm_pose_axis", + marker_type="axis", + axis_xpos=left_arm_pose, + axis_size=0.005, + axis_len=0.15, + arena_index=0, + ) + ) + + cartesian_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.PLANE_SAMPLING, + plane_normal=torch.tensor([0.0, 0.0, 1.0]), + plane_point=torch.tensor([0.0, 0.0, 1.2]), + # plane_bounds=torch.tensor([[-0.5, 0.5], [-0.5, 0.5]]), + reference_pose=left_arm_pose[0] + .cpu() + .numpy(), # Use computed left arm pose as reference + visualization=VisualizationConfig( + show_unreachable_points=False, vis_type="axis" + ), + control_part_name="left_arm", + ) + wa_cartesian = WorkspaceAnalyzer( + robot=robot, config=cartesian_config, sim_manager=sim + ) + results_cartesian = wa_cartesian.analyze(num_samples=1500, visualize=True) + print(f"\nCartesian Space Results:") + print( + f" Reachable points: {results_cartesian['num_reachable']} / {results_cartesian['num_samples']}" + ) + print(f" Analysis time: {results_cartesian['analysis_time']:.2f}s") + print(f" Metrics: {results_cartesian['metrics']}") + + embed(header="Workspace Analyzer Test Environment") diff --git a/pyproject.toml b/pyproject.toml index 07fbd83..043ee60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dynamic = ["version"] # Core install dependencies (kept from requirements.txt). Some VCS links are # specified using PEP 508 direct references where present. dependencies = [ - "dexsim_engine", + "dexsim_engine==0.3.8", "setuptools>=78.1.1", "gymnasium==0.29.1", "casadi==3.7.1", diff --git a/scripts/data_gen.sh b/scripts/data_gen.sh index 598afb8..fa3f6eb 100755 --- a/scripts/data_gen.sh +++ b/scripts/data_gen.sh @@ -1,13 +1,26 @@ #!/bin/bash -NUM_PROCESSES=3 # Set this to the number of parallel processes you want - -for ((i=0; i None: ), ) + def _setup_sensors(self, **kwargs): + """Setup the sensors in the environment. + + Returns: + Dict[str, BaseSensor]: A dictionary mapping sensor UIDs to sensor instances. + """ + + # TODO: support sensor attachment to the robot. + + from embodichain.lab.sim.sensors import CameraCfg, Camera + + sensors = {} + camera: Camera = self.sim.add_sensor( + sensor_cfg=CameraCfg( + width=640, + height=480, + intrinsics=[400, 400, 320, 240], + extrinsics=CameraCfg.ExtrinsicsCfg( + eye=[0, 0, 3.0], + target=[0, 0, 0], + ), + near=0.01, + far=10.0, + ) + ) + sensors[camera.uid] = camera + return sensors + def _update_sim_state(self, **kwargs) -> None: pose = torch.eye(4, device=self.device) pose = pose.unsqueeze_(0).repeat(self.num_envs, 1, 1) @@ -118,6 +151,9 @@ def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs: parser.add_argument( "--num_envs", type=int, default=1, help="number of environments to run" ) + parser.add_argument( + "--gpu_id", type=int, default=0, help="GPU ID to run the environment on" + ) parser.add_argument( "--device", type=str, @@ -132,11 +168,12 @@ def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs: num_envs=args.num_envs, headless=args.headless, device=args.device, + gpu_id=args.gpu_id, ) - for episode in range(10): + for episode in range(1000): print("Episode:", episode) - env.reset() + obs, info = env.reset() start_time = time.time() total_steps = 0 diff --git a/scripts/tutorials/sim/create_rigid_object_group.py b/scripts/tutorials/sim/create_rigid_object_group.py index 223c72d..c385bdc 100644 --- a/scripts/tutorials/sim/create_rigid_object_group.py +++ b/scripts/tutorials/sim/create_rigid_object_group.py @@ -66,15 +66,13 @@ def main(): physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) sim_device=args.device, enable_rt=args.enable_rt, # Enable ray tracing for better visuals + num_envs=args.num_envs, + arena_space=3.0, ) # Create the simulation instance sim = SimulationManager(sim_cfg) - # Build multiple arenas if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=3.0) - physics_attrs = RigidBodyAttributesCfg( mass=1.0, dynamic_friction=0.5, diff --git a/scripts/tutorials/sim/create_robot.py b/scripts/tutorials/sim/create_robot.py index 29a4606..7ae2c28 100644 --- a/scripts/tutorials/sim/create_robot.py +++ b/scripts/tutorials/sim/create_robot.py @@ -69,13 +69,10 @@ def main(): arena_space=3.0, enable_rt=args.enable_rt, physics_dt=1.0 / 100.0, + num_envs=args.num_envs, ) sim = SimulationManager(config) - # Build multiple environments if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs) - # Create robot configuration robot = create_robot(sim) diff --git a/scripts/tutorials/sim/create_scene.py b/scripts/tutorials/sim/create_scene.py index be353d2..d89b92d 100644 --- a/scripts/tutorials/sim/create_scene.py +++ b/scripts/tutorials/sim/create_scene.py @@ -64,15 +64,13 @@ def main(): physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) sim_device=args.device, enable_rt=args.enable_rt, # Enable ray tracing for better visuals + num_envs=args.num_envs, + arena_space=3.0, ) # Create the simulation instance sim = SimulationManager(sim_cfg) - # Build multiple arenas if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=3.0) - # Add objects to the scene cube: RigidObject = sim.add_rigid_object( cfg=RigidObjectCfg( diff --git a/scripts/tutorials/sim/create_sensor.py b/scripts/tutorials/sim/create_sensor.py index d35e2cd..52db7b8 100644 --- a/scripts/tutorials/sim/create_sensor.py +++ b/scripts/tutorials/sim/create_sensor.py @@ -103,13 +103,10 @@ def main(): arena_space=3.0, enable_rt=args.enable_rt, physics_dt=1.0 / 100.0, + num_envs=args.num_envs, ) sim = SimulationManager(config) - # Build multiple environments if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs) - # Create robot configuration robot = create_robot(sim) diff --git a/scripts/tutorials/sim/create_softbody.py b/scripts/tutorials/sim/create_softbody.py index 55a368d..1aa732d 100644 --- a/scripts/tutorials/sim/create_softbody.py +++ b/scripts/tutorials/sim/create_softbody.py @@ -71,9 +71,6 @@ def main(): # Create the simulation instance sim = SimulationManager(sim_cfg) - # Build multiple arenas if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=3.0) print("[INFO]: Scene setup complete!") # add softbody to the scene diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index fae7996..74ec5d9 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -67,10 +67,6 @@ def main(): sim = SimulationManager(sim_cfg) sim.set_manual_update(False) - # Build multiple arenas if requested - if args.num_envs > 1: - sim.build_multiple_arenas(args.num_envs, space=3.0) - # Get UR10 URDF path urdf_path = get_data_path("UniversalRobots/UR10/UR10.urdf") diff --git a/scripts/tutorials/sim/motion_generator.py b/scripts/tutorials/sim/motion_generator.py index c09a1ed..ca2fca2 100644 --- a/scripts/tutorials/sim/motion_generator.py +++ b/scripts/tutorials/sim/motion_generator.py @@ -78,7 +78,6 @@ def main(): # Initialize simulation sim = SimulationManager(SimulationManagerCfg(headless=False, sim_device="cpu")) - sim.build_multiple_arenas(1) sim.set_manual_update(False) # Robot configuration diff --git a/scripts/tutorials/sim/srs_solver.py b/scripts/tutorials/sim/srs_solver.py index 5d9922b..314040e 100644 --- a/scripts/tutorials/sim/srs_solver.py +++ b/scripts/tutorials/sim/srs_solver.py @@ -38,7 +38,6 @@ def main(): ) ) - sim.build_multiple_arenas(1) sim.set_manual_update(False) robot: Robot = sim.add_robot(cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"})) diff --git a/tests/gym/envs/test_base_env.py b/tests/gym/envs/test_base_env.py index 4857b1b..e9d1f3a 100644 --- a/tests/gym/envs/test_base_env.py +++ b/tests/gym/envs/test_base_env.py @@ -45,7 +45,6 @@ class RandomReachEnv(BaseEnv): def __init__( self, - num_envs=1, drive_type="force", headless=False, device="cpu", @@ -57,7 +56,7 @@ def __init__( sim_cfg=SimulationManagerCfg( headless=headless, arena_space=2.0, sim_device=device ), - num_envs=num_envs, + num_envs=NUM_ENVS, ) super().__init__( @@ -184,3 +183,11 @@ def setup_method(self): class TestBaseEnvCUDA(BaseEnvTest): def setup_method(self): self.setup_simulation("cuda") + + +if __name__ == "__main__": + # Execute the tests when the script is run directly + test_cpu = TestBaseEnvCPU() + test_cpu.setup_method() + test_cpu.test_env_rollout() + test_cpu.teardown_method() diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index fe93e63..072987f 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -15,7 +15,6 @@ # ---------------------------------------------------------------------------- import os -from numpy import mat import torch import pytest @@ -37,9 +36,10 @@ class BaseArticulationTest: """Shared test logic for CPU and CUDA.""" def setup_simulation(self, sim_device): - config = SimulationManagerCfg(headless=True, sim_device=sim_device) + config = SimulationManagerCfg( + headless=True, sim_device=sim_device, num_envs=NUM_ARENAS + ) self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(NUM_ARENAS) art_path = get_data_path(ART_PATH) assert os.path.isfile(art_path) diff --git a/tests/sim/objects/test_light.py b/tests/sim/objects/test_light.py index fe6ef55..ec72664 100644 --- a/tests/sim/objects/test_light.py +++ b/tests/sim/objects/test_light.py @@ -23,9 +23,8 @@ class TestLight: def setup_method(self): # Setup SimulationManager - config = SimulationManagerCfg(headless=True, sim_device="cpu") + config = SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=10) self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(10) # Create batch of lights cfg_dict = { diff --git a/tests/sim/objects/test_rigid_object.py b/tests/sim/objects/test_rigid_object.py index 86ed73d..d43a3ad 100644 --- a/tests/sim/objects/test_rigid_object.py +++ b/tests/sim/objects/test_rigid_object.py @@ -40,9 +40,10 @@ class BaseRigidObjectTest: """Shared test logic for CPU and CUDA.""" def setup_simulation(self, sim_device): - config = SimulationManagerCfg(headless=True, sim_device=sim_device) + config = SimulationManagerCfg( + headless=True, sim_device=sim_device, num_envs=NUM_ARENAS + ) self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(NUM_ARENAS) duck_path = get_data_path(DUCK_PATH) assert os.path.isfile(duck_path) @@ -57,6 +58,9 @@ def setup_simulation(self, sim_device): "shape_type": "Mesh", "fpath": duck_path, }, + "attrs": { + "mass": 1.0, + }, "body_type": "dynamic", } self.duck: RigidObject = self.sim.add_rigid_object( diff --git a/tests/sim/objects/test_rigid_object_group.py b/tests/sim/objects/test_rigid_object_group.py index e9a818e..36f6092 100644 --- a/tests/sim/objects/test_rigid_object_group.py +++ b/tests/sim/objects/test_rigid_object_group.py @@ -35,9 +35,10 @@ class BaseRigidObjectGroupTest: """Shared test logic for CPU and CUDA.""" def setup_simulation(self, sim_device): - config = SimulationManagerCfg(headless=True, sim_device=sim_device) + config = SimulationManagerCfg( + headless=True, sim_device=sim_device, num_envs=NUM_ARENAS + ) self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(NUM_ARENAS) duck_path = get_data_path(DUCK_PATH) assert os.path.isfile(duck_path) diff --git a/tests/sim/objects/test_robot.py b/tests/sim/objects/test_robot.py index 5420020..773de76 100644 --- a/tests/sim/objects/test_robot.py +++ b/tests/sim/objects/test_robot.py @@ -52,9 +52,8 @@ class BaseRobotTest: def setup_simulation(self, sim_device): # Set up simulation with specified device (CPU or CUDA) - config = SimulationManagerCfg(headless=True, sim_device=sim_device) + config = SimulationManagerCfg(headless=True, sim_device=sim_device, num_envs=10) self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(10) # NUM_ARENAS = 10 cfg = DexforceW1Cfg.from_dict( { @@ -255,6 +254,38 @@ def test_setter_and_getter_with_control_part(self): ) self.robot.set_qpos(qpos=dummy_qpos, name="left_arm") + def test_robot_cfg_merge(self): + from copy import deepcopy + from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg + + cfg = deepcopy(self.robot.cfg) + + cfg_dict = { + "drive_pros": { + "max_effort": { + "(LEFT|RIGHT)_HAND_(THUMB[12]|INDEX|MIDDLE|RING|PINKY)": 1.0, + }, + }, + "solver_cfg": { + "left_arm": { + "tcp": np.eye(4), + } + }, + } + + cfg = merge_robot_cfg(cfg, cfg_dict) + + assert ( + cfg.drive_pros.max_effort[ + "(LEFT|RIGHT)_HAND_(THUMB[12]|INDEX|MIDDLE|RING|PINKY)" + ] + == 1.0 + ), "Drive properties merge failed." + + assert np.allclose( + cfg.solver_cfg["left_arm"].tcp, np.eye(4) + ), "Solver config merge failed." + def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() diff --git a/tests/sim/objects/test_soft_object.py b/tests/sim/objects/test_soft_object.py index 3099c2d..c4d9c54 100644 --- a/tests/sim/objects/test_soft_object.py +++ b/tests/sim/objects/test_soft_object.py @@ -40,6 +40,8 @@ def setup_simulation(self): physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) sim_device="cuda", enable_rt=False, # Enable ray tracing for better visuals + num_envs=4, + arena_space=3.0, ) # Create the simulation instance @@ -49,8 +51,7 @@ def setup_simulation(self): # Enable manual physics update for precise control self.n_envs = 4 - # Build multiple arenas if requested - self.sim.build_multiple_arenas(self.n_envs, space=3.0) + # add softbody to the scene self.cow: SoftObject = self.sim.add_soft_object( cfg=SoftObjectCfg( diff --git a/tests/sim/planners/test_motion_generator.py b/tests/sim/planners/test_motion_generator.py index 67e17dc..f321ffe 100644 --- a/tests/sim/planners/test_motion_generator.py +++ b/tests/sim/planners/test_motion_generator.py @@ -40,7 +40,6 @@ class BaseTestMotionGenerator(object): def setup_class(cls): cls.config = SimulationManagerCfg(headless=True, sim_device="cpu") cls.robot_sim = SimulationManager(cls.config) - cls.robot_sim.build_multiple_arenas(1) cls.robot_sim.set_manual_update(False) cfg_dict = { diff --git a/tests/sim/sensors/test_camera.py b/tests/sim/sensors/test_camera.py index b899994..2de23f2 100644 --- a/tests/sim/sensors/test_camera.py +++ b/tests/sim/sensors/test_camera.py @@ -33,10 +33,9 @@ class CameraTest: def setup_simulation(self, sim_device, enable_rt): # Setup SimulationManager config = SimulationManagerCfg( - headless=True, sim_device=sim_device, enable_rt=enable_rt + headless=True, sim_device=sim_device, enable_rt=enable_rt, num_envs=NUM_ENVS ) self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(NUM_ENVS) # Create batch of cameras cfg_dict = { "sensor_type": "Camera", @@ -144,11 +143,21 @@ def setup_method(self): self.setup_simulation("cpu", enable_rt=False) +class TestCameraRaster(CameraTest): + def setup_method(self): + self.setup_simulation("cuda", enable_rt=False) + + class TestCameraFastRT(CameraTest): def setup_method(self): self.setup_simulation("cpu", enable_rt=True) +class TestCameraFastRT(CameraTest): + def setup_method(self): + self.setup_simulation("cuda", enable_rt=True) + + if __name__ == "__main__": test = CameraTest() test.setup_simulation("cpu", enable_rt=False) diff --git a/tests/sim/sensors/test_stereo.py b/tests/sim/sensors/test_stereo.py index 737bae6..0119a2d 100644 --- a/tests/sim/sensors/test_stereo.py +++ b/tests/sim/sensors/test_stereo.py @@ -27,10 +27,9 @@ class StereoCameraTest: def setup_simulation(self, sim_device, enable_rt): # Setup SimulationManager config = SimulationManagerCfg( - headless=True, sim_device=sim_device, enable_rt=enable_rt + headless=True, sim_device=sim_device, enable_rt=enable_rt, num_envs=NUM_ENVS ) self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(4) # Create batch of cameras cfg_dict = { "sensor_type": "StereoCamera", @@ -150,6 +149,16 @@ def setup_method(self): self.setup_simulation("cpu", enable_rt=False) +class TestStereoCameraRaster(StereoCameraTest): + def setup_method(self): + self.setup_simulation("cuda", enable_rt=False) + + class TestStereoCameraFastRT(StereoCameraTest): def setup_method(self): self.setup_simulation("cpu", enable_rt=True) + + +class TestStereoCameraFastRT(StereoCameraTest): + def setup_method(self): + self.setup_simulation("cuda", enable_rt=True) diff --git a/tests/sim/solvers/test_differential_solver.py b/tests/sim/solvers/test_differential_solver.py index 4103f10..3ce4b73 100644 --- a/tests/sim/solvers/test_differential_solver.py +++ b/tests/sim/solvers/test_differential_solver.py @@ -33,7 +33,6 @@ def setup_simulation(self, solver_type: str): # Set up simulation with specified device (CPU or CUDA) config = SimulationManagerCfg(headless=True, sim_device="cpu") self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(1) # Load robot URDF file urdf = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") diff --git a/tests/sim/solvers/test_opw_solver.py b/tests/sim/solvers/test_opw_solver.py index 6df71a1..bf6a7c4 100644 --- a/tests/sim/solvers/test_opw_solver.py +++ b/tests/sim/solvers/test_opw_solver.py @@ -30,7 +30,6 @@ class BaseSolverTest: def setup_simulation(self): config = SimulationManagerCfg(headless=True, sim_device="cpu") self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(1) self.sim.set_manual_update(False) cfg_dict = { diff --git a/tests/sim/solvers/test_pink_solver.py b/tests/sim/solvers/test_pink_solver.py index 590ed6d..537374a 100644 --- a/tests/sim/solvers/test_pink_solver.py +++ b/tests/sim/solvers/test_pink_solver.py @@ -33,7 +33,6 @@ def setup_simulation(self, solver_type: str): # Set up simulation with specified device (CPU or CUDA) config = SimulationManagerCfg(headless=True, sim_device="cpu") self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(1) self.sim.set_manual_update(False) # Load robot URDF file diff --git a/tests/sim/solvers/test_pinocchio_solver.py b/tests/sim/solvers/test_pinocchio_solver.py index db9dd36..3650a85 100644 --- a/tests/sim/solvers/test_pinocchio_solver.py +++ b/tests/sim/solvers/test_pinocchio_solver.py @@ -33,7 +33,6 @@ def setup_simulation(self, solver_type: str): # Set up simulation with specified device (CPU or CUDA) config = SimulationManagerCfg(headless=True, sim_device="cpu") self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(1) self.sim.set_manual_update(False) # Load robot URDF file diff --git a/tests/sim/solvers/test_pytorch_solver.py b/tests/sim/solvers/test_pytorch_solver.py index 630a787..7551770 100644 --- a/tests/sim/solvers/test_pytorch_solver.py +++ b/tests/sim/solvers/test_pytorch_solver.py @@ -33,7 +33,6 @@ def setup_simulation(self, solver_type: str): # Set up simulation with specified device (CPU or CUDA) config = SimulationManagerCfg(headless=True, sim_device="cpu") self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(1) # Load robot URDF file urdf = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index 4d2a356..96a644b 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -132,7 +132,6 @@ def setup_simulation(self, solver_type: str, device: str = "cpu"): # Set up simulation with specified device (CPU or CUDA) config = SimulationManagerCfg(headless=True, sim_device=device) self.sim = SimulationManager(config) - self.sim.build_multiple_arenas(1) # Load robot URDF file urdf = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") diff --git a/tests/sim/utility/test_workspace_analyze.py b/tests/sim/utility/test_workspace_analyze.py new file mode 100644 index 0000000..2226057 --- /dev/null +++ b/tests/sim/utility/test_workspace_analyze.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import pytest +import numpy as np + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.objects import Robot +from embodichain.lab.sim.robots import CobotMagicCfg +from embodichain.lab.sim.utility.workspace_analyzer.configs import ( + SamplingConfig, + VisualizationConfig, + DimensionConstraint, +) + + +# Base test class for workspace analyzer +class BaseWorkspaceAnalyzeTest: + sim = None # Define as a class attribute + + def setup_simulation(self): + config = SimulationManagerCfg(headless=True, sim_device="cpu") + self.sim = SimulationManager(config) + self.sim.set_manual_update(False) + + cfg_dict = { + "uid": "CobotMagic", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [ + -0.3, + 0.3, + 1.0, + 1.0, + -1.2, + -1.2, + 0.0, + 0.0, + 0.6, + 0.6, + 0.0, + 0.0, + 0.05, + 0.05, + 0.05, + 0.05, + ], + "solver_cfg": { + "left_arm": { + "class_type": "OPWSolver", + "end_link_name": "left_link6", + "root_link_name": "left_arm_base", + "tcp": [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.143], [0, 0, 0, 1]], + }, + "right_arm": { + "class_type": "OPWSolver", + "end_link_name": "right_link6", + "root_link_name": "right_arm_base", + "tcp": [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.143], [0, 0, 0, 1]], + }, + }, + } + + self.robot: Robot = self.sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) + + @classmethod + def teardown_class(cls): + if cls.sim is not None: + try: + cls.sim.destroy() + print("sim destroyed successfully") + except Exception as e: + print(f"Error during sim.destroy(): {e}") + + +class TestJointWorkspaceAnalyzeTest(BaseWorkspaceAnalyzeTest): + def setup_method(self): + self.setup_simulation() + + def test_joint_workspace_analyze(self): + from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, + AnalysisMode, + ) + + joint_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.JOINT_SPACE, + sampling=SamplingConfig(num_samples=100), + control_part_name="left_arm", + ) + wa_joint = WorkspaceAnalyzer( + robot=self.robot, + config=joint_config, + sim_manager=self.sim, + ) + results = wa_joint.analyze(num_samples=100) + + assert "workspace_points" in results + assert results["workspace_points"].shape[0] > 0 + assert results["mode"] == "joint_space" + print( + f"Joint space results: {results['num_valid']}/{results['num_samples']} valid points" + ) + + +class TestCartesianWorkspaceAnalyzeTest(BaseWorkspaceAnalyzeTest): + def setup_method(self): + self.setup_simulation() + + def test_cartesian_workspace_analyze(self): + from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, + AnalysisMode, + ) + + # Get reference pose for IK orientation + left_qpos = torch.tensor([-0.3, 0.3, 1.0, 1.0, -1.2, -1.2]) + reference_pose = self.robot.compute_fk( + qpos=left_qpos.unsqueeze(0), + name="left_arm", + to_matrix=True, + ) + + cartesian_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.CARTESIAN_SPACE, + sampling=SamplingConfig( + num_samples=50 + ), # Smaller sample for faster testing + constraint=DimensionConstraint( + min_bounds=[-0.5, -0.5, 0.8], + max_bounds=[0.5, 0.5, 1.5], + ), + reference_pose=reference_pose[0].cpu().numpy(), + ik_samples_per_point=1, # Single IK attempt per point for speed + control_part_name="left_arm", + ) + + wa_cartesian = WorkspaceAnalyzer( + robot=self.robot, + config=cartesian_config, + sim_manager=self.sim, + ) + results = wa_cartesian.analyze(num_samples=50) + + assert "all_points" in results + assert "reachable_points" in results + assert results["all_points"].shape[0] > 0 + assert results["mode"] == "cartesian_space" + print( + f"Cartesian space results: {results['num_reachable']}/{results['num_samples']} reachable points" + ) + + +class TestPlaneWorkspaceAnalyzeTest(BaseWorkspaceAnalyzeTest): + def setup_method(self): + self.setup_simulation() + + def test_plane_workspace_analyze(self): + from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, + AnalysisMode, + ) + + # Get reference pose for IK orientation + left_qpos = torch.tensor([-0.3, 0.3, 1.0, 1.0, -1.2, -1.2]) + reference_pose = self.robot.compute_fk( + qpos=left_qpos.unsqueeze(0), + name="left_arm", + to_matrix=True, + ) + + plane_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.PLANE_SAMPLING, + sampling=SamplingConfig(num_samples=30), # Small sample for testing + plane_normal=torch.tensor([0.0, 0.0, 1.0]), # XY plane + plane_point=torch.tensor([0.0, 0.0, 1.2]), # At height 1.2m + reference_pose=reference_pose[0].cpu().numpy(), + ik_samples_per_point=1, # Single IK attempt per point for speed + visualization=VisualizationConfig( + enabled=False, # Disable visualization for testing + ), + control_part_name="left_arm", + ) + + wa_plane = WorkspaceAnalyzer( + robot=self.robot, + config=plane_config, + sim_manager=self.sim, + ) + results = wa_plane.analyze(num_samples=30) + + assert "all_points" in results + assert "reachable_points" in results + assert results["all_points"].shape[0] > 0 + assert results["mode"] == "plane_sampling" + assert "plane_sampling_config" in results + print( + f"Plane sampling results: {results['num_reachable']}/{results['num_samples']} reachable points" + ) + print( + f"Plane configuration: Normal={results['plane_sampling_config']['plane_normal']}, Point={results['plane_sampling_config']['plane_point']}" + ) + + +class TestWorkspaceAnalyzerComprehensive(BaseWorkspaceAnalyzeTest): + """Comprehensive test class for testing consistency between different modes""" + + def setup_method(self): + self.setup_simulation() + + def test_all_modes_consistency(self): + """Test basic functionality and result consistency of all three modes""" + from embodichain.lab.sim.utility.workspace_analyzer.workspace_analyzer import ( + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, + AnalysisMode, + ) + + # Get reference pose + left_qpos = torch.tensor([-0.3, 0.3, 1.0, 1.0, -1.2, -1.2]) + reference_pose = self.robot.compute_fk( + qpos=left_qpos.unsqueeze(0), + name="left_arm", + to_matrix=True, + ) + + num_samples = 20 # Small sample for fast testing + results = {} + + # Test joint space mode + joint_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.JOINT_SPACE, + sampling=SamplingConfig(num_samples=100), + control_part_name="left_arm", + ) + wa_joint = WorkspaceAnalyzer( + robot=self.robot, + config=joint_config, + sim_manager=self.sim, + ) + results["joint"] = wa_joint.analyze(num_samples=num_samples) + + # Test Cartesian space mode + cartesian_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.CARTESIAN_SPACE, + sampling=SamplingConfig(num_samples=num_samples), + constraint=DimensionConstraint( + min_bounds=[-0.5, -0.5, 0.8], + max_bounds=[0.5, 0.5, 1.5], + ), + reference_pose=reference_pose[0].cpu().numpy(), + ik_samples_per_point=1, + control_part_name="left_arm", + ) + wa_cartesian = WorkspaceAnalyzer( + robot=self.robot, + config=cartesian_config, + sim_manager=self.sim, + ) + results["cartesian"] = wa_cartesian.analyze(num_samples=num_samples) + + # Test plane sampling mode + plane_config = WorkspaceAnalyzerConfig( + mode=AnalysisMode.PLANE_SAMPLING, + sampling=SamplingConfig(num_samples=num_samples), + plane_normal=torch.tensor([0.0, 0.0, 1.0]), + plane_point=torch.tensor([0.0, 0.0, 1.2]), + reference_pose=reference_pose[0].cpu().numpy(), + ik_samples_per_point=1, + control_part_name="left_arm", + ) + wa_plane = WorkspaceAnalyzer( + robot=self.robot, + config=plane_config, + sim_manager=self.sim, + ) + results["plane"] = wa_plane.analyze(num_samples=num_samples) + + # Verify all modes return valid results + for mode_name, result in results.items(): + assert "analysis_time" in result + assert result["analysis_time"] > 0 + assert "mode" in result + print( + f"{mode_name.capitalize()} mode: {result['mode']}, Time: {result['analysis_time']:.2f}s" + ) + + # Verify joint space mode + assert results["joint"]["mode"] == "joint_space" + assert results["joint"]["num_valid"] > 0 + + # Verify Cartesian space mode + assert results["cartesian"]["mode"] == "cartesian_space" + assert "reachable_points" in results["cartesian"] + + # Verify plane sampling mode + assert results["plane"]["mode"] == "plane_sampling" + assert "plane_sampling_config" in results["plane"] + + print("\n=== Comprehensive Test Summary ===") + print( + f"Joint space: {results['joint']['num_valid']}/{results['joint']['num_samples']} valid points" + ) + print( + f"Cartesian space: {results['cartesian']['num_reachable']}/{results['cartesian']['num_samples']} reachable points" + ) + print( + f"Plane sampling: {results['plane']['num_reachable']}/{results['plane']['num_samples']} reachable points" + ) + + +if __name__ == "__main__": + np.set_printoptions(precision=5, suppress=True) + pytest_args = ["-v", "-s", __file__] + pytest.main(pytest_args)