270 lines
13 KiB
Python
270 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
DiscoRL ↔ Gym 集成 - 快速参考
|
||
|
||
使用方式:
|
||
python scripts/QUICK_START.py
|
||
"""
|
||
|
||
quick_ref = """
|
||
|
||
╔════════════════════════════════════════════════════════════════════════════╗
|
||
║ DiscoRL × Gym 集成 - 快速参考卡 ║
|
||
╚════════════════════════════════════════════════════════════════════════════╝
|
||
|
||
┌──────────────────────────────────────────────────────────────────────────┐
|
||
│ 1️⃣ 验证安装 (5 分钟) │
|
||
└──────────────────────────────────────────────────────────────────────────┘
|
||
|
||
$ python scripts/test_disco_setup.py
|
||
|
||
预期输出:
|
||
✓ All core components working!
|
||
✓ 所有 5 个测试通过
|
||
|
||
|
||
┌──────────────────────────────────────────────────────────────────────────┐
|
||
│ 2️⃣ 在 CartPole 上训练 (1 分钟) │
|
||
└──────────────────────────────────────────────────────────────────────────┘
|
||
|
||
$ python scripts/train_disco_cartpole.py
|
||
|
||
配置:
|
||
- batch_size=4
|
||
- trajectory_length=32
|
||
- num_iterations=50
|
||
|
||
预期输出:
|
||
✓ Training Complete
|
||
✓ Final avg reward: ~0.97
|
||
|
||
|
||
┌──────────────────────────────────────────────────────────────────────────┐
|
||
│ 3️⃣ 验证集成 (30 秒) │
|
||
└──────────────────────────────────────────────────────────────────────────┘
|
||
|
||
$ python scripts/poc_integration.py
|
||
|
||
预期输出:
|
||
✓ Success! DiscoRL ↔ Gym integration works!
|
||
|
||
|
||
╔════════════════════════════════════════════════════════════════════════════╗
|
||
║ 核心代码段
|
||
╚════════════════════════════════════════════════════════════════════════════╝
|
||
|
||
┌──────────────────────────────────────────────────────────────────────────┐
|
||
│ A) 环境设置 │
|
||
└──────────────────────────────────────────────────────────────────────────┘
|
||
|
||
from disco_cartpole_env import DiscoCartPoleEnv
|
||
|
||
env = DiscoCartPoleEnv(batch_size=4, max_steps=500)
|
||
obs_spec = env.single_observation_spec()
|
||
act_spec = env.single_action_spec()
|
||
|
||
|
||
┌──────────────────────────────────────────────────────────────────────────┐
|
||
│ B) 代理创建 │
|
||
└──────────────────────────────────────────────────────────────────────────┘
|
||
|
||
from disco_rl import agent as disco_agent
|
||
|
||
agent_settings = disco_agent.get_settings_disco()
|
||
agent = disco_agent.Agent(
|
||
single_observation_spec=obs_spec,
|
||
single_action_spec=act_spec,
|
||
agent_settings=agent_settings,
|
||
batch_axis_name=None,
|
||
)
|
||
|
||
learner_state = agent.initial_learner_state(rng_key)
|
||
actor_state = agent.initial_actor_state(rng_key)
|
||
|
||
|
||
┌──────────────────────────────────────────────────────────────────────────┐
|
||
│ C) 数据收集 │
|
||
└──────────────────────────────────────────────────────────────────────────┘
|
||
|
||
state, timestep = env.reset(rng_key=subkey)
|
||
|
||
for t in range(trajectory_length):
|
||
# 代理推理
|
||
actor_timestep, actor_state = agent.actor_step(
|
||
learner_state.params,
|
||
rng,
|
||
timestep,
|
||
actor_state,
|
||
)
|
||
|
||
# 环境步进
|
||
state, timestep = env.step(state, actor_timestep.actions)
|
||
|
||
# 记录数据
|
||
observations.append(timestep.observation['observation'])
|
||
actions.append(actor_timestep.actions)
|
||
rewards.append(timestep.reward)
|
||
# ... 等
|
||
|
||
|
||
┌──────────────────────────────────────────────────────────────────────────┐
|
||
│ D) 参数更新 │
|
||
└──────────────────────────────────────────────────────────────────────────┘
|
||
|
||
from disco_rl import types
|
||
|
||
rollout = types.ActorRollout(
|
||
observations=jnp.stack(observations),
|
||
actions=jnp.stack(actions),
|
||
rewards=jnp.stack(rewards),
|
||
discounts=jnp.stack(discounts),
|
||
agent_outs=agent_outs_stacked,
|
||
logits=jnp.stack(logits),
|
||
states=actor_state,
|
||
)
|
||
|
||
new_learner_state, new_actor_state, logs = agent.learner_step(
|
||
rng=rng,
|
||
rollout=rollout,
|
||
learner_state=learner_state,
|
||
agent_net_state=actor_state,
|
||
update_rule_params=update_rule_params,
|
||
is_meta_training=False,
|
||
)
|
||
|
||
|
||
╔════════════════════════════════════════════════════════════════════════════╗
|
||
║ 常见问题与答案
|
||
╚════════════════════════════════════════════════════════════════════════════╝
|
||
|
||
Q: 如何用于自定义环境?
|
||
A: 1. cp disco_cartpole_env.py disco_custom_env.py
|
||
2. 修改 __init__ 中的环境创建逻辑
|
||
3. 调整 action_spec 和 observation_spec
|
||
|
||
Q: 如何加载预训练权重?
|
||
A: 预训练权重加载目前在开发中
|
||
临时解决: 使用随机初始化的参数
|
||
|
||
Q: 如何扩展到多个 GPU?
|
||
A: 1. 移除 os.environ['JAX_PLATFORMS'] = 'cpu'
|
||
2. 使用 jax.device_count() 获取设备数
|
||
3. 在 agent.learner_step 中设置 batch_axis_name='devices'
|
||
|
||
Q: 性能太慢怎么办?
|
||
A: • 增加 batch_size (更多并行环境)
|
||
• 减少 trajectory_length
|
||
• 使用 GPU (移除 CPU-only 设置)
|
||
• 减少网络大小 (调整 agent_settings)
|
||
|
||
Q: CartPole 不难吗?
|
||
A: CartPole 是验证集成的好工具
|
||
一旦工作，应用到实际环境:
|
||
• gym_env_250326_erase.py (自定义任务)
|
||
• 或任何其他 Gym 兼容环境
|
||
|
||
|
||
╔════════════════════════════════════════════════════════════════════════════╗
|
||
║ 文件参考
|
||
╚════════════════════════════════════════════════════════════════════════════╝
|
||
|
||
核心文件 (必需):
|
||
disco_cartpole_env.py 环境适配器 ← 为自定义环境修改这个
|
||
|
||
工具文件:
|
||
disco_weights.py 权重加载
|
||
train_disco_cartpole.py 训练循环
|
||
test_disco_setup.py 测试套件
|
||
|
||
文档:
|
||
INTEGRATION_GUIDE.py 详细指南
|
||
COMPLETION_SUMMARY.py 完成报告
|
||
QUICK_START.py 本文件
|
||
|
||
配置文件 (如需):
|
||
config_*.json 在 configs/ 中
|
||
|
||
|
||
╔════════════════════════════════════════════════════════════════════════════╗
|
||
║ 关键数据类型
|
||
╚════════════════════════════════════════════════════════════════════════════╝
|
||
|
||
types.EnvironmentTimestep:
|
||
observation: dict {'observation': Array([B, ...], float32)}
|
||
step_type: Array([B], int32) 0=MID, 1=LAST
|
||
reward: Array([B], float32)
|
||
|
||
types.ActorTimestep:
|
||
observations: dict
|
||
actions: Array([B], int32) 动作索引
|
||
agent_outs: dict 策略网络输出
|
||
logits: Array([B, num_actions])
|
||
...
|
||
|
||
types.ActorRollout:
|
||
observations, actions, rewards, discounts, agent_outs, logits, states
|
||
(所有都在时间维度堆叠: [T, B, ...])
|
||
|
||
|
||
╔════════════════════════════════════════════════════════════════════════════╗
|
||
║ 环境规格 (CartPole)
|
||
╚════════════════════════════════════════════════════════════════════════════╝
|
||
|
||
观测空间: Box(4,) ← 杆角度、角速度、车位置、车速度
|
||
动作空间: Discrete(2) ← 0=向左推, 1=向右推
|
||
奖励: +1.0 ← 每一步 (最多 500 步)
|
||
完成: 当角度 > 24° 或位置超出界限
|
||
|
||
|
||
╔════════════════════════════════════════════════════════════════════════════╗
|
||
║ 常用命令
|
||
╚════════════════════════════════════════════════════════════════════════════╝
|
||
|
||
# 快速验证
|
||
python scripts/test_disco_setup.py
|
||
|
||
# 完整训练 (50 iter, 4 batch)
|
||
python scripts/train_disco_cartpole.py
|
||
|
||
# 端到端演示
|
||
python scripts/poc_integration.py
|
||
|
||
# 查看详细文档
|
||
python scripts/INTEGRATION_GUIDE.py | less
|
||
|
||
# 查看完成报告
|
||
python scripts/COMPLETION_SUMMARY.py | less
|
||
|
||
# 运行此快速参考
|
||
python scripts/QUICK_START.py
|
||
|
||
|
||
╔════════════════════════════════════════════════════════════════════════════╗
|
||
║ 总结
|
||
╚════════════════════════════════════════════════════════════════════════════╝
|
||
|
||
✓ DiscoRL (JAX) ↔ Gym (任何环境)
|
||
|
||
✓ 完整的训练循环
|
||
|
||
✓ 预验证的代码
|
||
|
||
✓ 可立即复用的模板
|
||
|
||
准备开始? 运行:
|
||
python scripts/test_disco_setup.py
|
||
|
||
有问题? 查看:
|
||
python scripts/INTEGRATION_GUIDE.py
|
||
|
||
"""
|
||
|
||
print(quick_ref)
|
||
|
||
# 如果用户想保存
|
||
import sys
|
||
if len(sys.argv) > 1 and sys.argv[1] == '--save':
|
||
with open('/home/frank14f/Frank_LBM/scripts/QUICK_START.txt', 'w') as f:
|
||
f.write(quick_ref)
|
||
print("✓ Saved to QUICK_START.txt")
|