# train_xense

**Repository Path**: blouce/train_xense

## Basic Information

- **Project Name**: train_xense
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-05-19
- **Last Updated**: 2026-05-19

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 机器人末端执行器预测网络
# 机器人主从位置补偿预测网络 (Master Position Delta Predictor)

这是一个专门设计的 PyTorch 神经网络实现，用于在无本体遥操作场景下，利用从端（Follower）的位姿和触觉反馈，精确预测主端（Master）的位置补偿。

## 📋 项目概述

### 核心任务
在遥操作中，当从端受到力觉反馈时，主端位置相对于从端会产生细微偏移。本项目通过训练一个残差 MLP 网络来捕捉这种“力-位”映射关系。

### 输入与输出设计
为了提高模型精度和训练稳定性，本项目采用**增量预测**策略。

- **输入 (15D 特征)**:
  - `tactile_right_right`: 6D 触觉传感器数据 (Fx, Fy, Fz, Tx, Ty, Tz)
  - `follow_right_position`: 从端 3D 位置 (x, y, z)
  - `follow_right_rotation`: 从端 3D 欧拉角 (通过 Sin-Cos 或 6D 编码转换为 6维向量)
- **输出 (3D 目标)**:
  - `pos_delta`: 主端与从端的位置差值，即 `master_right_position - follow_right_position`。

### 关键策略
1. **维度缩减**: 仅预测 3D 位置。旋转部分（Rotation）和夹抓（Gripper）在推理阶段直接从 Follower 透传，不参与神经网络训练。
2. **残差学习 (Residual Learning)**: 网络学习的是“位置偏差”而非“绝对坐标”，极大降低了学习熵值。
3. **特征解耦**: 触觉数据在进入网络前进行归一化，防止力矩数值淹没位姿特征。

## 🚀 网络架构

1.  **特征编码层**: 将 15D 输入投影至 256D 隐藏空间。
2.  **残差块 (4 个 Residual Blocks)**:
    - 结构: `Linear -> LayerNorm -> GeLU -> Dropout(0.1) -> Linear`
    - 每个块均包含 Skip Connection。
3.  **回归头**: 最终输出层映射为 3D (`delta_x, delta_y, delta_z`)。

## 📁 数据准备

### JSON 格式要求
JSON 文件应包含 `data` 数组，每帧包含：
- `tactile_right_right`: [6维]
- `follow_right_position`: [3维]
- `follow_right_rotation`: [3维, RPY]
- `master_right_position`: [3维, 用于计算训练 Target]

### 数据转换逻辑
在训练过程中，Dataset 类会自动执行：
`target = master_right_position - follow_right_position`

## 🛠️ 核心脚本说明

### 1. 训练脚本 (`train_from_json.py`)
- **损失函数**: 专门针对位置偏移的 `MSELoss`。
- **性能指标**: 
  - **R² (决定系数)**: 目标是让测试集位置预测的 R² > 0.85。
  - **MAE**: 平均绝对误差（单位：米）。

### 2. 推理与填充脚本 (`infer_and_fill.py`)
该脚本用于处理原始 JSON，生成供 PI0 训练的完整数据。
**填充逻辑**:
- `master_right_position = follow_right_position + model_predict(pos_delta)`
- `master_right_rotation = follow_right_rotation` (直接复制从端旋转)
- `master_right_gripper = follow_right_gripper` (直接复制从端夹抓)

## 📈 实验指南

### 如何提升精度？
- **滑动窗口**: 如果单帧预测 $R^2$ 低于 0.7，请开启 `--window-size 10` 参数，引入时序上下文。
- **旋转编码**: 输入的欧拉角必须转换为连续表示（6D 或 Sin-Cos），否则模型在某些姿态下会失效。
- **Loss 权重**: 如果位置的 X/Y 轴误差远大于 Z 轴，可以考虑在 MSELoss 中引入轴向加权。

## 📦 依赖环境
- Python 3.8+
- PyTorch 2.0+
- Numpy, Matplotlib, Scikit-learn

## 📝 许可证
MIT License

如果 JSON 里保存的是一段长度大于 10 的时序，脚本会按 10 帧窗口切样本；如果正好是 10 帧，也可以直接训练。

## 📁 文件结构

```
train_xense/
├── neural_network.py          # 核心模型定义
├── advanced_examples.py        # 高级使用示例
├── train_from_json.py         # 直接从 JSON 文件夹训练
├── requirements.txt            # 依赖项
└── README.md                   # 本文档
```

## 💻 代码详解

### 核心类

#### `RobotEndEffectorPredictor`
主要的神经网络模型

```python
from neural_network import RobotEndEffectorPredictor

model = RobotEndEffectorPredictor(
    input_frames=10,
    input_dim_per_frame=15,
    hidden_dim=256,
    num_residual_blocks=4,
    output_dim=10,
    dropout_rate=0.1
)

# 前向传播
output = model(input_tensor)  # (batch, 10, 15) → (batch, 10)
```

#### `ResidualMLPBlock`
单个残差块的实现

```python
from neural_network import ResidualMLPBlock

block = ResidualMLPBlock(hidden_dim=256, dropout_rate=0.1)
output = block(input_tensor)
```

#### `RoboticDataset`
数据集和预处理

```python
from neural_network import RoboticDataset

dataset = RoboticDataset(X_data, y_data, normalize=True)
# 支持自动反归一化
denormalized = dataset.denormalize_targets(predictions)
```

#### `Trainer`
训练管理器

```python
from neural_network import Trainer

trainer = Trainer(model, device='cuda', learning_rate=1e-3)
history = trainer.fit(train_loader, val_loader, epochs=100, patience=20)
predictions = trainer.predict(test_input)
```

### 工具函数

#### 欧拉角转换
```python
from neural_network import euler_to_rotation_6d, rotation_6d_to_matrix

# 欧拉角 (弧度) → 6D 表示
euler_angles = torch.tensor([[0.1, 0.2, 0.3]])  # (batch, 3)
rotation_6d = euler_to_rotation_6d(euler_angles)  # (batch, 6)

# 6D 表示 → 旋转矩阵
rotation_matrix = rotation_6d_to_matrix(rotation_6d)  # (batch, 3, 3)
```

#### 模型管理
```python
from advanced_examples import save_model, load_model

# 保存模型
save_model(model, "./checkpoints", "my_model")

# 加载模型
loaded_model = load_model("./checkpoints/my_model.pt")
```

#### 性能评估
```python
from advanced_examples import PerformanceEvaluator

evaluator = PerformanceEvaluator()
metrics = evaluator.calculate_metrics(predictions, targets)
evaluator.print_metrics(metrics, "验证集")
```

## 🎯 使用示例

### 完整的训练流程

```python
import torch
from torch.utils.data import DataLoader
import numpy as np
from neural_network import (
    RobotEndEffectorPredictor,
    RoboticDataset,
    Trainer
)

# 1. 准备数据
X_train = np.random.randn(1000, 10, 15).astype(np.float32)
y_train = np.random.randn(1000, 9).astype(np.float32)
X_val = np.random.randn(200, 10, 15).astype(np.float32)
y_val = np.random.randn(200, 9).astype(np.float32)

# 2. 创建数据集和加载器
train_dataset = RoboticDataset(X_train, y_train, normalize=True)
val_dataset = RoboticDataset(X_val, y_val, normalize=False)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

# 3. 创建模型
model = RobotEndEffectorPredictor(
    input_frames=10,
    input_dim_per_frame=15,
    hidden_dim=256,
    num_residual_blocks=4,
    output_dim=10
)

# 4. 训练
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
trainer = Trainer(model, device=device, learning_rate=1e-3)

history = trainer.fit(
    train_loader,
    val_loader,
    epochs=100,
    patience=15,
    verbose=True
)

# 5. 预测
test_input = torch.randn(5, 10, 15)
predictions = trainer.predict(test_input)
print(f"预测形状: {predictions.shape}")
```

### 欧拉角到 6D 旋转表示转换

```python
import torch
from neural_network import euler_to_rotation_6d, rotation_6d_to_matrix

# 创建欧拉角 (单位：弧度)
euler_angles = torch.tensor([
    [0.0, 0.0, 0.0],           # 无旋转
    [3.14159/2, 0.0, 0.0],     # 绕 X 轴 90 度
    [0.0, 0.0, 3.14159/4],     # 绕 Z 轴 45 度
])

# 转换为 6D 表示
rotation_6d = euler_to_rotation_6d(euler_angles)
print(f"6D 旋转表示: {rotation_6d}")

# 转换为旋转矩阵
rotation_matrix = rotation_6d_to_matrix(rotation_6d)
print(f"旋转矩阵形状: {rotation_matrix.shape}")  # (3, 3, 3)

# 验证正交性
print(f"行列式: {torch.det(rotation_matrix)}")  # 应接近 1
```

### 模型保存和加载

```python
from advanced_examples import save_model, load_model

# 训练模型...

# 保存
save_model(model, "./checkpoints", "robot_model_v1")

# 加载
loaded_model = load_model("./checkpoints/robot_model_v1.pt")

# 使用加载的模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
loaded_model.to(device)
predictions = loaded_model(test_input.to(device))
```

## 📊 性能指标

训练过程会计算以下指标：

- **MSE** (Mean Squared Error): 均方误差
- **RMSE** (Root Mean Squared Error): 均方根误差
- **MAE** (Mean Absolute Error): 平均绝对误差
- **R²**: 决定系数

```python
from advanced_examples import PerformanceEvaluator

evaluator = PerformanceEvaluator()
metrics = evaluator.calculate_metrics(predictions, targets)

print(f"RMSE: {metrics['rmse_mean']:.6f}")
print(f"MAE:  {metrics['mae_mean']:.6f}")
print(f"R²:   {metrics['r2_mean']:.6f}")
```

## 🔧 高级配置

### 调整网络架构

```python
# 增加隐藏层维度
model = RobotEndEffectorPredictor(
    input_frames=10,
    input_dim_per_frame=15,
    hidden_dim=512,           # 从 256 增加到 512
    num_residual_blocks=6,    # 增加块数
    output_dim=10,
    dropout_rate=0.2          # 增加 dropout
)
```

### 自定义优化器

```python
trainer = Trainer(model, learning_rate=5e-4)

# 使用 SGD 代替 Adam
trainer.optimizer = torch.optim.SGD(
    model.parameters(),
    lr=1e-3,
    momentum=0.9
)
```

### 学习率调度

```python
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    trainer.optimizer,
    T_max=100
)

# 在训练循环中使用
for epoch in range(epochs):
    train_loss = trainer.train_epoch(train_loader)
    scheduler.step()
```

## 📈 可视化

### 训练历史

```python
from advanced_examples import plot_training_history

plot_training_history(history, save_path="training_history.png")
```

### 预测 vs 目标

```python
from advanced_examples import plot_predictions_vs_targets

output_names = ['Pos X', 'Pos Y', 'Pos Z', 'Rot 1-6']
plot_predictions_vs_targets(
    predictions, 
    targets,
    output_names=output_names,
    save_path="predictions.png"
)
```

## 🐛 常见问题

### Q: 模型不收敛怎么办？
A: 尝试以下方法：
- 降低学习率（例如 1e-4）
- 增加训练 epoch 数
- 检查数据是否正确归一化
- 增加隐藏层维度
- 减少 dropout 率


## 📝 许可证

MIT License

## 📧 反馈和贡献

欢迎提交问题和改进建议！

## 🔗 相关资源

- [PyTorch 官方文档](https://pytorch.org/docs/)
- [旋转表示方法](https://github.com/papagina/RotationContinuity)
- [残差网络论文](https://arxiv.org/abs/1512.03385)

---

**最后更新**: 2026 年 5 月 14 日