Hi, thanks for open-sourcing the full Go2 implementation, I've also been working on getting an extreme-parkour-style policy onto the Go2, and having a complete teacher–student pipeline with the estimator and deployment code to reference is genuinely useful.
While reproducing the Stage-2 student distillation I ran into two things I'd like to check my understanding on:
1. Gradient flow to the depth backbone in the gated BC+RL update
As far as I can trace, with the default config (consistency_gating_k = 2.0):
- during rollout, the depth latents used by the RL/BC update are stored detached, rl_depth_latent_buf.append(depth_latent.clone().detach()) in on_policy_runner.py,
- update_depth_actor_gated recomputes the action mean from those detached latents, so λ·L_RL + (1−λ)·L_BC only reaches the actor head,
- and the separate latent-alignment update (update_depth_encoder) is disabled by default (enable_latent_loss = False, latent_loss_weight = 0.0 in legged_robot_config.py).
If I'm reading this correctly, the only loss that backpropagates through the FiLM-CNN + Mamba depth backbone is the yaw regression term. Is that intentional for the published results, or were they trained with enable_latent_loss = True (or some other configuration where the encoder gets more supervision)? Section III-C / Eq. 11–12 of the paper reads as if the gated objective trains the student end-to-end, so I want to make sure I'm not missing where the encoder gradients are supposed to come from.
2. NaN divergence around ~8k iterations
Training the student with the released gated defaults, the Mamba depth encoder diverged to NaN around iteration 7.5k–8k and never recovered. Did you observe anything similar during development, and do you rely on any stabilization beyond the max_grad_norm clip (e.g. a tighter encoder LR, warm-up, weight decay, or a norm bound on the SSM blocks)?
3. Exact training configuration for the published student
Could you share the exact flags/config used for the Go2 student results in the paper? In particular:
- --max_iterations (the README says 30000, but Fig. 8's x-axis appears to run to ~50k),
- whether --delay was enabled during the main run, or only for the post-convergence fine-tune mentioned in the README tips,
- enable_latent_loss / latent_loss_weight,
- whether the estimator was warm-started via --load_estimator_checkpoint.
Thanks. Happy to share more details from my runs if useful.
Hi, thanks for open-sourcing the full Go2 implementation, I've also been working on getting an extreme-parkour-style policy onto the Go2, and having a complete teacher–student pipeline with the estimator and deployment code to reference is genuinely useful.
While reproducing the Stage-2 student distillation I ran into two things I'd like to check my understanding on:
1. Gradient flow to the depth backbone in the gated BC+RL update
As far as I can trace, with the default config (consistency_gating_k = 2.0):
If I'm reading this correctly, the only loss that backpropagates through the FiLM-CNN + Mamba depth backbone is the yaw regression term. Is that intentional for the published results, or were they trained with enable_latent_loss = True (or some other configuration where the encoder gets more supervision)? Section III-C / Eq. 11–12 of the paper reads as if the gated objective trains the student end-to-end, so I want to make sure I'm not missing where the encoder gradients are supposed to come from.
2. NaN divergence around ~8k iterations
Training the student with the released gated defaults, the Mamba depth encoder diverged to NaN around iteration 7.5k–8k and never recovered. Did you observe anything similar during development, and do you rely on any stabilization beyond the max_grad_norm clip (e.g. a tighter encoder LR, warm-up, weight decay, or a norm bound on the SSM blocks)?
3. Exact training configuration for the published student
Could you share the exact flags/config used for the Go2 student results in the paper? In particular:
Thanks. Happy to share more details from my runs if useful.