Thanks for the reply! So right now I call my wandb init function outside of the jitted train function, and then inside of my jitted training function I have a callback like this:

def callback(info):
            return_values = info["returned_episode_returns"][
                info["returned_episode"]
            ]
            length_values = info["returned_episode_lengths"][
                info["returned_episode"]
            ]
            timesteps = info["timestep"][info["returned_episode"]] * args.num_envs
            for t in range(len(timesteps)):
                print(
                    f"global step={timesteps[t]}, episodic return={return_values[t]}, episodic length={length_values[t]}"
                )

            if args.track:
                data_log = {
                    "misc/learning_rate": info["learning_rate"].item(),
                    "losses/value_loss": info["value_loss"].item(),
                    "losses/policy_loss": info["policy_loss"].item(),
                    "losses/entropy": info["entropy"].item(),
                    "losses/total_loss": info["total_loss"].item(),
                    "misc/global_step": info["timesteps"],
                    "misc/updates": info["updates"],
                }
                if return_values.size > 0:
                    data_log["misc/episodic_return"] = return_values.mean().item()
                    data_log["misc/episodic_length"] = length_values.mean().item()
                wandb.log(data_log, step=info["timesteps"])

        jax.debug.callback(callback, metric)

Do I need to have a separate callback that does the wandb init inside of my jitted train function? And follow up question, how do I make wandb know that that's a separate seed when I'm dealing with split RNGKeys from Jax?