TorchRL项目中SyncDataCollector的CUDA同步问题解析

问题背景

在TorchRL项目的实际使用中，开发者发现当使用SyncDataCollector进行数据收集时，系统会在约18万步后崩溃，并出现"resource_tracker: There appear to be 4 leaked semaphore objects to clean up at shutdown"的警告信息。这一问题在使用DMControlEnv和Maniskill3等环境时尤为明显。

问题现象

主要表现包括：

1. 训练过程中突然崩溃
2. 控制台输出关于泄露信号量对象的警告
3. 在使用LazyMemmapStorage时问题更为频繁
4. 在评估阶段使用env.rollout时出现CUDA非法内存访问错误

根本原因分析

经过深入调查，发现问题的核心在于CUDA同步机制与特定环境(如Isaac Lab、Maniskill3等)的兼容性问题：

1. CUDA同步冲突：SyncDataCollector内部设置了自定义CUDA流，而某些物理引擎(如PhysX)使用默认CUDA流，导致同步问题

2. 内存管理问题：当使用LazyMemmapStorage时，磁盘空间不足会加剧问题的出现

3. 非阻塞传输风险：TensorDict在将数据从GPU传输到CPU时默认使用non_blocking=True，可能导致数据损坏

解决方案

针对这一问题，TorchRL团队提出了多种解决方案：

1. 禁用CUDA同步

在SyncDataCollector中新增no_cuda_sync参数，允许用户关闭CUDA同步功能：

SyncDataCollector(
    create_env_fn=env,
    policy=policy,
    no_cuda_sync=True,  # 新增参数
    ...
)

2. 使用替代存储方案

将LazyMemmapStorage替换为LazyTensorStorage可以缓解部分问题：

replay_buffer = TensorDictReplayBuffer(
    storage=LazyTensorStorage(**storage_kwargs),  # 替代LazyMemmapStorage
    ...
)

3. 显式控制数据传输

在进行评估时，显式指定数据传输方式：

rollouts = self.eval_env.rollout(
    ...
).to(device="cpu", non_blocking=False)  # 强制同步传输

技术细节深入

CUDA流同步机制

TorchRL默认使用自定义CUDA流来提高性能，但这与某些物理引擎的CUDA使用方式冲突。当两个不同的流尝试访问相同的内存区域时，如果没有正确的同步机制，就会导致非法内存访问。

内存映射文件问题

LazyMemmapStorage依赖内存映射文件技术，当系统磁盘空间不足时：

1. 无法正确映射内存区域
2. 导致信号量泄露
3. 最终引发系统级错误

数据传输安全

TensorDict默认使用non_blocking=True进行D2H(Device to Host)传输，这种异步方式虽然快，但需要开发者自行确保同步。对于物理引擎等特殊环境，建议使用同步传输(non_blocking=False)来保证数据一致性。

最佳实践建议

1. 环境兼容性检查：在使用物理引擎(PhysX等)时，优先考虑禁用CUDA同步

2. 存储方案选择：确保有足够磁盘空间时再使用LazyMemmapStorage，否则选择LazyTensorStorage

3. 评估阶段安全措施：在评估循环中显式使用同步数据传输

4. 监控资源使用：定期检查GPU内存和磁盘空间使用情况

总结

TorchRL中的SyncDataCollector为高效数据收集提供了强大支持，但在特定环境下需要注意CUDA同步和内存管理问题。通过合理配置no_cuda_sync参数、选择合适的存储方案以及控制数据传输方式，可以有效避免训练过程中的崩溃问题，确保强化学习训练的稳定进行。