mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-29 15:39:56 +00:00
add robot models xml etc
remove isaaclab, add test scripts added simulator to the hub remove mujoco env
This commit is contained in:
Martino Russi
@@ -173,7 +173,3 @@ outputs/
|
||||
|
||||
# Dev folders
|
||||
.cache/*
|
||||
*.stl
|
||||
*.urdf
|
||||
*.xml
|
||||
*.part
|
||||
|
||||
@@ -1,256 +0,0 @@
|
||||
# Camera System for MuJoCo G1 Simulator
|
||||
|
||||
## Overview
|
||||
|
||||
The simulator has two cameras defined:
|
||||
|
||||
### 1. **`head_camera`** - Robot Ego-View
|
||||
- **Location**: Attached to `torso_link` body
|
||||
- **Position**: `[0.06, 0.0, 0.45]` relative to torso (6cm forward, 45cm up)
|
||||
- **Orientation**: `euler="0 -0.8 -1.57"` (facing forward, slightly tilted down)
|
||||
- **FOV**: 90 degrees
|
||||
- **Purpose**: First-person view from the robot's perspective (like a head-mounted camera)
|
||||
|
||||
### 2. **`global_view`** - Third-Person View
|
||||
- **Location**: Fixed in world coordinates
|
||||
- **Position**: `[2.910, -5.040, 3.860]` (behind and above the robot)
|
||||
- **Purpose**: External observer view for visualization
|
||||
|
||||
## How Camera Publishing Works
|
||||
|
||||
The camera system uses a **zero-copy architecture** with three components:
|
||||
|
||||
```
|
||||
┌─────────────────────┐
|
||||
│ MuJoCo Simulator │
|
||||
│ (Main Process) │
|
||||
│ │
|
||||
│ 1. Render cameras │
|
||||
│ 2. Copy to shmem │──┐
|
||||
└─────────────────────┘ │
|
||||
│ Shared Memory
|
||||
│ (fast IPC)
|
||||
┌────▼────────────────┐
|
||||
│ Image Publisher │
|
||||
│ (Subprocess) │
|
||||
│ │
|
||||
│ 3. Encode images │
|
||||
│ 4. ZMQ publish │
|
||||
└─────────┬───────────┘
|
||||
│
|
||||
│ TCP (ZMQ)
|
||||
│ port 5555
|
||||
▼
|
||||
┌─────────────────────┐
|
||||
│ Your Policy/Client │
|
||||
│ (Subscribe) │
|
||||
└─────────────────────┘
|
||||
```
|
||||
|
||||
### Key Technologies:
|
||||
- **MuJoCo Renderer**: Captures RGB images from virtual cameras
|
||||
- **Shared Memory (`multiprocessing.shared_memory`)**: Zero-copy transfer between processes
|
||||
- **ZMQ (ZeroMQ)**: Network socket for publishing images (TCP)
|
||||
- **No ROS2 required!** Pure Python multiprocessing
|
||||
|
||||
## Usage
|
||||
|
||||
### Basic Simulation (No Camera Publishing)
|
||||
```bash
|
||||
python run_sim.py
|
||||
```
|
||||
|
||||
### With Camera Publishing
|
||||
```bash
|
||||
# Publish head camera on default port 5555
|
||||
python run_sim.py --publish-images
|
||||
|
||||
# Publish multiple cameras
|
||||
python run_sim.py --publish-images --cameras head_camera global_view
|
||||
|
||||
# Custom port
|
||||
python run_sim.py --publish-images --camera-port 6000
|
||||
```
|
||||
|
||||
### Viewing Camera Streams
|
||||
|
||||
In a **separate terminal**, run the camera viewer:
|
||||
|
||||
```bash
|
||||
# Basic usage (default: localhost:5555)
|
||||
python view_cameras.py
|
||||
|
||||
# Custom host/port
|
||||
python view_cameras.py --host 192.168.1.100 --port 6000
|
||||
|
||||
# Save images to directory
|
||||
python view_cameras.py --save ./camera_recordings
|
||||
|
||||
# Adjust display rate
|
||||
python view_cameras.py --fps 60
|
||||
```
|
||||
|
||||
**Keyboard Controls:**
|
||||
- `q`: Quit viewer
|
||||
- `s`: Save snapshot of current frame
|
||||
|
||||
**Example Workflow:**
|
||||
```bash
|
||||
# Terminal 1: Start simulator with camera publishing
|
||||
python run_sim.py --publish-images
|
||||
|
||||
# Terminal 2: View the camera feed
|
||||
python view_cameras.py
|
||||
```
|
||||
|
||||
### Receiving Images in Your Code
|
||||
|
||||
```python
|
||||
import zmq
|
||||
import numpy as np
|
||||
from gr00t_wbc.control.sensor.sensor_server import ImageMessageSchema
|
||||
|
||||
# Connect to camera publisher
|
||||
context = zmq.Context()
|
||||
socket = context.socket(zmq.SUB)
|
||||
socket.connect("tcp://localhost:5555")
|
||||
socket.setsockopt(zmq.SUBSCRIBE, b"") # Subscribe to all messages
|
||||
|
||||
while True:
|
||||
# Receive serialized image data
|
||||
data = socket.recv_pyobj()
|
||||
|
||||
# Decode images
|
||||
if "head_camera" in data:
|
||||
# Decode image (returns numpy array HxWx3 uint8)
|
||||
img = decode_image(data["head_camera"])
|
||||
|
||||
# Use image for your policy
|
||||
process_observation(img)
|
||||
```
|
||||
|
||||
## Camera Configuration
|
||||
|
||||
Edit `config.yaml` to change camera settings:
|
||||
|
||||
```yaml
|
||||
IMAGE_DT: 0.033333 # Publishing rate (30 Hz)
|
||||
ENABLE_OFFSCREEN: false # Enable for camera rendering
|
||||
MP_START_METHOD: "spawn" # Multiprocessing method
|
||||
```
|
||||
|
||||
Or programmatically in `run_sim.py`:
|
||||
|
||||
```python
|
||||
camera_configs = {
|
||||
"head_camera": {
|
||||
"height": 480,
|
||||
"width": 640
|
||||
},
|
||||
"custom_camera": {
|
||||
"height": 224,
|
||||
"width": 224
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Adding Custom Cameras
|
||||
|
||||
Edit `assets/g1_29dof_with_hand.xml` or `assets/scene_43dof.xml`:
|
||||
|
||||
```xml
|
||||
<!-- Camera attached to robot body -->
|
||||
<body name="torso_link" pos="0 0 0.019">
|
||||
<camera name="my_camera" pos="0.1 0.0 0.5" euler="0 0 0" fovy="60"/>
|
||||
</body>
|
||||
|
||||
<!-- Camera in world coordinates -->
|
||||
<worldbody>
|
||||
<camera name="side_view" pos="0 -3.0 1.5" xyaxes="1 0 0 0 0.5 0.866"/>
|
||||
</worldbody>
|
||||
```
|
||||
|
||||
Then publish it:
|
||||
```bash
|
||||
python run_sim.py --publish-images --cameras my_camera
|
||||
```
|
||||
|
||||
## Performance Notes
|
||||
|
||||
- **Rendering overhead**: ~5-10ms per camera per frame @ 640x480
|
||||
- **Publishing overhead**: ~2-3ms for encoding + network
|
||||
- Image publishing runs in **separate subprocess** to not block simulation
|
||||
- Uses **shared memory** for fast inter-process image transfer
|
||||
- Target: 30 FPS camera publishing while maintaining 500 Hz simulation
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### No images received?
|
||||
1. Check if offscreen rendering is enabled (`--publish-images` flag)
|
||||
2. Verify ZMQ port is not blocked
|
||||
3. Check camera exists in scene XML
|
||||
|
||||
### Images are delayed?
|
||||
- Reduce `IMAGE_DT` in config
|
||||
- Lower camera resolution
|
||||
- Use fewer cameras
|
||||
|
||||
### "Camera not found" error?
|
||||
- Verify camera name in XML matches config
|
||||
- Check XML syntax is valid
|
||||
- Ensure MuJoCo model loads successfully
|
||||
|
||||
## Quick Reference
|
||||
|
||||
### File Structure
|
||||
```
|
||||
mujoco_sim_g1/
|
||||
├── run_sim.py # Simulator launcher
|
||||
├── view_cameras.py # Camera viewer (this file!)
|
||||
├── config.yaml # Simulator config
|
||||
├── assets/
|
||||
│ ├── scene_43dof.xml # Scene with global_view camera
|
||||
│ └── g1_29dof_with_hand.xml # Robot model with head_camera
|
||||
└── sim/
|
||||
├── base_sim.py # MuJoCo environment
|
||||
├── sensor_utils.py # ZMQ camera server/client
|
||||
└── image_publish_utils.py # Multiprocessing image publisher
|
||||
```
|
||||
|
||||
### Camera Definitions
|
||||
|
||||
Edit these files to modify cameras:
|
||||
|
||||
**`assets/g1_29dof_with_hand.xml`** - Robot-attached cameras:
|
||||
```xml
|
||||
<body name="torso_link" pos="0 0 0.019">
|
||||
<camera name="head_camera" pos="0.06 0.0 0.45" euler="0 -0.8 -1.57" fovy="90"/>
|
||||
</body>
|
||||
```
|
||||
|
||||
**`assets/scene_43dof.xml`** - World-frame cameras:
|
||||
```xml
|
||||
<worldbody>
|
||||
<camera name="global_view" pos="2.910 -5.040 3.860" xyaxes="0.866 0.500 0.000 -0.250 0.433 0.866" fovy="45"/>
|
||||
</worldbody>
|
||||
```
|
||||
|
||||
### Complete Example
|
||||
|
||||
```bash
|
||||
# Terminal 1: Start simulator with camera publishing
|
||||
cd mujoco_sim_g1
|
||||
python run_sim.py --publish-images --cameras head_camera global_view
|
||||
|
||||
# Terminal 2: View cameras in real-time
|
||||
python view_cameras.py
|
||||
|
||||
# Terminal 3: Use in your policy (Python code)
|
||||
from sim.sensor_utils import SensorClient, ImageUtils
|
||||
client = SensorClient()
|
||||
client.start_client("localhost", 5555)
|
||||
data = client.receive_message()
|
||||
img = ImageUtils.decode_image(data["head_camera"])
|
||||
# img is now numpy array (H, W, 3) in BGR format
|
||||
```
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
Thu Nov 20 17:42:09 2025
|
||||
WARNING: Nan, Inf or huge value in QACC at DOF 0. The simulation is unstable. Time = 124.1560.
|
||||
|
||||
@@ -1,167 +0,0 @@
|
||||
# Standalone MuJoCo Simulator for Unitree G1
|
||||
|
||||
This is a standalone MuJoCo physics simulator for the Unitree G1 robot, extracted from the GR00T-WholeBodyControl repository.
|
||||
|
||||
## Features
|
||||
|
||||
- **Physics Simulation**: Runs G1 robot in MuJoCo at 500Hz (2ms timestep)
|
||||
- **DDS Communication**: Uses Unitree SDK2 DDS for robot communication
|
||||
- **Compatible**: Works with existing `unitree_g1.py` control code via DDS
|
||||
- **Visualization**: Real-time 3D visualization of robot motion
|
||||
|
||||
## Directory Structure
|
||||
|
||||
```
|
||||
mujoco_sim_g1/
|
||||
├── requirements.txt # Python dependencies
|
||||
├── run_sim.py # Main launcher script
|
||||
├── config.yaml # Simulation configuration
|
||||
├── sim/ # Simulation modules
|
||||
│ ├── base_sim.py
|
||||
│ ├── simulator_factory.py
|
||||
│ ├── unitree_sdk2py_bridge.py
|
||||
│ └── ...
|
||||
└── assets/ # Robot models
|
||||
├── g1_29dof_with_hand.xml
|
||||
└── meshes/*.STL
|
||||
```
|
||||
|
||||
## Installation
|
||||
|
||||
### 1. Create Virtual Environment (Recommended)
|
||||
|
||||
```bash
|
||||
cd mujoco_sim_g1
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate # On Windows: venv\Scripts\activate
|
||||
```
|
||||
|
||||
### 2. Install Dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**Note**: If you encounter issues with `rclpy`, you can comment out the ROS2 imports in `sim/base_sim.py` (lines 11, 609-617) if you don't need camera publishing.
|
||||
|
||||
## Usage
|
||||
|
||||
### Basic Usage
|
||||
|
||||
```bash
|
||||
# Activate environment
|
||||
source venv/bin/activate
|
||||
|
||||
# Run the simulator
|
||||
python run_sim.py
|
||||
```
|
||||
|
||||
The simulator will:
|
||||
1. Load the G1 robot model
|
||||
2. Initialize DDS communication on domain 0
|
||||
3. Open a MuJoCo visualization window
|
||||
4. Start listening for motor commands via DDS
|
||||
|
||||
### Running with Your Robot Control Code
|
||||
|
||||
Once the simulator is running, you can control it from another terminal:
|
||||
|
||||
```bash
|
||||
# In another terminal
|
||||
cd /home/yope/Documents/lerobot
|
||||
conda activate unitree_lerobot
|
||||
|
||||
# Run your existing control code
|
||||
python test_locomotion_minimal.py
|
||||
```
|
||||
|
||||
Your `unitree_g1.py` code will automatically connect to the simulator via DDS!
|
||||
|
||||
## Configuration
|
||||
|
||||
Edit `config.yaml` to customize:
|
||||
|
||||
- **SIMULATE_DT**: Simulation timestep (default: 0.002s = 500Hz)
|
||||
- **DOMAIN_ID**: DDS domain ID (default: 0)
|
||||
- **ENABLE_ONSCREEN**: Show visualization (default: true)
|
||||
- **USE_JOYSTICK**: Enable gamepad control (default: false)
|
||||
- **ROBOT_SCENE**: Path to MuJoCo XML scene
|
||||
|
||||
### PD Gains
|
||||
|
||||
The simulator uses the following PD gains (matching NVIDIA GR00T):
|
||||
|
||||
**Legs (indices 0-11):**
|
||||
- Hip joints: KP=150, KD=2
|
||||
- Knee joints: KP=300, KD=4
|
||||
- Ankle joints: KP=40, KD=2
|
||||
|
||||
**Waist (indices 12-14):**
|
||||
- All waist joints: KP=250, KD=5
|
||||
|
||||
**Arms (indices 15-28):**
|
||||
- Shoulders: KP=100, KD=2-5
|
||||
- Elbows/Wrists: KP=20-40, KD=1-2
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### ImportError: cannot import name 'ChannelFactoryInitialize'
|
||||
|
||||
```bash
|
||||
pip install --upgrade unitree-sdk2py
|
||||
```
|
||||
|
||||
### ROS2/rclpy errors
|
||||
|
||||
If you don't need camera publishing, edit `sim/base_sim.py`:
|
||||
- Comment out line 11: `import rclpy`
|
||||
- Comment out lines 609-617 (ROS2 initialization)
|
||||
|
||||
### Meshes not found
|
||||
|
||||
Make sure mesh paths in `assets/g1_29dof_with_hand.xml` are relative:
|
||||
```xml
|
||||
<mesh file="meshes/pelvis.STL"/> <!-- Not absolute path -->
|
||||
```
|
||||
|
||||
### Robot falls immediately
|
||||
|
||||
Check that:
|
||||
1. PD gains match NVIDIA's values (see config.yaml)
|
||||
2. Velocity command scaling is correct (ang_vel_scale=0.25, cmd_scale=[2.0, 2.0, 0.25])
|
||||
3. Observations for missing joints are zeroed out (indices 12, 14, 20, 21, 27, 28)
|
||||
|
||||
## Technical Details
|
||||
|
||||
### Communication
|
||||
|
||||
The simulator publishes:
|
||||
- **`rt/lowstate`**: Robot state (joint positions, velocities, IMU, etc.)
|
||||
- **`rt/wirelesscontroller`**: Wireless remote controller state (if joystick enabled)
|
||||
|
||||
The simulator subscribes to:
|
||||
- **`rt/lowcmd`**: Motor commands (position, velocity, torque, KP, KD)
|
||||
|
||||
### Coordinate Frames
|
||||
|
||||
- **World frame**: Z-up
|
||||
- **Joint ordering**: 29 DOF (12 legs + 3 waist + 14 arms)
|
||||
- **IMU**: Quaternion in [w, x, y, z] format
|
||||
|
||||
### Performance
|
||||
|
||||
- Simulation runs at ~500Hz (2ms timestep)
|
||||
- Viewer updates at ~50Hz (20ms)
|
||||
- Typical CPU usage: 20-40% on single core
|
||||
|
||||
## Files from GR00T-WholeBodyControl
|
||||
|
||||
This standalone simulator was extracted from:
|
||||
- `gr00t_wbc/control/envs/g1/sim/` (simulation modules)
|
||||
- `gr00t_wbc/control/robot_model/model_data/g1/` (robot model files)
|
||||
- `gr00t_wbc/control/main/teleop/configs/g1_29dof_gear_wbc.yaml` (configuration)
|
||||
|
||||
## License
|
||||
|
||||
Follows the license of the original GR00T-WholeBodyControl repository.
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 67 KiB |
@@ -1,423 +0,0 @@
|
||||
GEAR_WBC_CONFIG: "gr00t_wbc/sim2mujoco/resources/robots/g1/g1_gear_wbc.yaml"
|
||||
|
||||
# copy from g1_43dof_hist.yaml
|
||||
ROBOT_TYPE: 'g1_29dof' # Robot name, "go2", "b2", "b2w", "h1", "go2w", "g1"
|
||||
ROBOT_SCENE: "assets/scene_43dof.xml" # Robot scene with 3-finger dex hands
|
||||
|
||||
DOMAIN_ID: 0 # Domain id
|
||||
# Network Interface, "lo" for simulation and the one with "192.168.123.222" for real robot
|
||||
INTERFACE: "lo"
|
||||
SIMULATOR: "mujoco" # "robocasa"
|
||||
|
||||
USE_JOYSTICK: 0 # Simulate Unitree WirelessController using a gamepad (0: disable, 1: enable)
|
||||
JOYSTICK_TYPE: "xbox" # support "xbox" and "switch" gamepad layout
|
||||
JOYSTICK_DEVICE: 0 # Joystick number
|
||||
|
||||
FREE_BASE: False
|
||||
enable_waist: False # Enable waist joints control (False = band on torso, True = band on pelvis)
|
||||
|
||||
PRINT_SCENE_INFORMATION: True # Print link, joint and sensors information of robot
|
||||
ENABLE_ELASTIC_BAND: True # Virtual spring band, used for lifting h1
|
||||
|
||||
SIMULATE_DT: 0.002 # Simulation timestep (500Hz)
|
||||
ENABLE_ONSCREEN: true # Show visualization window
|
||||
ENABLE_OFFSCREEN: false # Offscreen rendering for cameras (set to true to enable camera publishing)
|
||||
IMAGE_DT: 0.033333 # Image publishing rate ~30Hz
|
||||
MP_START_METHOD: "spawn" # Multiprocessing start method for image subprocess
|
||||
VIEWER_DT: 0.02 # Viewer update time
|
||||
REWARD_DT: 0.02
|
||||
USE_SENSOR: False
|
||||
USE_HISTORY: True
|
||||
USE_HISTORY_LOCO: True
|
||||
USE_HISTORY_MIMIC: True
|
||||
|
||||
GAIT_PERIOD: 0.9 # 1.25
|
||||
|
||||
MOTOR2JOINT: [0, 1, 2, 3, 4, 5,
|
||||
6, 7, 8, 9, 10, 11,
|
||||
12, 13, 14,
|
||||
15, 16, 17, 18, 19, 20, 21,
|
||||
22, 23, 24, 25, 26, 27, 28]
|
||||
|
||||
JOINT2MOTOR: [0, 1, 2, 3, 4, 5,
|
||||
6, 7, 8, 9, 10, 11,
|
||||
12, 13, 14,
|
||||
15, 16, 17, 18, 19, 20, 21,
|
||||
22, 23, 24, 25, 26, 27, 28]
|
||||
|
||||
|
||||
UNITREE_LEGGED_CONST:
|
||||
HIGHLEVEL: 0xEE
|
||||
LOWLEVEL: 0xFF
|
||||
TRIGERLEVEL: 0xF0
|
||||
PosStopF: 2146000000.0
|
||||
VelStopF: 16000.0
|
||||
MODE_MACHINE: 5
|
||||
MODE_PR: 0
|
||||
|
||||
JOINT_KP: [
|
||||
100, 100, 100, 200, 20, 20,
|
||||
100, 100, 100, 200, 20, 20,
|
||||
400, 400, 400,
|
||||
90, 60, 20, 60, 4, 4, 4,
|
||||
90, 60, 20, 60, 4, 4, 4
|
||||
]
|
||||
|
||||
|
||||
JOINT_KD: [
|
||||
2.5, 2.5, 2.5, 5, 0.2, 0.1,
|
||||
2.5, 2.5, 2.5, 5, 0.2, 0.1,
|
||||
5.0, 5.0, 5.0,
|
||||
2.0, 1.0, 0.4, 1.0, 0.2, 0.2, 0.2,
|
||||
2.0, 1.0, 0.4, 1.0, 0.2, 0.2, 0.2
|
||||
]
|
||||
|
||||
# arm kp
|
||||
# soft kp, safe, test it first
|
||||
# 50, 50, 20, 20, 10, 10, 10
|
||||
# hard kp, use only if policy is safe
|
||||
# 200, 200, 80, 80, 50, 50, 50,
|
||||
|
||||
# MOTOR_KP: [
|
||||
# 100, 100, 100, 200, 20, 20,
|
||||
# 100, 100, 100, 200, 20, 20,
|
||||
# 400, 400, 400,
|
||||
# 50, 50, 20, 20, 10, 10, 10,
|
||||
# 50, 50, 20, 20, 10, 10, 10
|
||||
# ]
|
||||
|
||||
MOTOR_KP: [
|
||||
150, 150, 150, 200, 40, 40,
|
||||
150, 150, 150, 200, 40, 40,
|
||||
800, 250, 800,
|
||||
100, 100, 40, 40, 100, 100, 100,
|
||||
100, 100, 40, 40, 100, 100, 100
|
||||
]
|
||||
|
||||
MOTOR_KD: [
|
||||
2, 2, 2, 4, 2, 2,
|
||||
2, 2, 2, 4, 2, 2,
|
||||
10, 5, 10,
|
||||
5, 5, 2, 2, 5, 5, 5,
|
||||
5, 5, 2, 2, 5, 5, 5
|
||||
]
|
||||
|
||||
# MOTOR_KP: [
|
||||
# 100, 100, 100, 200, 20, 20,
|
||||
# 100, 100, 100, 200, 20, 20,
|
||||
# 400, 400, 400,
|
||||
# 90, 60, 20, 60, 4, 4, 4,
|
||||
# 90, 60, 20, 60, 4, 4, 4
|
||||
# ]
|
||||
|
||||
|
||||
# MOTOR_KD: [
|
||||
# 2.5, 2.5, 2.5, 5, 0.2, 0.1,
|
||||
# 2.5, 2.5, 2.5, 5, 0.2, 0.1,
|
||||
# 5.0, 5.0, 5.0,
|
||||
# 2.0, 1.0, 0.4, 1.0, 0.2, 0.2, 0.2,
|
||||
# 2.0, 1.0, 0.4, 1.0, 0.2, 0.2, 0.2
|
||||
# ]
|
||||
|
||||
|
||||
WeakMotorJointIndex:
|
||||
left_hip_yaw_joint: 0
|
||||
left_hip_roll_joint: 1
|
||||
left_hip_pitch_joint: 2
|
||||
left_knee_joint: 3
|
||||
left_ankle_pitch_joint: 4
|
||||
left_ankle_roll_joint: 5
|
||||
right_hip_yaw_joint: 6
|
||||
right_hip_roll_joint: 7
|
||||
right_hip_pitch_joint: 8
|
||||
right_knee_joint: 9
|
||||
right_ankle_pitch_joint: 10
|
||||
right_ankle_roll_joint: 11
|
||||
waist_yaw_joint : 12
|
||||
waist_roll_joint : 13
|
||||
waist_pitch_joint : 14
|
||||
left_shoulder_pitch_joint: 15
|
||||
left_shoulder_roll_joint: 16
|
||||
left_shoulder_yaw_joint: 17
|
||||
left_elbow_joint: 18
|
||||
left_wrist_roll_joint: 19
|
||||
left_wrist_pitch_joint: 20
|
||||
left_wrist_yaw_joint: 21
|
||||
right_shoulder_pitch_joint: 22
|
||||
right_shoulder_roll_joint: 23
|
||||
right_shoulder_yaw_joint: 24
|
||||
right_elbow_joint: 25
|
||||
right_wrist_roll_joint: 26
|
||||
right_wrist_pitch_joint: 27
|
||||
right_wrist_yaw_joint: 28
|
||||
|
||||
NUM_MOTORS: 29
|
||||
NUM_JOINTS: 29
|
||||
NUM_HAND_MOTORS: 7
|
||||
NUM_HAND_JOINTS: 7
|
||||
NUM_UPPER_BODY_JOINTS: 17
|
||||
|
||||
DEFAULT_DOF_ANGLES: [
|
||||
-0.1, # left_hip_pitch_joint
|
||||
0.0, # left_hip_roll_joint
|
||||
0.0, # left_hip_yaw_joint
|
||||
0.3, # left_knee_joint
|
||||
-0.2, # left_ankle_pitch_joint
|
||||
0.0, # left_ankle_roll_joint
|
||||
-0.1, # right_hip_pitch_joint
|
||||
0.0, # right_hip_roll_joint
|
||||
0.0, # right_hip_yaw_joint
|
||||
0.3, # right_knee_joint
|
||||
-0.2, # right_ankle_pitch_joint
|
||||
0.0, # right_ankle_roll_joint
|
||||
0.0, # waist_yaw_joint
|
||||
0.0, # waist_roll_joint
|
||||
0.0, # waist_pitch_joint
|
||||
0.0, # left_shoulder_pitch_joint
|
||||
0.0, # left_shoulder_roll_joint
|
||||
0.0, # left_shoulder_yaw_joint
|
||||
0.0, # left_elbow_joint
|
||||
0.0, # left_wrist_roll_joint
|
||||
0.0, # left_wrist_pitch_joint
|
||||
0.0, # left_wrist_yaw_joint
|
||||
0.0, # right_shoulder_pitch_joint
|
||||
0.0, # right_shoulder_roll_joint
|
||||
0.0, # right_shoulder_yaw_joint
|
||||
0.0, # right_elbow_joint
|
||||
0.0, # right_wrist_roll_joint
|
||||
0.0, # right_wrist_pitch_joint
|
||||
0.0 # right_wrist_yaw_joint
|
||||
]
|
||||
|
||||
DEFAULT_MOTOR_ANGLES: [
|
||||
-0.1, # left_hip_pitch_joint
|
||||
0.0, # left_hip_roll_joint
|
||||
0.0, # left_hip_yaw_joint
|
||||
0.3, # left_knee_joint
|
||||
-0.2, # left_ankle_pitch_joint
|
||||
0.0, # left_ankle_roll_joint
|
||||
-0.1, # right_hip_pitch_joint
|
||||
0.0, # right_hip_roll_joint
|
||||
0.0, # right_hip_yaw_joint
|
||||
0.3, # right_knee_joint
|
||||
-0.2, # right_ankle_pitch_joint
|
||||
0.0, # right_ankle_roll_joint
|
||||
0.0, # waist_yaw_joint
|
||||
0.0, # waist_roll_joint
|
||||
0.0, # waist_pitch_joint
|
||||
0.0, # left_shoulder_pitch_joint
|
||||
0.0, # left_shoulder_roll_joint
|
||||
0.0, # left_shoulder_yaw_joint
|
||||
0.0, # left_elbow_joint
|
||||
0.0, # left_wrist_roll_joint
|
||||
0.0, # left_wrist_pitch_joint
|
||||
0.0, # left_wrist_yaw_joint
|
||||
0.0, # right_shoulder_pitch_joint
|
||||
0.0, # right_shoulder_roll_joint
|
||||
0.0, # right_shoulder_yaw_joint
|
||||
0.0, # right_elbow_joint
|
||||
0.0, # right_wrist_roll_joint
|
||||
0.0, # right_wrist_pitch_joint
|
||||
0.0 # right_wrist_yaw_joint
|
||||
]
|
||||
|
||||
motor_pos_lower_limit_list: [-2.5307, -0.5236, -2.7576, -0.087267, -0.87267, -0.2618,
|
||||
-2.5307, -2.9671, -2.7576, -0.087267, -0.87267, -0.2618,
|
||||
-2.618, -0.52, -0.52,
|
||||
-3.0892, -1.5882, -2.618, -1.0472,
|
||||
-1.972222054, -1.61443, -1.61443,
|
||||
-3.0892, -2.2515, -2.618, -1.0472,
|
||||
-1.972222054, -1.61443, -1.61443]
|
||||
motor_pos_upper_limit_list: [2.8798, 2.9671, 2.7576, 2.8798, 0.5236, 0.2618,
|
||||
2.8798, 0.5236, 2.7576, 2.8798, 0.5236, 0.2618,
|
||||
2.618, 0.52, 0.52,
|
||||
2.6704, 2.2515, 2.618, 2.0944,
|
||||
1.972222054, 1.61443, 1.61443,
|
||||
2.6704, 1.5882, 2.618, 2.0944,
|
||||
1.972222054, 1.61443, 1.61443]
|
||||
motor_vel_limit_list: [32.0, 32.0, 32.0, 20.0, 37.0, 37.0,
|
||||
32.0, 32.0, 32.0, 20.0, 37.0, 37.0,
|
||||
32.0, 37.0, 37.0,
|
||||
37.0, 37.0, 37.0, 37.0,
|
||||
37.0, 22.0, 22.0,
|
||||
37.0, 37.0, 37.0, 37.0,
|
||||
37.0, 22.0, 22.0]
|
||||
motor_effort_limit_list: [88.0, 88.0, 88.0, 139.0, 50.0, 50.0,
|
||||
88.0, 88.0, 88.0, 139.0, 50.0, 50.0,
|
||||
88.0, 50.0, 50.0,
|
||||
25.0, 25.0, 25.0, 25.0,
|
||||
25.0, 5.0, 5.0,
|
||||
2.45, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7,
|
||||
25.0, 25.0, 25.0, 25.0,
|
||||
25.0, 5.0, 5.0,
|
||||
2.45, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7]
|
||||
history_config: {
|
||||
base_ang_vel: 4,
|
||||
projected_gravity: 4,
|
||||
command_lin_vel: 4,
|
||||
command_ang_vel: 4,
|
||||
command_base_height: 4,
|
||||
command_stand: 4,
|
||||
ref_upper_dof_pos: 4,
|
||||
dof_pos: 4,
|
||||
dof_vel: 4,
|
||||
actions: 4,
|
||||
# phase_time: 4,
|
||||
ref_motion_phase: 4,
|
||||
sin_phase: 4,
|
||||
cos_phase: 4
|
||||
}
|
||||
history_loco_config: {
|
||||
base_ang_vel: 4,
|
||||
projected_gravity: 4,
|
||||
command_lin_vel: 4,
|
||||
command_ang_vel: 4,
|
||||
# command_base_height: 4,
|
||||
command_stand: 4,
|
||||
ref_upper_dof_pos: 4,
|
||||
dof_pos: 4,
|
||||
dof_vel: 4,
|
||||
actions: 4,
|
||||
# phase_time: 4,
|
||||
sin_phase: 4,
|
||||
cos_phase: 4
|
||||
}
|
||||
history_loco_height_config: {
|
||||
base_ang_vel: 4,
|
||||
projected_gravity: 4,
|
||||
command_lin_vel: 4,
|
||||
command_ang_vel: 4,
|
||||
command_base_height: 4,
|
||||
command_stand: 4,
|
||||
ref_upper_dof_pos: 4,
|
||||
dof_pos: 4,
|
||||
dof_vel: 4,
|
||||
actions: 4,
|
||||
# phase_time: 4,
|
||||
sin_phase: 4,
|
||||
cos_phase: 4
|
||||
}
|
||||
history_mimic_config: {
|
||||
base_ang_vel: 4,
|
||||
projected_gravity: 4,
|
||||
dof_pos: 4,
|
||||
dof_vel: 4,
|
||||
actions: 4,
|
||||
ref_motion_phase: 4,
|
||||
}
|
||||
obs_dims: {
|
||||
base_lin_vel: 3,
|
||||
base_ang_vel: 3,
|
||||
projected_gravity: 3,
|
||||
command_lin_vel: 2,
|
||||
command_ang_vel: 1,
|
||||
command_stand: 1,
|
||||
command_base_height: 1,
|
||||
ref_upper_dof_pos: 17, # upper body actions
|
||||
dof_pos: 29,
|
||||
dof_vel: 29,
|
||||
# actions: 12, # lower body actions
|
||||
actions: 29, # full body actions
|
||||
phase_time: 1,
|
||||
ref_motion_phase: 1, # mimic motion phase
|
||||
sin_phase: 1,
|
||||
cos_phase: 1,
|
||||
}
|
||||
obs_loco_dims: {
|
||||
base_lin_vel: 3,
|
||||
base_ang_vel: 3,
|
||||
projected_gravity: 3,
|
||||
command_lin_vel: 2,
|
||||
command_ang_vel: 1,
|
||||
command_stand: 1,
|
||||
command_base_height: 1,
|
||||
ref_upper_dof_pos: 17, # upper body actions
|
||||
dof_pos: 29,
|
||||
dof_vel: 29,
|
||||
actions: 12, # lower body actions
|
||||
phase_time: 1,
|
||||
sin_phase: 1,
|
||||
cos_phase: 1,
|
||||
}
|
||||
obs_mimic_dims: {
|
||||
base_lin_vel: 3,
|
||||
base_ang_vel: 3,
|
||||
projected_gravity: 3,
|
||||
dof_pos: 29,
|
||||
dof_vel: 29,
|
||||
actions: 29, # full body actions
|
||||
ref_motion_phase: 1, # mimic motion phase
|
||||
}
|
||||
obs_scales: {
|
||||
base_lin_vel: 2.0,
|
||||
base_ang_vel: 0.25,
|
||||
projected_gravity: 1.0,
|
||||
command_lin_vel: 1,
|
||||
command_ang_vel: 1,
|
||||
command_stand: 1,
|
||||
command_base_height: 2, # Yuanhang: it's 2, not 1!
|
||||
ref_upper_dof_pos: 1.0,
|
||||
dof_pos: 1.0,
|
||||
dof_vel: 0.05,
|
||||
history: 1.0,
|
||||
history_loco: 1.0,
|
||||
history_mimic: 1.0,
|
||||
actions: 1.0,
|
||||
phase_time: 1.0,
|
||||
ref_motion_phase: 1.0,
|
||||
sin_phase: 1.0,
|
||||
cos_phase: 1.0
|
||||
}
|
||||
|
||||
loco_upper_body_dof_pos: [
|
||||
0.0, 0.0, 0.0, # waist
|
||||
0.0, 0.3, 0.0, 1.0, # left shoulder and elbow
|
||||
0.0, 0.0, 0.0, # left wrist
|
||||
0.0, -0.3, 0.0, 1.0, # right shoulder and elbow
|
||||
0.0, 0.0, 0.0 # right wrist
|
||||
]
|
||||
|
||||
robot_dofs: {
|
||||
"g1_29dof": [1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1],
|
||||
"g1_29dof_anneal_23dof": [1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1,
|
||||
1, 1, 1, 1, 0, 0, 0,
|
||||
1, 1, 1, 1, 0, 0, 0],
|
||||
}
|
||||
|
||||
mimic_robot_types: {
|
||||
|
||||
"APT_level1": "g1_29dof_anneal_23dof",
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
# 01281657
|
||||
mimic_models: {
|
||||
"APT_level1": "20250116_225127-TairanTestbed_G129dofANNEAL23dof_dm_APT_video_APT_level1_MinimalFriction-0.3_RfiTrue_Far0.325_RESUME_LARGENOISE-motion_tracking-g1_29dof_anneal_23dof/exported/model_176500.onnx",
|
||||
|
||||
}
|
||||
|
||||
|
||||
|
||||
start_upper_body_dof_pos: {
|
||||
|
||||
"APT_level1":
|
||||
[0.19964170455932617, 0.07710712403059006, -0.2882401943206787,
|
||||
0.21672365069389343, 0.15629297494888306, -0.5167576670646667, 0.5782126784324646,
|
||||
0.0, 0.0, 0.0,
|
||||
0.25740593671798706, -0.2504104673862457, 0.22500675916671753, 0.5127624273300171,
|
||||
0.0, 0.0, 0.0],
|
||||
|
||||
}
|
||||
|
||||
motion_length_s: {
|
||||
"APT_level1": 7.66,
|
||||
|
||||
}
|
||||
@@ -1,13 +0,0 @@
|
||||
mujoco>=3.0.0
|
||||
numpy>=1.24.0
|
||||
pyyaml>=6.0
|
||||
unitree-sdk2py>=1.0.0
|
||||
loguru>=0.7.0
|
||||
|
||||
# Camera publishing dependencies
|
||||
opencv-python>=4.8.0
|
||||
pyzmq>=25.0.0
|
||||
msgpack>=1.0.0
|
||||
msgpack-numpy>=0.4.8
|
||||
matplotlib>=3.5.0 # For live camera viewer
|
||||
|
||||
@@ -1,77 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Standalone MuJoCo simulator for Unitree G1"""
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add sim module to path
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
import yaml
|
||||
from sim.simulator_factory import SimulatorFactory, init_channel
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Unitree G1 MuJoCo Simulator")
|
||||
parser.add_argument("--publish-images", action="store_true",
|
||||
help="Enable camera image publishing via ZMQ (requires offscreen rendering)")
|
||||
parser.add_argument("--camera-port", type=int, default=5555,
|
||||
help="ZMQ port for camera publishing (default: 5555)")
|
||||
parser.add_argument("--cameras", type=str, nargs="+", default=None,
|
||||
help="Camera names to publish (default: head_camera)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load config
|
||||
config_path = Path(__file__).parent / "config.yaml"
|
||||
with open(config_path) as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
# Override config with command line args
|
||||
enable_offscreen = args.publish_images or config.get("ENABLE_OFFSCREEN", False)
|
||||
|
||||
print("="*60)
|
||||
print("🤖 Starting Unitree G1 MuJoCo Simulator")
|
||||
print("="*60)
|
||||
print(f"📁 Scene: {config['ROBOT_SCENE']}")
|
||||
print(f"⏱️ Timestep: {config['SIMULATE_DT']}s ({int(1/config['SIMULATE_DT'])} Hz)")
|
||||
print(f"👁️ Visualization: {'ON' if config.get('ENABLE_ONSCREEN', True) else 'OFF'}")
|
||||
|
||||
# Configure cameras if requested
|
||||
camera_configs = {}
|
||||
if enable_offscreen:
|
||||
camera_list = args.cameras or ["head_camera"]
|
||||
for cam_name in camera_list:
|
||||
camera_configs[cam_name] = {"height": 480, "width": 640}
|
||||
print(f"📷 Cameras: {', '.join(camera_list)} → ZMQ port {args.camera_port}")
|
||||
|
||||
print("="*60)
|
||||
|
||||
# Initialize DDS channel
|
||||
init_channel(config=config)
|
||||
|
||||
# Create simulator
|
||||
sim = SimulatorFactory.create_simulator(
|
||||
config=config,
|
||||
env_name="default",
|
||||
onscreen=config.get("ENABLE_ONSCREEN", True),
|
||||
offscreen=enable_offscreen,
|
||||
camera_configs=camera_configs,
|
||||
)
|
||||
|
||||
# Start simulator (blocking)
|
||||
print("\nSimulator running. Press Ctrl+C to exit.")
|
||||
if enable_offscreen and args.publish_images:
|
||||
print(f"Camera images publishing on tcp://localhost:{args.camera_port}")
|
||||
try:
|
||||
SimulatorFactory.start_simulator(
|
||||
sim,
|
||||
as_thread=False,
|
||||
enable_image_publish=args.publish_images,
|
||||
camera_port=args.camera_port,
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down simulator...")
|
||||
sim.close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -1,803 +0,0 @@
|
||||
import argparse
|
||||
import pathlib
|
||||
from pathlib import Path
|
||||
import threading
|
||||
from threading import Lock, Thread
|
||||
from typing import Dict
|
||||
|
||||
import mujoco
|
||||
import mujoco.viewer
|
||||
import numpy as np
|
||||
try:
|
||||
import rclpy
|
||||
HAS_RCLPY = True
|
||||
except ImportError:
|
||||
HAS_RCLPY = False
|
||||
print("Warning: rclpy not found. Camera image publishing will be disabled.")
|
||||
from unitree_sdk2py.core.channel import ChannelFactoryInitialize
|
||||
import yaml
|
||||
|
||||
from .image_publish_utils import ImagePublishProcess
|
||||
from .metric_utils import check_contact, check_height
|
||||
from .sim_utilts import get_subtree_body_names
|
||||
from .unitree_sdk2py_bridge import ElasticBand, UnitreeSdk2Bridge
|
||||
|
||||
GR00T_WBC_ROOT = Path(__file__).resolve().parent.parent # Points to mujoco_sim_g1/
|
||||
|
||||
|
||||
class DefaultEnv:
|
||||
"""Base environment class that handles simulation environment setup and step"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Dict[str, any],
|
||||
env_name: str = "default",
|
||||
camera_configs: Dict[str, any] = None,
|
||||
onscreen: bool = False,
|
||||
offscreen: bool = False,
|
||||
):
|
||||
# Avoid mutable default argument gotcha
|
||||
if camera_configs is None:
|
||||
camera_configs = {}
|
||||
|
||||
# global_view is only set up for this specifc scene for now.
|
||||
if config["ROBOT_SCENE"] == "gr00t_wbc/control/robot_model/model_data/g1/scene_29dof.xml":
|
||||
camera_configs["global_view"] = {
|
||||
"height": 400,
|
||||
"width": 400,
|
||||
}
|
||||
self.config = config
|
||||
self.env_name = env_name
|
||||
self.num_body_dof = self.config["NUM_JOINTS"]
|
||||
self.num_hand_dof = self.config["NUM_HAND_JOINTS"]
|
||||
self.sim_dt = self.config["SIMULATE_DT"]
|
||||
self.obs = None
|
||||
self.torques = np.zeros(self.num_body_dof + self.num_hand_dof * 2)
|
||||
self.torque_limit = np.array(self.config["motor_effort_limit_list"])
|
||||
self.camera_configs = camera_configs
|
||||
|
||||
# Debug: print camera config
|
||||
if len(camera_configs) > 0:
|
||||
print(f"✓ DefaultEnv initialized with {len(camera_configs)} camera(s): {list(camera_configs.keys())}")
|
||||
|
||||
# Thread safety lock
|
||||
self.reward_lock = Lock()
|
||||
|
||||
# Unitree bridge will be initialized by the simulator
|
||||
self.unitree_bridge = None
|
||||
|
||||
# Store display mode
|
||||
self.onscreen = onscreen
|
||||
|
||||
# Initialize scene (defined in subclasses)
|
||||
self.init_scene()
|
||||
self.last_reward = 0
|
||||
|
||||
# Setup offscreen rendering if needed
|
||||
self.offscreen = offscreen
|
||||
if self.offscreen:
|
||||
self.init_renderers()
|
||||
self.image_dt = self.config.get("IMAGE_DT", 0.033333)
|
||||
|
||||
# Image publishing subprocess (initialized separately)
|
||||
self.image_publish_process = None
|
||||
|
||||
def init_scene(self):
|
||||
"""Initialize the default robot scene"""
|
||||
self.mj_model = mujoco.MjModel.from_xml_path(
|
||||
str(pathlib.Path(GR00T_WBC_ROOT) / self.config["ROBOT_SCENE"])
|
||||
)
|
||||
self.mj_data = mujoco.MjData(self.mj_model)
|
||||
self.mj_model.opt.timestep = self.sim_dt
|
||||
self.torso_index = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, "torso_link")
|
||||
self.root_body = "pelvis"
|
||||
# Enable the elastic band
|
||||
if self.config["ENABLE_ELASTIC_BAND"]:
|
||||
self.elastic_band = ElasticBand()
|
||||
if "g1" in self.config["ROBOT_TYPE"]:
|
||||
if self.config["enable_waist"]:
|
||||
self.band_attached_link = self.mj_model.body("pelvis").id
|
||||
else:
|
||||
self.band_attached_link = self.mj_model.body("torso_link").id
|
||||
elif "h1" in self.config["ROBOT_TYPE"]:
|
||||
self.band_attached_link = self.mj_model.body("torso_link").id
|
||||
else:
|
||||
self.band_attached_link = self.mj_model.body("base_link").id
|
||||
|
||||
if self.onscreen:
|
||||
self.viewer = mujoco.viewer.launch_passive(
|
||||
self.mj_model,
|
||||
self.mj_data,
|
||||
key_callback=self.elastic_band.MujuocoKeyCallback,
|
||||
show_left_ui=False,
|
||||
show_right_ui=False,
|
||||
)
|
||||
else:
|
||||
mujoco.mj_forward(self.mj_model, self.mj_data)
|
||||
self.viewer = None
|
||||
else:
|
||||
if self.onscreen:
|
||||
self.viewer = mujoco.viewer.launch_passive(
|
||||
self.mj_model, self.mj_data, show_left_ui=False, show_right_ui=False
|
||||
)
|
||||
else:
|
||||
mujoco.mj_forward(self.mj_model, self.mj_data)
|
||||
self.viewer = None
|
||||
|
||||
if self.viewer:
|
||||
# viewer camera
|
||||
self.viewer.cam.azimuth = 120 # Horizontal rotation in degrees
|
||||
self.viewer.cam.elevation = -30 # Vertical tilt in degrees
|
||||
self.viewer.cam.distance = 2.0 # Distance from camera to target
|
||||
self.viewer.cam.lookat = np.array([0, 0, 0.5]) # Point the camera is looking at
|
||||
|
||||
# Note that the actuator order is the same as the joint order in the mujoco model.
|
||||
self.body_joint_index = []
|
||||
self.left_hand_index = []
|
||||
self.right_hand_index = []
|
||||
for i in range(self.mj_model.njnt):
|
||||
name = self.mj_model.joint(i).name
|
||||
if any(
|
||||
[
|
||||
part_name in name
|
||||
for part_name in ["hip", "knee", "ankle", "waist", "shoulder", "elbow", "wrist"]
|
||||
]
|
||||
):
|
||||
self.body_joint_index.append(i)
|
||||
elif "left_hand" in name:
|
||||
self.left_hand_index.append(i)
|
||||
elif "right_hand" in name:
|
||||
self.right_hand_index.append(i)
|
||||
|
||||
assert len(self.body_joint_index) == self.config["NUM_JOINTS"], \
|
||||
f"Expected {self.config['NUM_JOINTS']} body joints, got {len(self.body_joint_index)}"
|
||||
# Hand joints are optional (some models don't have hands)
|
||||
if self.config.get("NUM_HAND_JOINTS", 0) > 0:
|
||||
expected_hands = self.config["NUM_HAND_JOINTS"]
|
||||
if len(self.left_hand_index) != expected_hands or len(self.right_hand_index) != expected_hands:
|
||||
print(f"Warning: Expected {expected_hands} hand joints, got left={len(self.left_hand_index)}, right={len(self.right_hand_index)}")
|
||||
print("Continuing without hands...")
|
||||
|
||||
self.body_joint_index = np.array(self.body_joint_index)
|
||||
self.left_hand_index = np.array(self.left_hand_index)
|
||||
self.right_hand_index = np.array(self.right_hand_index)
|
||||
|
||||
def init_renderers(self):
|
||||
# Initialize camera renderers
|
||||
self.renderers = {}
|
||||
for camera_name, camera_config in self.camera_configs.items():
|
||||
renderer = mujoco.Renderer(
|
||||
self.mj_model, height=camera_config["height"], width=camera_config["width"]
|
||||
)
|
||||
self.renderers[camera_name] = renderer
|
||||
|
||||
def start_image_publish_subprocess(self, start_method: str = "spawn", camera_port: int = 5555):
|
||||
"""Start image publishing subprocess using ZMQ"""
|
||||
# Use spawn method for better GIL isolation, or configured method
|
||||
if len(self.camera_configs) == 0:
|
||||
print(
|
||||
"Warning: No camera configs provided, image publishing subprocess will not be started"
|
||||
)
|
||||
return
|
||||
start_method = self.config.get("MP_START_METHOD", "spawn")
|
||||
self.image_publish_process = ImagePublishProcess(
|
||||
camera_configs=self.camera_configs,
|
||||
image_dt=self.image_dt,
|
||||
zmq_port=camera_port,
|
||||
start_method=start_method,
|
||||
verbose=self.config.get("verbose", False),
|
||||
)
|
||||
self.image_publish_process.start_process()
|
||||
print(f"✓ Started image publishing subprocess on ZMQ port {camera_port}")
|
||||
|
||||
def compute_body_torques(self) -> np.ndarray:
|
||||
"""Compute body torques based on the current robot state"""
|
||||
body_torques = np.zeros(self.num_body_dof)
|
||||
if self.unitree_bridge is not None and self.unitree_bridge.low_cmd:
|
||||
for i in range(self.unitree_bridge.num_body_motor):
|
||||
if self.unitree_bridge.use_sensor:
|
||||
body_torques[i] = (
|
||||
self.unitree_bridge.low_cmd.motor_cmd[i].tau
|
||||
+ self.unitree_bridge.low_cmd.motor_cmd[i].kp
|
||||
* (self.unitree_bridge.low_cmd.motor_cmd[i].q - self.mj_data.sensordata[i])
|
||||
+ self.unitree_bridge.low_cmd.motor_cmd[i].kd
|
||||
* (
|
||||
self.unitree_bridge.low_cmd.motor_cmd[i].dq
|
||||
- self.mj_data.sensordata[i + self.unitree_bridge.num_body_motor]
|
||||
)
|
||||
)
|
||||
else:
|
||||
body_torques[i] = (
|
||||
self.unitree_bridge.low_cmd.motor_cmd[i].tau
|
||||
+ self.unitree_bridge.low_cmd.motor_cmd[i].kp
|
||||
* (
|
||||
self.unitree_bridge.low_cmd.motor_cmd[i].q
|
||||
- self.mj_data.qpos[self.body_joint_index[i] + 7 - 1]
|
||||
)
|
||||
+ self.unitree_bridge.low_cmd.motor_cmd[i].kd
|
||||
* (
|
||||
self.unitree_bridge.low_cmd.motor_cmd[i].dq
|
||||
- self.mj_data.qvel[self.body_joint_index[i] + 6 - 1]
|
||||
)
|
||||
)
|
||||
return body_torques
|
||||
|
||||
def compute_hand_torques(self) -> np.ndarray:
|
||||
"""Compute hand torques based on the current robot state"""
|
||||
left_hand_torques = np.zeros(self.num_hand_dof)
|
||||
right_hand_torques = np.zeros(self.num_hand_dof)
|
||||
if self.unitree_bridge is not None and self.unitree_bridge.low_cmd:
|
||||
for i in range(self.unitree_bridge.num_hand_motor):
|
||||
left_hand_torques[i] = (
|
||||
self.unitree_bridge.left_hand_cmd.motor_cmd[i].tau
|
||||
+ self.unitree_bridge.left_hand_cmd.motor_cmd[i].kp
|
||||
* (
|
||||
self.unitree_bridge.left_hand_cmd.motor_cmd[i].q
|
||||
- self.mj_data.qpos[self.left_hand_index[i] + 7 - 1]
|
||||
)
|
||||
+ self.unitree_bridge.left_hand_cmd.motor_cmd[i].kd
|
||||
* (
|
||||
self.unitree_bridge.left_hand_cmd.motor_cmd[i].dq
|
||||
- self.mj_data.qvel[self.left_hand_index[i] + 6 - 1]
|
||||
)
|
||||
)
|
||||
right_hand_torques[i] = (
|
||||
self.unitree_bridge.right_hand_cmd.motor_cmd[i].tau
|
||||
+ self.unitree_bridge.right_hand_cmd.motor_cmd[i].kp
|
||||
* (
|
||||
self.unitree_bridge.right_hand_cmd.motor_cmd[i].q
|
||||
- self.mj_data.qpos[self.right_hand_index[i] + 7 - 1]
|
||||
)
|
||||
+ self.unitree_bridge.right_hand_cmd.motor_cmd[i].kd
|
||||
* (
|
||||
self.unitree_bridge.right_hand_cmd.motor_cmd[i].dq
|
||||
- self.mj_data.qvel[self.right_hand_index[i] + 6 - 1]
|
||||
)
|
||||
)
|
||||
return np.concatenate((left_hand_torques, right_hand_torques))
|
||||
|
||||
def compute_body_qpos(self) -> np.ndarray:
|
||||
"""Compute body joint positions based on the current command"""
|
||||
body_qpos = np.zeros(self.num_body_dof)
|
||||
if self.unitree_bridge is not None and self.unitree_bridge.low_cmd:
|
||||
for i in range(self.unitree_bridge.num_body_motor):
|
||||
body_qpos[i] = self.unitree_bridge.low_cmd.motor_cmd[i].q
|
||||
return body_qpos
|
||||
|
||||
def compute_hand_qpos(self) -> np.ndarray:
|
||||
"""Compute hand joint positions based on the current command"""
|
||||
hand_qpos = np.zeros(self.num_hand_dof * 2)
|
||||
if self.unitree_bridge is not None and self.unitree_bridge.low_cmd:
|
||||
for i in range(self.unitree_bridge.num_hand_motor):
|
||||
hand_qpos[i] = self.unitree_bridge.left_hand_cmd.motor_cmd[i].q
|
||||
hand_qpos[i + self.num_hand_dof] = self.unitree_bridge.right_hand_cmd.motor_cmd[i].q
|
||||
return hand_qpos
|
||||
|
||||
def prepare_obs(self) -> Dict[str, any]:
|
||||
"""Prepare observation dictionary from the current robot state"""
|
||||
obs = {}
|
||||
obs["floating_base_pose"] = self.mj_data.qpos[:7]
|
||||
obs["floating_base_vel"] = self.mj_data.qvel[:6]
|
||||
obs["floating_base_acc"] = self.mj_data.qacc[:6]
|
||||
obs["secondary_imu_quat"] = self.mj_data.xquat[self.torso_index]
|
||||
obs["secondary_imu_vel"] = self.mj_data.cvel[self.torso_index]
|
||||
obs["body_q"] = self.mj_data.qpos[self.body_joint_index + 7 - 1]
|
||||
obs["body_dq"] = self.mj_data.qvel[self.body_joint_index + 6 - 1]
|
||||
obs["body_ddq"] = self.mj_data.qacc[self.body_joint_index + 6 - 1]
|
||||
obs["body_tau_est"] = self.mj_data.actuator_force[self.body_joint_index - 1]
|
||||
if self.num_hand_dof > 0:
|
||||
obs["left_hand_q"] = self.mj_data.qpos[self.left_hand_index + 7 - 1]
|
||||
obs["left_hand_dq"] = self.mj_data.qvel[self.left_hand_index + 6 - 1]
|
||||
obs["left_hand_ddq"] = self.mj_data.qacc[self.left_hand_index + 6 - 1]
|
||||
obs["left_hand_tau_est"] = self.mj_data.actuator_force[self.left_hand_index - 1]
|
||||
obs["right_hand_q"] = self.mj_data.qpos[self.right_hand_index + 7 - 1]
|
||||
obs["right_hand_dq"] = self.mj_data.qvel[self.right_hand_index + 6 - 1]
|
||||
obs["right_hand_ddq"] = self.mj_data.qacc[self.right_hand_index + 6 - 1]
|
||||
obs["right_hand_tau_est"] = self.mj_data.actuator_force[self.right_hand_index - 1]
|
||||
obs["time"] = self.mj_data.time
|
||||
return obs
|
||||
|
||||
def sim_step(self):
|
||||
self.obs = self.prepare_obs()
|
||||
self.unitree_bridge.PublishLowState(self.obs)
|
||||
if self.unitree_bridge.joystick:
|
||||
self.unitree_bridge.PublishWirelessController()
|
||||
if self.config["ENABLE_ELASTIC_BAND"]:
|
||||
if self.elastic_band.enable:
|
||||
# Get Cartesian pose and velocity of the band_attached_link
|
||||
pose = np.concatenate(
|
||||
[
|
||||
self.mj_data.xpos[self.band_attached_link], # link position in world
|
||||
self.mj_data.xquat[
|
||||
self.band_attached_link
|
||||
], # link quaternion in world [w,x,y,z]
|
||||
np.zeros(6), # placeholder for velocity
|
||||
]
|
||||
)
|
||||
|
||||
# Get velocity in world frame
|
||||
mujoco.mj_objectVelocity(
|
||||
self.mj_model,
|
||||
self.mj_data,
|
||||
mujoco.mjtObj.mjOBJ_BODY,
|
||||
self.band_attached_link,
|
||||
pose[7:13],
|
||||
0, # 0 for world frame
|
||||
)
|
||||
|
||||
# Reorder velocity from [ang, lin] to [lin, ang]
|
||||
pose[7:10], pose[10:13] = pose[10:13], pose[7:10].copy()
|
||||
self.mj_data.xfrc_applied[self.band_attached_link] = self.elastic_band.Advance(pose)
|
||||
else:
|
||||
# explicitly resetting the force when the band is not enabled
|
||||
self.mj_data.xfrc_applied[self.band_attached_link] = np.zeros(6)
|
||||
body_torques = self.compute_body_torques()
|
||||
hand_torques = self.compute_hand_torques()
|
||||
self.torques[self.body_joint_index - 1] = body_torques
|
||||
if self.num_hand_dof > 0:
|
||||
self.torques[self.left_hand_index - 1] = hand_torques[: self.num_hand_dof]
|
||||
self.torques[self.right_hand_index - 1] = hand_torques[self.num_hand_dof :]
|
||||
|
||||
self.torques = np.clip(self.torques, -self.torque_limit, self.torque_limit)
|
||||
|
||||
if self.config["FREE_BASE"]:
|
||||
self.mj_data.ctrl = np.concatenate((np.zeros(6), self.torques))
|
||||
else:
|
||||
self.mj_data.ctrl = self.torques
|
||||
mujoco.mj_step(self.mj_model, self.mj_data)
|
||||
# self.check_self_collision()
|
||||
|
||||
def kinematics_step(self):
|
||||
"""
|
||||
Run kinematics only: compute the qpos of the robot and directly set the qpos.
|
||||
For debugging purposes.
|
||||
"""
|
||||
if self.unitree_bridge is not None:
|
||||
self.unitree_bridge.PublishLowState(self.prepare_obs())
|
||||
if self.unitree_bridge.joystick:
|
||||
self.unitree_bridge.PublishWirelessController()
|
||||
|
||||
if self.config["ENABLE_ELASTIC_BAND"]:
|
||||
if self.elastic_band.enable:
|
||||
# Get Cartesian pose and velocity of the band_attached_link
|
||||
pose = np.concatenate(
|
||||
[
|
||||
self.mj_data.xpos[self.band_attached_link], # link position in world
|
||||
self.mj_data.xquat[
|
||||
self.band_attached_link
|
||||
], # link quaternion in world [w,x,y,z]
|
||||
np.zeros(6), # placeholder for velocity
|
||||
]
|
||||
)
|
||||
|
||||
# Get velocity in world frame
|
||||
mujoco.mj_objectVelocity(
|
||||
self.mj_model,
|
||||
self.mj_data,
|
||||
mujoco.mjtObj.mjOBJ_BODY,
|
||||
self.band_attached_link,
|
||||
pose[7:13],
|
||||
0, # 0 for world frame
|
||||
)
|
||||
|
||||
# Reorder velocity from [ang, lin] to [lin, ang]
|
||||
pose[7:10], pose[10:13] = pose[10:13], pose[7:10].copy()
|
||||
|
||||
self.mj_data.xfrc_applied[self.band_attached_link] = self.elastic_band.Advance(pose)
|
||||
else:
|
||||
# explicitly resetting the force when the band is not enabled
|
||||
self.mj_data.xfrc_applied[self.band_attached_link] = np.zeros(6)
|
||||
|
||||
body_qpos = self.compute_body_qpos() # (num_body_dof,)
|
||||
hand_qpos = self.compute_hand_qpos() # (num_hand_dof * 2,)
|
||||
|
||||
self.mj_data.qpos[self.body_joint_index + 7 - 1] = body_qpos
|
||||
self.mj_data.qpos[self.left_hand_index + 7 - 1] = hand_qpos[: self.num_hand_dof]
|
||||
self.mj_data.qpos[self.right_hand_index + 7 - 1] = hand_qpos[self.num_hand_dof :]
|
||||
|
||||
mujoco.mj_kinematics(self.mj_model, self.mj_data)
|
||||
mujoco.mj_comPos(self.mj_model, self.mj_data)
|
||||
|
||||
def apply_perturbation(self, key):
|
||||
"""Apply perturbation to the robot"""
|
||||
# Add velocity perturbations in body frame
|
||||
perturbation_x_body = 0.0 # forward/backward in body frame
|
||||
perturbation_y_body = 0.0 # left/right in body frame
|
||||
if key == "up":
|
||||
perturbation_x_body = 1.0 # forward
|
||||
elif key == "down":
|
||||
perturbation_x_body = -1.0 # backward
|
||||
elif key == "left":
|
||||
perturbation_y_body = 1.0 # left
|
||||
elif key == "right":
|
||||
perturbation_y_body = -1.0 # right
|
||||
|
||||
# Transform body frame velocity to world frame using MuJoCo's rotation
|
||||
vel_body = np.array([perturbation_x_body, perturbation_y_body, 0.0])
|
||||
vel_world = np.zeros(3)
|
||||
base_quat = self.mj_data.qpos[3:7] # [w, x, y, z] quaternion
|
||||
|
||||
# Use MuJoCo's robust quaternion rotation (handles invalid quaternions automatically)
|
||||
mujoco.mju_rotVecQuat(vel_world, vel_body, base_quat)
|
||||
|
||||
# Apply to base linear velocity in world frame
|
||||
self.mj_data.qvel[0] += vel_world[0] # world X velocity
|
||||
self.mj_data.qvel[1] += vel_world[1] # world Y velocity
|
||||
|
||||
# Update dynamics after velocity change
|
||||
mujoco.mj_forward(self.mj_model, self.mj_data)
|
||||
|
||||
def update_viewer(self):
|
||||
if self.viewer is not None:
|
||||
self.viewer.sync()
|
||||
|
||||
def update_viewer_camera(self):
|
||||
if self.viewer is not None:
|
||||
if self.viewer.cam.type == mujoco.mjtCamera.mjCAMERA_TRACKING:
|
||||
self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
|
||||
else:
|
||||
self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
|
||||
|
||||
def update_reward(self):
|
||||
"""Calculate reward. Should be implemented by subclasses."""
|
||||
with self.reward_lock:
|
||||
self.last_reward = 0
|
||||
|
||||
def get_reward(self):
|
||||
"""Thread-safe way to get the last calculated reward."""
|
||||
with self.reward_lock:
|
||||
return self.last_reward
|
||||
|
||||
def set_unitree_bridge(self, unitree_bridge):
|
||||
"""Set the unitree bridge from the simulator"""
|
||||
self.unitree_bridge = unitree_bridge
|
||||
|
||||
def get_privileged_obs(self):
|
||||
"""Get privileged observation. Should be implemented by subclasses."""
|
||||
return {}
|
||||
|
||||
def update_render_caches(self):
|
||||
"""Update render cache and shared memory for subprocess."""
|
||||
render_caches = {}
|
||||
for camera_name, camera_config in self.camera_configs.items():
|
||||
renderer = self.renderers[camera_name]
|
||||
if "params" in camera_config:
|
||||
renderer.update_scene(self.mj_data, camera=camera_config["params"])
|
||||
else:
|
||||
renderer.update_scene(self.mj_data, camera=camera_name)
|
||||
render_caches[camera_name + "_image"] = renderer.render()
|
||||
|
||||
# Update shared memory if image publishing process is available
|
||||
if self.image_publish_process is not None:
|
||||
self.image_publish_process.update_shared_memory(render_caches)
|
||||
|
||||
return render_caches
|
||||
|
||||
def handle_keyboard_button(self, key):
|
||||
if self.elastic_band is not None:
|
||||
self.elastic_band.handle_keyboard_button(key)
|
||||
|
||||
if key == "backspace":
|
||||
self.reset()
|
||||
if key == "v":
|
||||
self.update_viewer_camera()
|
||||
if key in ["up", "down", "left", "right"]:
|
||||
self.apply_perturbation(key)
|
||||
|
||||
def check_fall(self):
|
||||
"""Check if the robot has fallen"""
|
||||
self.fall = False
|
||||
if self.mj_data.qpos[2] < 0.2:
|
||||
self.fall = True
|
||||
print(f"Warning: Robot has fallen, height: {self.mj_data.qpos[2]:.3f} m")
|
||||
|
||||
if self.fall:
|
||||
self.reset()
|
||||
|
||||
def check_self_collision(self):
|
||||
"""Check for self-collision of the robot"""
|
||||
robot_bodies = get_subtree_body_names(self.mj_model, self.mj_model.body(self.root_body).id)
|
||||
self_collision, contact_bodies = check_contact(
|
||||
self.mj_model, self.mj_data, robot_bodies, robot_bodies, return_all_contact_bodies=True
|
||||
)
|
||||
if self_collision:
|
||||
print(f"Warning: Self-collision detected: {contact_bodies}")
|
||||
return self_collision
|
||||
|
||||
def reset(self):
|
||||
mujoco.mj_resetData(self.mj_model, self.mj_data)
|
||||
|
||||
|
||||
class CubeEnv(DefaultEnv):
|
||||
"""Environment with a cube object for pick and place tasks"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Dict[str, any],
|
||||
onscreen: bool = False,
|
||||
offscreen: bool = False,
|
||||
):
|
||||
# Override the robot scene
|
||||
config = config.copy() # Create a copy to avoid modifying the original
|
||||
config["ROBOT_SCENE"] = "gr00t_wbc/control/robot_model/model_data/g1/pnp_cube_43dof.xml"
|
||||
super().__init__(config, "cube", {}, onscreen, offscreen)
|
||||
|
||||
def update_reward(self):
|
||||
"""Calculate reward based on gripper contact with cube and cube height"""
|
||||
right_hand_body = [
|
||||
"right_hand_thumb_2_link",
|
||||
"right_hand_middle_1_link",
|
||||
"right_hand_index_1_link",
|
||||
]
|
||||
gripper_cube_contact = check_contact(
|
||||
self.mj_model, self.mj_data, right_hand_body, "cube_body"
|
||||
)
|
||||
cube_lifted = check_height(self.mj_model, self.mj_data, "cube", 0.85, 2.0)
|
||||
|
||||
with self.reward_lock:
|
||||
self.last_reward = gripper_cube_contact & cube_lifted
|
||||
|
||||
|
||||
class BoxEnv(DefaultEnv):
|
||||
"""Environment with a box object for manipulation tasks"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Dict[str, any],
|
||||
onscreen: bool = False,
|
||||
offscreen: bool = False,
|
||||
):
|
||||
# Override the robot scene
|
||||
config = config.copy() # Create a copy to avoid modifying the original
|
||||
config["ROBOT_SCENE"] = "gr00t_wbc/control/robot_model/model_data/g1/lift_box_43dof.xml"
|
||||
super().__init__(config, "box", {}, onscreen, offscreen)
|
||||
|
||||
def reward(self):
|
||||
"""Calculate reward based on gripper contact with cube and cube height"""
|
||||
left_hand_body = [
|
||||
"left_hand_thumb_2_link",
|
||||
"left_hand_middle_1_link",
|
||||
"left_hand_index_1_link",
|
||||
]
|
||||
right_hand_body = [
|
||||
"right_hand_thumb_2_link",
|
||||
"right_hand_middle_1_link",
|
||||
"right_hand_index_1_link",
|
||||
]
|
||||
gripper_box_contact = check_contact(self.mj_model, self.mj_data, left_hand_body, "box_body")
|
||||
gripper_box_contact &= check_contact(
|
||||
self.mj_model, self.mj_data, right_hand_body, "box_body"
|
||||
)
|
||||
box_lifted = check_height(self.mj_model, self.mj_data, "box", 0.92, 2.0)
|
||||
|
||||
print("gripper_box_contact: ", gripper_box_contact, "box_lifted: ", box_lifted)
|
||||
|
||||
with self.reward_lock:
|
||||
self.last_reward = gripper_box_contact & box_lifted
|
||||
return self.last_reward
|
||||
|
||||
|
||||
class BottleEnv(DefaultEnv):
|
||||
"""Environment with a cylinder object for manipulation tasks"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Dict[str, any],
|
||||
onscreen: bool = False,
|
||||
offscreen: bool = False,
|
||||
):
|
||||
# Override the robot scene
|
||||
config = config.copy() # Create a copy to avoid modifying the original
|
||||
config["ROBOT_SCENE"] = "gr00t_wbc/control/robot_model/model_data/g1/pnp_bottle_43dof.xml"
|
||||
camera_configs = {
|
||||
"egoview": {
|
||||
"height": 400,
|
||||
"width": 400,
|
||||
},
|
||||
}
|
||||
super().__init__(
|
||||
config, "cylinder", camera_configs, onscreen, offscreen
|
||||
)
|
||||
|
||||
self.bottle_body = self.mj_model.body("bottle_body")
|
||||
self.bottle_geom = self.mj_model.geom("bottle")
|
||||
|
||||
if self.viewer is not None:
|
||||
self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
|
||||
self.viewer.cam.fixedcamid = self.mj_model.camera("egoview").id
|
||||
|
||||
def update_reward(self):
|
||||
"""Calculate reward based on gripper contact with cylinder and cylinder height"""
|
||||
pass
|
||||
|
||||
def get_privileged_obs(self):
|
||||
obs_pos = self.mj_data.xpos[self.bottle_body.id]
|
||||
obs_quat = self.mj_data.xquat[self.bottle_body.id]
|
||||
return {"bottle_pos": obs_pos, "bottle_quat": obs_quat}
|
||||
|
||||
|
||||
class BaseSimulator:
|
||||
"""Base simulator class that handles initialization and running of simulations"""
|
||||
|
||||
def __init__(self, config: Dict[str, any], env_name: str = "default", **kwargs):
|
||||
self.config = config
|
||||
self.env_name = env_name
|
||||
|
||||
# Initialize ROS 2 node (optional, only if rclpy is available)
|
||||
if HAS_RCLPY:
|
||||
if not rclpy.ok():
|
||||
rclpy.init()
|
||||
self.node = rclpy.create_node("sim_mujoco")
|
||||
self.thread = threading.Thread(target=rclpy.spin, args=(self.node,), daemon=True)
|
||||
self.thread.start()
|
||||
else:
|
||||
self.thread = None
|
||||
executor = rclpy.get_global_executor()
|
||||
self.node = executor.get_nodes()[0] # will only take the first node
|
||||
else:
|
||||
self.node = None
|
||||
self.thread = None
|
||||
|
||||
# Set update frequencies
|
||||
self.sim_dt = self.config["SIMULATE_DT"]
|
||||
self.reward_dt = self.config.get("REWARD_DT", 0.02)
|
||||
self.image_dt = self.config.get("IMAGE_DT", 0.033333)
|
||||
self.viewer_dt = self.config.get("VIEWER_DT", 0.02)
|
||||
|
||||
# Create the appropriate environment based on name
|
||||
if env_name == "default":
|
||||
self.sim_env = DefaultEnv(config, env_name, **kwargs)
|
||||
elif env_name == "pnp_cube":
|
||||
self.sim_env = CubeEnv(config, **kwargs)
|
||||
elif env_name == "lift_box":
|
||||
self.sim_env = BoxEnv(config, **kwargs)
|
||||
elif env_name == "pnp_bottle":
|
||||
self.sim_env = BottleEnv(config, **kwargs)
|
||||
else:
|
||||
raise ValueError(f"Invalid environment name: {env_name}")
|
||||
|
||||
# Initialize the DDS communication layer - should be safe to call multiple times
|
||||
|
||||
try:
|
||||
if self.config.get("INTERFACE", None):
|
||||
ChannelFactoryInitialize(self.config["DOMAIN_ID"], self.config["INTERFACE"])
|
||||
else:
|
||||
ChannelFactoryInitialize(self.config["DOMAIN_ID"])
|
||||
except Exception as e:
|
||||
# If it fails because it's already initialized, that's okay
|
||||
print(f"Note: Channel factory initialization attempt: {e}")
|
||||
|
||||
# Initialize the unitree bridge and pass it to the environment
|
||||
self.init_unitree_bridge()
|
||||
self.sim_env.set_unitree_bridge(self.unitree_bridge)
|
||||
|
||||
# Initialize additional components
|
||||
self.init_subscriber()
|
||||
self.init_publisher()
|
||||
|
||||
self.sim_thread = None
|
||||
|
||||
def start_as_thread(self):
|
||||
# Create simulation thread
|
||||
self.sim_thread = Thread(target=self.start)
|
||||
self.sim_thread.start()
|
||||
|
||||
def start_image_publish_subprocess(self, start_method: str = "spawn", camera_port: int = 5555):
|
||||
"""Start the image publish subprocess"""
|
||||
self.sim_env.start_image_publish_subprocess(start_method, camera_port)
|
||||
|
||||
def init_subscriber(self):
|
||||
"""Initialize subscribers. Can be overridden by subclasses."""
|
||||
pass
|
||||
|
||||
def init_publisher(self):
|
||||
"""Initialize publishers. Can be overridden by subclasses."""
|
||||
pass
|
||||
|
||||
def init_unitree_bridge(self):
|
||||
"""Initialize the unitree SDK bridge"""
|
||||
self.unitree_bridge = UnitreeSdk2Bridge(self.config)
|
||||
if self.config["USE_JOYSTICK"]:
|
||||
self.unitree_bridge.SetupJoystick(
|
||||
device_id=self.config["JOYSTICK_DEVICE"], js_type=self.config["JOYSTICK_TYPE"]
|
||||
)
|
||||
|
||||
def start(self):
|
||||
"""Main simulation loop"""
|
||||
import time
|
||||
sim_cnt = 0
|
||||
last_time = time.time()
|
||||
|
||||
print(f"Starting simulation loop. Viewer: {self.sim_env.viewer is not None}")
|
||||
|
||||
try:
|
||||
while (
|
||||
self.sim_env.viewer and self.sim_env.viewer.is_running()
|
||||
) or self.sim_env.viewer is None:
|
||||
# Run simulation step
|
||||
self.sim_env.sim_step()
|
||||
|
||||
# Update viewer at viewer rate
|
||||
if sim_cnt % int(self.viewer_dt / self.sim_dt) == 0:
|
||||
self.sim_env.update_viewer()
|
||||
|
||||
# Calculate reward at reward rate
|
||||
if sim_cnt % int(self.reward_dt / self.sim_dt) == 0:
|
||||
self.sim_env.update_reward()
|
||||
|
||||
# Update render caches at image rate
|
||||
if sim_cnt % int(self.image_dt / self.sim_dt) == 0:
|
||||
self.sim_env.update_render_caches()
|
||||
|
||||
# Sleep to maintain correct rate (simple timing without ROS)
|
||||
elapsed = time.time() - last_time
|
||||
sleep_time = max(0, self.sim_dt - elapsed)
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
last_time = time.time()
|
||||
|
||||
sim_cnt += 1
|
||||
|
||||
print(f"Loop exited. Viewer running: {self.sim_env.viewer.is_running() if self.sim_env.viewer else 'No viewer'}")
|
||||
except KeyboardInterrupt:
|
||||
# User pressed Ctrl+C - exit cleanly
|
||||
print("Keyboard interrupt received")
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"Exception in simulation loop: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
self.close()
|
||||
|
||||
def __del__(self):
|
||||
"""Clean up resources when simulator is deleted"""
|
||||
self.close()
|
||||
|
||||
def reset(self):
|
||||
"""Reset the simulation. Can be overridden by subclasses."""
|
||||
self.sim_env.reset()
|
||||
|
||||
def close(self):
|
||||
"""Close the simulation. Can be overridden by subclasses."""
|
||||
try:
|
||||
# Close viewer
|
||||
if hasattr(self.sim_env, "viewer") and self.sim_env.viewer is not None:
|
||||
self.sim_env.viewer.close()
|
||||
|
||||
# Shutdown ROS (if available)
|
||||
if HAS_RCLPY and rclpy.ok():
|
||||
rclpy.shutdown()
|
||||
except Exception as e:
|
||||
print(f"Warning during close: {e}")
|
||||
|
||||
def get_privileged_obs(self):
|
||||
obs = self.sim_env.get_privileged_obs()
|
||||
# TODO: add ros2 topic to get privileged obs
|
||||
return obs
|
||||
|
||||
def handle_keyboard_button(self, key):
|
||||
# Only handles keyboard buttons for default env.
|
||||
if self.env_name == "default":
|
||||
self.sim_env.handle_keyboard_button(key)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Robot")
|
||||
parser.add_argument(
|
||||
"--config",
|
||||
type=str,
|
||||
default="./gr00t_wbc/control/main/teleop/configs/g1_29dof_gear_wbc.yaml",
|
||||
help="config file",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
with open(args.config, "r") as file:
|
||||
config = yaml.load(file, Loader=yaml.FullLoader)
|
||||
|
||||
if config.get("INTERFACE", None):
|
||||
ChannelFactoryInitialize(config["DOMAIN_ID"], config["INTERFACE"])
|
||||
else:
|
||||
ChannelFactoryInitialize(config["DOMAIN_ID"])
|
||||
|
||||
simulation = BaseSimulator(config)
|
||||
simulation.start_as_thread()
|
||||
@@ -1,257 +0,0 @@
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import shared_memory
|
||||
import time
|
||||
from typing import Any, Dict
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .sensor_utils import ImageMessageSchema, ImageUtils, SensorServer
|
||||
|
||||
|
||||
def get_multiprocessing_info(verbose: bool = True):
|
||||
"""Get information about multiprocessing start methods"""
|
||||
|
||||
if verbose:
|
||||
print(f"Available start methods: {mp.get_all_start_methods()}")
|
||||
return mp.get_start_method()
|
||||
|
||||
|
||||
class ImagePublishProcess:
|
||||
"""Subprocess for publishing images using shared memory and ZMQ"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
camera_configs: Dict[str, Any],
|
||||
image_dt: float,
|
||||
zmq_port: int = 5555,
|
||||
start_method: str = "spawn",
|
||||
verbose: bool = False,
|
||||
):
|
||||
self.camera_configs = camera_configs
|
||||
self.image_dt = image_dt
|
||||
self.zmq_port = zmq_port
|
||||
self.verbose = verbose
|
||||
self.shared_memory_blocks = {}
|
||||
self.shared_memory_info = {}
|
||||
self.process = None
|
||||
|
||||
# Use specific context to avoid global state pollution
|
||||
self.mp_context = mp.get_context(start_method)
|
||||
if self.verbose:
|
||||
print(f"Using multiprocessing context: {start_method}")
|
||||
|
||||
self.stop_event = self.mp_context.Event()
|
||||
self.data_ready_event = self.mp_context.Event()
|
||||
|
||||
# Ensure events start in correct state
|
||||
self.stop_event.clear()
|
||||
self.data_ready_event.clear()
|
||||
|
||||
if self.verbose:
|
||||
print(f"Initial stop_event state: {self.stop_event.is_set()}")
|
||||
print(f"Initial data_ready_event state: {self.data_ready_event.is_set()}")
|
||||
|
||||
# Calculate shared memory requirements for each camera
|
||||
for camera_name, camera_config in camera_configs.items():
|
||||
height = camera_config["height"]
|
||||
width = camera_config["width"]
|
||||
# RGB image: height * width * 3 (uint8)
|
||||
size = height * width * 3
|
||||
|
||||
# Create shared memory block
|
||||
shm = shared_memory.SharedMemory(create=True, size=size)
|
||||
self.shared_memory_blocks[camera_name] = shm
|
||||
self.shared_memory_info[camera_name] = {
|
||||
"name": shm.name,
|
||||
"size": size,
|
||||
"shape": (height, width, 3),
|
||||
"dtype": np.uint8,
|
||||
}
|
||||
|
||||
def start_process(self):
|
||||
"""Start the image publishing subprocess"""
|
||||
if self.verbose:
|
||||
print(f"Starting subprocess with stop_event state: {self.stop_event.is_set()}")
|
||||
self.process = self.mp_context.Process(
|
||||
target=self._image_publish_worker,
|
||||
args=(
|
||||
self.shared_memory_info,
|
||||
self.image_dt,
|
||||
self.zmq_port,
|
||||
self.stop_event,
|
||||
self.data_ready_event,
|
||||
self.verbose,
|
||||
),
|
||||
)
|
||||
self.process.start()
|
||||
if self.verbose:
|
||||
print(f"Subprocess started, PID: {self.process.pid}")
|
||||
|
||||
def update_shared_memory(self, render_caches: Dict[str, np.ndarray]):
|
||||
"""Update shared memory with new rendered images"""
|
||||
images_updated = 0
|
||||
for camera_name in self.camera_configs.keys():
|
||||
image_key = f"{camera_name}_image"
|
||||
if image_key in render_caches:
|
||||
image = render_caches[image_key]
|
||||
|
||||
# Ensure image is uint8 and has correct shape
|
||||
if image.dtype != np.uint8:
|
||||
image = (image * 255).astype(np.uint8)
|
||||
|
||||
# Get shared memory array
|
||||
shm = self.shared_memory_blocks[camera_name]
|
||||
shared_array = np.ndarray(
|
||||
self.shared_memory_info[camera_name]["shape"],
|
||||
dtype=self.shared_memory_info[camera_name]["dtype"],
|
||||
buffer=shm.buf,
|
||||
)
|
||||
|
||||
# Copy image data to shared memory atomically
|
||||
np.copyto(shared_array, image)
|
||||
images_updated += 1
|
||||
|
||||
# Signal that new data is ready only after all images are written
|
||||
if images_updated > 0:
|
||||
if self.verbose:
|
||||
print(f"Main process: Updated {images_updated} images, setting data_ready_event")
|
||||
self.data_ready_event.set()
|
||||
elif self.verbose:
|
||||
print(
|
||||
"Main process: No images to update. "
|
||||
"please check if camera configs are provided and the renderer is properly initialized"
|
||||
)
|
||||
|
||||
def stop(self):
|
||||
"""Stop the image publishing subprocess"""
|
||||
if self.verbose:
|
||||
print("Stopping image publishing subprocess...")
|
||||
self.stop_event.set()
|
||||
|
||||
if self.process and self.process.is_alive():
|
||||
# Give the process time to clean up gracefully
|
||||
self.process.join(timeout=5)
|
||||
if self.process.is_alive():
|
||||
if self.verbose:
|
||||
print("Subprocess didn't stop gracefully, terminating...")
|
||||
self.process.terminate()
|
||||
self.process.join(timeout=2)
|
||||
if self.process.is_alive():
|
||||
if self.verbose:
|
||||
print("Force killing subprocess...")
|
||||
self.process.kill()
|
||||
self.process.join()
|
||||
|
||||
# Clean up shared memory
|
||||
for camera_name, shm in self.shared_memory_blocks.items():
|
||||
try:
|
||||
shm.close()
|
||||
shm.unlink()
|
||||
if self.verbose:
|
||||
print(f"Cleaned up shared memory for {camera_name}")
|
||||
except Exception as e:
|
||||
if self.verbose:
|
||||
print(f"Warning: Failed to cleanup shared memory for {camera_name}: {e}")
|
||||
|
||||
self.shared_memory_blocks.clear()
|
||||
if self.verbose:
|
||||
print("Image publishing subprocess stopped and cleaned up")
|
||||
|
||||
@staticmethod
|
||||
def _image_publish_worker(
|
||||
shared_memory_info, image_dt, zmq_port, stop_event, data_ready_event, verbose
|
||||
):
|
||||
"""Worker function that runs in the subprocess"""
|
||||
# Import dependencies within worker (needed for multiprocessing spawn mode)
|
||||
from .sensor_utils import ImageMessageSchema, ImageUtils, SensorServer
|
||||
|
||||
if verbose:
|
||||
print(f"Worker started! PID: {__import__('os').getpid()}")
|
||||
print(f"Worker stop_event state at start: {stop_event.is_set()}")
|
||||
print(f"Worker data_ready_event state at start: {data_ready_event.is_set()}")
|
||||
|
||||
try:
|
||||
# Initialize ZMQ sensor server
|
||||
sensor_server = SensorServer()
|
||||
sensor_server.start_server(port=zmq_port)
|
||||
|
||||
# Connect to shared memory blocks
|
||||
shared_arrays = {}
|
||||
shm_blocks = {}
|
||||
for camera_name, info in shared_memory_info.items():
|
||||
shm = shared_memory.SharedMemory(name=info["name"])
|
||||
shm_blocks[camera_name] = shm
|
||||
shared_arrays[camera_name] = np.ndarray(
|
||||
info["shape"], dtype=info["dtype"], buffer=shm.buf
|
||||
)
|
||||
|
||||
print(
|
||||
f"Image publishing subprocess started with {len(shared_arrays)} cameras on ZMQ port {zmq_port}"
|
||||
)
|
||||
|
||||
loop_count = 0
|
||||
last_data_time = time.time()
|
||||
|
||||
while not stop_event.is_set():
|
||||
loop_count += 1
|
||||
|
||||
# Wait for new data with shorter timeout for better responsiveness
|
||||
timeout = min(image_dt, 0.1) # Max 100ms timeout
|
||||
data_available = data_ready_event.wait(timeout=timeout)
|
||||
|
||||
current_time = time.time()
|
||||
|
||||
if data_available:
|
||||
data_ready_event.clear()
|
||||
if loop_count % 50 == 0:
|
||||
print("Image publish frequency: ", 1 / (current_time - last_data_time))
|
||||
last_data_time = current_time
|
||||
|
||||
# Collect all camera images and serialize them
|
||||
try:
|
||||
# Copy all images atomically at once
|
||||
image_copies = {name: arr.copy() for name, arr in shared_arrays.items()}
|
||||
|
||||
# Create message with all camera images
|
||||
message_dict = {
|
||||
"images": image_copies,
|
||||
"timestamps": {name: current_time for name in image_copies.keys()},
|
||||
}
|
||||
|
||||
# Create ImageMessageSchema and serialize
|
||||
image_msg = ImageMessageSchema(
|
||||
timestamps=message_dict.get("timestamps"),
|
||||
images=message_dict.get("images", None),
|
||||
)
|
||||
|
||||
# Serialize and send via ZMQ
|
||||
serialized_data = image_msg.serialize()
|
||||
|
||||
# Add individual camera images to the message
|
||||
for camera_name, image_copy in image_copies.items():
|
||||
serialized_data[f"{camera_name}"] = ImageUtils.encode_image(image_copy)
|
||||
|
||||
sensor_server.send_message(serialized_data)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error publishing images: {e}")
|
||||
|
||||
elif verbose and loop_count % 10 == 0:
|
||||
print(f"Subprocess: Still waiting for data... (iteration {loop_count})")
|
||||
|
||||
# Small sleep to prevent busy waiting when no data
|
||||
if not data_available:
|
||||
time.sleep(0.001)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("Image publisher interrupted by user")
|
||||
finally:
|
||||
# Clean up
|
||||
try:
|
||||
for shm in shm_blocks.values():
|
||||
shm.close()
|
||||
sensor_server.stop_server()
|
||||
except Exception as e:
|
||||
print(f"Error during subprocess cleanup: {e}")
|
||||
if verbose:
|
||||
print("Image publish subprocess stopped")
|
||||
@@ -1,71 +0,0 @@
|
||||
from typing import List, Tuple
|
||||
|
||||
import mujoco
|
||||
|
||||
from .sim_utilts import get_body_geom_ids
|
||||
|
||||
|
||||
def check_contact(
|
||||
mj_model: mujoco.MjModel,
|
||||
mj_data: mujoco.MjData,
|
||||
bodies_1: List[str] | str,
|
||||
bodies_2: List[str] | str,
|
||||
return_all_contact_bodies: bool = False,
|
||||
) -> Tuple[bool, List[Tuple[str, str]]] | bool:
|
||||
"""
|
||||
Finds contact between two body groups. Any geom in the body is considered to be in contact.
|
||||
Args:
|
||||
mj_model (MujocoModel): Current simulation object
|
||||
mj_data (MjData): Current simulation data
|
||||
bodies_1 (str or list of int): an individual body name or list of body names.
|
||||
bodies_2 (str or list of int): another individual body name or list of body names.
|
||||
Returns:
|
||||
bool: True if any body in @bodies_1 is in contact with any body in @bodies_2.
|
||||
"""
|
||||
if isinstance(bodies_1, str):
|
||||
bodies_1 = [bodies_1]
|
||||
if isinstance(bodies_2, str):
|
||||
bodies_2 = [bodies_2]
|
||||
|
||||
geoms_1 = [get_body_geom_ids(mj_model, mj_model.body(g).id) for g in bodies_1]
|
||||
geoms_1 = [g for geom_list in geoms_1 for g in geom_list]
|
||||
geoms_2 = [get_body_geom_ids(mj_model, mj_model.body(g).id) for g in bodies_2]
|
||||
geoms_2 = [g for geom_list in geoms_2 for g in geom_list]
|
||||
contact_bodies = []
|
||||
for i in range(mj_data.ncon):
|
||||
contact = mj_data.contact[i]
|
||||
# check contact geom in geoms
|
||||
c1_in_g1 = contact.geom1 in geoms_1
|
||||
c2_in_g2 = contact.geom2 in geoms_2 if geoms_2 is not None else True
|
||||
# check contact geom in geoms (flipped)
|
||||
c2_in_g1 = contact.geom2 in geoms_1
|
||||
c1_in_g2 = contact.geom1 in geoms_2 if geoms_2 is not None else True
|
||||
if (c1_in_g1 and c2_in_g2) or (c1_in_g2 and c2_in_g1):
|
||||
contact_bodies.append(
|
||||
(
|
||||
mj_model.body(mj_model.geom(contact.geom1).bodyid).name,
|
||||
mj_model.body(mj_model.geom(contact.geom2).bodyid).name,
|
||||
)
|
||||
)
|
||||
if not return_all_contact_bodies:
|
||||
break
|
||||
if return_all_contact_bodies:
|
||||
return len(contact_bodies) > 0, set(contact_bodies)
|
||||
else:
|
||||
return len(contact_bodies) > 0
|
||||
|
||||
|
||||
def check_height(
|
||||
mj_model: mujoco.MjModel,
|
||||
mj_data: mujoco.MjData,
|
||||
geom_name: str,
|
||||
lower_bound: float = -float("inf"),
|
||||
upper_bound: float = float("inf"),
|
||||
):
|
||||
"""
|
||||
Checks if the height of a geom is greater than a given height.
|
||||
"""
|
||||
geom_id = mj_model.geom(geom_name).id
|
||||
return (
|
||||
mj_data.geom_xpos[geom_id][2] < upper_bound and mj_data.geom_xpos[geom_id][2] > lower_bound
|
||||
)
|
||||
@@ -1,130 +0,0 @@
|
||||
"""Standalone sensor utilities for camera image publishing via ZMQ"""
|
||||
import base64
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict
|
||||
|
||||
import cv2
|
||||
import msgpack
|
||||
import msgpack_numpy as m
|
||||
import numpy as np
|
||||
import zmq
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageMessageSchema:
|
||||
"""
|
||||
Standardized message schema for image data.
|
||||
Used to serialize/deserialize image data for network transmission.
|
||||
"""
|
||||
|
||||
timestamps: Dict[str, float]
|
||||
"""Dictionary of timestamps, keyed by image identifier (e.g., {"ego_view": 123.45})"""
|
||||
images: Dict[str, np.ndarray]
|
||||
"""Dictionary of images, keyed by image identifier (e.g., {"ego_view": array})"""
|
||||
|
||||
def serialize(self) -> Dict[str, Any]:
|
||||
"""Serialize the message for transmission."""
|
||||
serialized_msg = {"timestamps": self.timestamps, "images": {}}
|
||||
for key, image in self.images.items():
|
||||
serialized_msg["images"][key] = ImageUtils.encode_image(image)
|
||||
return serialized_msg
|
||||
|
||||
@staticmethod
|
||||
def deserialize(data: Dict[str, Any]) -> "ImageMessageSchema":
|
||||
"""Deserialize received message data."""
|
||||
timestamps = data.get("timestamps", {})
|
||||
images = {}
|
||||
for key, value in data.get("images", {}).items():
|
||||
if isinstance(value, str):
|
||||
images[key] = ImageUtils.decode_image(value)
|
||||
else:
|
||||
images[key] = value
|
||||
return ImageMessageSchema(timestamps=timestamps, images=images)
|
||||
|
||||
def asdict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary format."""
|
||||
return {"timestamps": self.timestamps, "images": self.images}
|
||||
|
||||
|
||||
class SensorServer:
|
||||
"""ZMQ-based sensor server for publishing camera images"""
|
||||
|
||||
def start_server(self, port: int):
|
||||
self.context = zmq.Context()
|
||||
self.socket = self.context.socket(zmq.PUB)
|
||||
self.socket.setsockopt(zmq.SNDHWM, 20) # high water mark
|
||||
self.socket.setsockopt(zmq.LINGER, 0)
|
||||
self.socket.bind(f"tcp://*:{port}")
|
||||
print(f"Sensor server running at tcp://*:{port}")
|
||||
|
||||
self.message_sent = 0
|
||||
self.message_dropped = 0
|
||||
|
||||
def stop_server(self):
|
||||
self.socket.close()
|
||||
self.context.term()
|
||||
|
||||
def send_message(self, data: Dict[str, Any]):
|
||||
try:
|
||||
packed = msgpack.packb(data, use_bin_type=True)
|
||||
self.socket.send(packed, flags=zmq.NOBLOCK)
|
||||
except zmq.Again:
|
||||
self.message_dropped += 1
|
||||
print(f"[Warning] message dropped: {self.message_dropped}")
|
||||
self.message_sent += 1
|
||||
|
||||
if self.message_sent % 100 == 0:
|
||||
print(
|
||||
f"[Sensor server] Message sent: {self.message_sent}, message dropped: {self.message_dropped}"
|
||||
)
|
||||
|
||||
|
||||
class SensorClient:
|
||||
"""ZMQ-based sensor client for subscribing to camera images"""
|
||||
|
||||
def start_client(self, server_ip: str, port: int):
|
||||
self.context = zmq.Context()
|
||||
self.socket = self.context.socket(zmq.SUB)
|
||||
self.socket.setsockopt_string(zmq.SUBSCRIBE, "")
|
||||
self.socket.setsockopt(zmq.CONFLATE, True) # last msg only.
|
||||
self.socket.setsockopt(zmq.RCVHWM, 3) # queue size 3 for receive buffer
|
||||
self.socket.connect(f"tcp://{server_ip}:{port}")
|
||||
|
||||
def stop_client(self):
|
||||
self.socket.close()
|
||||
self.context.term()
|
||||
|
||||
def receive_message(self):
|
||||
packed = self.socket.recv()
|
||||
return msgpack.unpackb(packed, object_hook=m.decode)
|
||||
|
||||
|
||||
class ImageUtils:
|
||||
"""Utilities for encoding/decoding images for network transmission"""
|
||||
|
||||
@staticmethod
|
||||
def encode_image(image: np.ndarray) -> str:
|
||||
"""Encode numpy image to base64-encoded JPEG string"""
|
||||
_, color_buffer = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
|
||||
return base64.b64encode(color_buffer).decode("utf-8")
|
||||
|
||||
@staticmethod
|
||||
def encode_depth_image(image: np.ndarray) -> str:
|
||||
"""Encode depth image to base64-encoded PNG string"""
|
||||
depth_compressed = cv2.imencode(".png", image)[1].tobytes()
|
||||
return base64.b64encode(depth_compressed).decode("utf-8")
|
||||
|
||||
@staticmethod
|
||||
def decode_image(image: str) -> np.ndarray:
|
||||
"""Decode base64-encoded JPEG string to numpy image"""
|
||||
color_data = base64.b64decode(image)
|
||||
color_array = np.frombuffer(color_data, dtype=np.uint8)
|
||||
return cv2.imdecode(color_array, cv2.IMREAD_COLOR)
|
||||
|
||||
@staticmethod
|
||||
def decode_depth_image(image: str) -> np.ndarray:
|
||||
"""Decode base64-encoded PNG string to depth image"""
|
||||
depth_data = base64.b64decode(image)
|
||||
depth_array = np.frombuffer(depth_data, dtype=np.uint8)
|
||||
return cv2.imdecode(depth_array, cv2.IMREAD_UNCHANGED)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user