Problems since using bc algorithms

Hello brax team,

thanks for your implementation of bc, there are some code problems I found when using the bc. I tried to fix them with AI, but I think you may have better solution for these problems. most of them are facing jax compile errors.







thanks,
Tau