主要是记录当前项目里 TensorBoard 视频和 Histograms 的记录方式。

TensorBoard 除了记录 reward 曲线和训练曲线,也可以看 Histograms 直方图,还可以把训练过程中的视频帧写进去。这里用天工 AMP 项目里的改法来记一下。原工程里不带这几项,是后面自己补的,所以这篇主要关注实现位置和变量对应关系。


当前实现方式

这次修改主要涉及三个文件:

1
2
3
legged_lab/scripts/train.py
rsl_rl/rsl_rl/runners/amp_on_policy_runner.py
legged_lab/envs/tienkung/tienkung_env.py

它们分别处理三件事:

· train.py 负责解析命令行参数,并把视频、直方图相关配置传给 runner。
· amp_on_policy_runner.py 负责在训练循环里写 Histograms 和视频。
· tienkung_env.py 负责提供 render(),并控制录视频时的相机视角。


训练脚本里增加参数

位置:legged_lab/scripts/train.py

这里加了几个命令行参数:

1
2
3
4
5
6
7
8
parser.add_argument("--video", action="store_true", default=False, help="Record TensorBoard videos during training.")
parser.add_argument("--video_length", type=int, default=200, help="Length of each recorded video in env steps.")
parser.add_argument(
"--video_interval", type=int, default=2000, help="Interval between TensorBoard video recordings in learning iterations."
)
parser.add_argument(
"--histogram_interval", type=int, default=100, help="Interval between TensorBoard histogram recordings."
)

这几个参数的意思比较直接:

· --video:是否开启训练视频记录。
· --video_length:每次录多少个 env step。
· --video_interval:每隔多少个 learning iteration 录一次视频。
· --histogram_interval:每隔多少个 learning iteration 写一次 Histograms。

创建环境时,如果开启 --video,环境要使用 rgb_array 渲染模式:

1
env = env_class(env_cfg, args_cli.headless, render_mode="rgb_array" if args_cli.video else None)

这里容易漏。如果不传 render_mode="rgb_array",后面 runner 调 env.render() 时可能拿不到 RGB 帧,TensorBoard 里也就没有视频内容。

然后把记录频率传给 runner:

1
2
3
runner_cfg["histogram_interval"] = args_cli.histogram_interval
runner_cfg["video_interval"] = args_cli.video_interval if args_cli.video else 0
runner_cfg["video_length"] = args_cli.video_length

这里有个小细节:不开 --video 时,video_interval 直接写成 0。runner 里可以用这个值判断是否跳过视频录制,不需要再单独维护一个 video enable 标志。


Runner 里记录 Histograms

位置:rsl_rl/rsl_rl/runners/amp_on_policy_runner.py

训练循环里,每隔 histogram_interval 个 learning iteration 记录一次:

1
2
if self._should_log_histograms(it):
self._log_histograms(it)

当前记录的主要内容是:

1
2
3
4
5
6
tensors = {
"Histograms/actions": storage.actions,
"Histograms/obs": storage.observations,
"Histograms/policy/action_mean": storage.mu,
"Histograms/value": storage.values,
}

按 RSL-RL 这类 rollout storage 的常见结构看,这些 tensor 一般是下面这种形状:

1
2
3
4
storage.actions       [num_steps, num_envs, num_actions]
storage.observations [num_steps, num_envs, num_obs]
storage.mu [num_steps, num_envs, num_actions]
storage.values [num_steps, num_envs, 1]

如果代码里先 flatten 再写 TensorBoard,形状可能会变成:

1
[num_steps * num_envs, dim]

但看 Histograms 时重点不是 batch 维度本身,而是每个维度的数值分布有没有异常。比如 action 是否长期顶到边界,observation 有没有突然炸掉,value 是否出现很大的偏移。

同时还记录 actor 和 critic 的参数分布、梯度分布:

1
2
3
4
5
for module_name, module in modules.items():
for name, param in module.named_parameters():
self.writer.add_histogram(f"Histograms/{module_name}/params/{safe_name}", param, iteration)
if param.grad is not None:
self.writer.add_histogram(f"Histograms/{module_name}/grads/{safe_name}", param.grad, iteration)

这里的 safe_name 一般是把参数名里的特殊字符处理掉,避免 TensorBoard tag 里出现不太好看的层级。比如 mlp.0.weight 这种名字,最后会变成一个更稳定的 tag。

梯度分布这里比较有用。reward 曲线飘的时候,单看 scalar loss 经常看不出来问题来源。Histograms 可以帮忙判断是 observation 输入分布不对,还是 action 输出顶死,还是某些层的梯度突然爆了。


Runner 里记录视频

位置:rsl_rl/rsl_rl/runners/amp_on_policy_runner.py

训练循环里,每隔 video_interval 个 learning iteration 录一段视频:

1
2
if self._should_log_video(it):
obs, privileged_obs, amp_obs = self._record_video(obs, privileged_obs, amp_obs, it)

这里 _record_video() 返回新的 obsprivileged_obsamp_obs,是因为录视频本质上也会继续 step 环境。录制过程会消耗环境状态,所以录完以后要把当前 observation 返回给训练循环继续用。

录制时调用环境的 render() 抓取 RGB 帧:

1
2
3
4
for _ in range(self.video_length):
frame = self.env.render()
if frame is not None:
frames.append(frame)

单帧通常是这样的格式:

1
2
3
frame: [H, W, 3]
dtype: uint8
channel: RGB

多帧堆起来之后,大概就是:

1
frames: [T, H, W, 3]

其中 T 对应 video_length,但实际写入时可能小于 video_length,因为 render() 返回 None 的帧会被跳过。

抓到的帧会编码成 GIF,并以 TensorBoard image summary 的方式写入:

1
self._add_tensorboard_gif("Videos/train", frames, iteration, fps=fps)

注意:当前实现没有使用 writer.add_video(),也没有依赖 moviepyVideos/train 这个 tag 通常在 TensorBoard 的 Images 面板里查看,不是在 Video 面板里看。


环境里提供 render 和相机视角

位置:legged_lab/envs/tienkung/tienkung_env.py

环境增加了 render_mode,用于支持 rgb_array

1
2
3
def __init__(..., headless, render_mode: str | None = None):
self.render_mode = render_mode
self.metadata = {"render_modes": [None, "human", "rgb_array"]}

render() 里先更新相机,再返回 RGB 图像:

1
2
3
4
def render(self, recompute: bool = False):
self._update_video_camera_view()
...
return rgb_data[:, :, :3]

这里返回 rgb_data[:, :, :3],是为了只保留 RGB 三个通道。如果底层拿到的是 RGBA,最后一个 alpha 通道不需要写进 TensorBoard。

录视频的角度由 _update_video_camera_view() 控制。当前相机跟随机器人 root 位置,使用一个固定偏移量看向机器人:

1
2
self._video_camera_eye_offset = np.array([-3.0, -3.0, 1.6])
self._video_camera_target_offset = np.array([0.0, 0.0, 0.7])

实际设置相机位置:

1
2
3
4
root_pos = self.robot.data.root_pos_w[0, :3].detach().cpu().numpy()
eye = root_pos + self._video_camera_eye_offset
target = root_pos + self._video_camera_target_offset
self.sim.set_camera_view(eye=eye, target=target, camera_prim_path=self._video_cam_prim_path)

这表示相机不是固定在世界原点,而是跟着机器人移动。

eye_offset 决定相机站在哪里看,target_offset 决定相机看向机器人身体的哪个位置。当前设置大概是从机器人斜后上方看过去,目标点在 root 上方一点,能看到身体姿态和步态变化。

这里也容易出问题:

· root_pos 用的是第 0 个 env 的机器人位置,所以录出来的是第一个环境。
· 如果 eye_offset 太低,地面或者身体会挡住视角。
· 如果 target_offset 太低,相机可能一直盯着脚底。
· 如果训练是 headless,需要确认底层仿真仍然支持 RGB 渲染。


训练命令

快速验证视频是否能写入:

1
python legged_lab/scripts/train.py --task=walk --headless --logger=tensorboard --num_envs=1024 --video --video_interval=50 --video_length=100 --histogram_interval=100 --resume True --load_run 2026-07-07_12-39-05 --checkpoint model_9400.pt

这个命令主要用来测试视频录制,video_interval=50 比较密,能更快在 TensorBoard 里看到结果。正常长时间训练时不建议录太频繁,视频编码和写盘都会占资源。

正常训练可以使用:

1
python legged_lab/scripts/train.py --task=walk --headless --logger=tensorboard --num_envs=1024 --video --video_interval=2000 --video_length=200 --histogram_interval=100

这里要注意几个配置之间的关系:

· num_envs=1024 只影响并行采样数量,不表示视频会录 1024 个环境。当前相机只跟第 0 个环境。
· video_length=200 表示每次录 200 个 env step,不是 200 个 learning iteration。
· video_interval=2000 表示每隔 2000 个 learning iteration 录一次。
· histogram_interval=100 表示每 100 个 learning iteration 记录一次分布。


TensorBoard 启动命令

推荐用当前 conda 环境的 Python 启动 TensorBoard:

1
python -m tensorboard.main --logdir .\logs --port 6006

也可以使用:

1
tensorboard --logdir .\logs --port 6006

如果 tensorboard 命令报 launcher 路径错误,优先用:

1
python -m tensorboard.main --logdir .\logs --port 6006

这个问题一般是因为当前环境和系统里注册的 tensorboard.exe 不一致。用 python -m 可以保证 TensorBoard 从当前 conda 环境里启动。


Histograms 里看什么

主要可以记录下面这些内容:

TensorBoard tag 含义
Histograms/actions 实际送入环境的动作。
Histograms/obs 输入 actor 的 observation。
Histograms/policy/action_mean actor/policy 输出的动作均值。
Histograms/value critic 对状态价值的估计。
Histograms/actor/params/... actor 网络各层参数分布。
Histograms/actor/grads/... actor 网络各层梯度分布。
Histograms/critic/params/... critic 网络各层参数分布。
Histograms/critic/grads/... critic 网络各层梯度分布。

这些内容可以直观看 observation、action 或者网络中某个量的范围、均值、方差、是否有异常值,以及它们随训练如何变化。

比如 actor 可以看到实际网络输出动作的范围以及分布。甚至可以单独记录某一个关节的动作变化。

actor 动作分布

也可以看某个学习层的 bias 梯度分布之类的内容。

bias 梯度分布

还可以看某一层的权重变化,用来判断训练过程中的数据与网络输出是否处于正常范围。定位异常时,一般会顺着 observation、action、value、网络参数、梯度这些位置看。

这里没有一个固定标准。比如梯度大不一定错,关键要看它是不是突然成片变成极端值,或者某几层长期接近 0。具体判断还是要结合项目慢慢积累。


视频效果

最后是 video 的效果。当前实现可以设定每隔多少个 iteration 录制一次,然后在 TensorBoard 里拖动不同 step,对比不同时期的训练状态。

训练视频记录

这个功能对机器人训练比较有用。reward 曲线只能说明数值在变,但视频能直接看策略是不是在用奇怪的姿态刷奖励,或者是不是某个阶段开始抖动、卡住、摔倒。