Need Help to execute 'continue learning' using trained SAC model with --load_checkpoint_path tag

Hello All,

I new to RL. I created a new script, `train_jax_sac.py` (which I attached), to run the environment using SAC in a setup similar to PPO. I could successfully run a fresh training using this file but have trouble using it with `--load_checkpoint_path ` tag though it was able to get the reward for the 1st eval.

ERROR-1: This was the first issue I encountered. I resolved it by adding the following snippet inside `.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py`, similar to what’s done in PPO :
`training_state = jax.device_put_replicated(
training_state, jax.local_devices()[:local_devices_to_use]
)`.

However, this fix led to ERROR-2— and on the contrary now it failed even when starting a fresh run.

ERROR-2: After continuing with the above change, I encountered a second issue. Even though I obtained a reward for the first evaluation, the training got stuck afterward, which made me suspect that my approach might not be correct.

Could you please review this and let me know if I’m missing something? If my approach is incorrect, I’d appreciate your guidance on how to correctly adapt SAC to run in the same setup as PPO.

I’ve attached my modified file, `train_jax_ppo.py`, which served as the reference for creating `train_jax_sac.py`.
[train_jax_sac.py](https://github.com/user-attachments/files/22887845/train_jax_sac.py)

--------------------------------------------------------------------------------------------
ERROR-1: `(vrodrigues) [vrodrigues@anode03 mujoco_playground]$ python learning/train_jax_sac.py --env_name AcrobotSwingup --suffix="test_sac_10M" --load_checkpoint_path=logs/AcrobotSwingup-20251013-124033-test_sac_10M/checkpoints/
Environment Config:
action_repeat: 1
ctrl_dt: 0.01
episode_length: 1000
impl: jax
nconmax: 0
njmax: 0
sim_dt: 0.01
vision: false

SAC Training Parameters:
action_repeat: 1
batch_size: 512
discounting: 0.99
episode_length: 1000
grad_updates_per_step: 8
learning_rate: 0.001
max_replay_size: 4194304
min_replay_size: 8192
network_factory:
q_network_layer_norm: true
normalize_observations: true
num_envs: 128
num_evals: 10
num_timesteps: 10000000
reward_scaling: 1.0

Experiment name: AcrobotSwingup-20251013-125002-test_sac_10M
Logs are being stored in: /home/vrodrigues/mujoco_playground/logs/AcrobotSwingup-20251013-125002-test_sac_10M
Restoring from: /home/vrodrigues/mujoco_playground/logs/AcrobotSwingup-20251013-124033-test_sac_10M/checkpoints/000010000640
Checkpoint path: /home/vrodrigues/mujoco_playground/logs/AcrobotSwingup-20251013-125002-test_sac_10M/checkpoints
Traceback (most recent call last):
File "/home/vrodrigues/mujoco_playground/learning/train_jax_sac.py", line 439, in
app.run(main)
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/absl/app.py", line 316, in run
_run_main(main, args)
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/absl/app.py", line 261, in _run_main
sys.exit(main(argv))
^^^^^^^^^^
File "/home/vrodrigues/mujoco_playground/learning/train_jax_sac.py", line 364, in main
make_inference_fn, params, _ = train_fn(
^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 540, in train
_unpmap(
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 70, in _unpmap
return jax.tree_util.tree_map(lambda x: x[0], v)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in tree_map
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 70, in
return jax.tree_util.tree_map(lambda x: x[0], v)
~^^^
IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed`

--------------------------------------------------------------------------------------------
ERROR=2:
(vrodrigues) [vrodrigues@anode03 mujoco_playground]$ python learning/train_jax_sac.py --env_name AcrobotSwingup --suffix="test_sac_10M_continue" --load_checkpoint_path=logs/AcrobotSwingup-20251013-12403
3-test_sac_10M/checkpoints/
Environment Config:
action_repeat: 1
ctrl_dt: 0.01
episode_length: 1000
impl: jax
nconmax: 0
njmax: 0
sim_dt: 0.01
vision: false

SAC Training Parameters:
action_repeat: 1
batch_size: 512
discounting: 0.99
episode_length: 1000
grad_updates_per_step: 8
learning_rate: 0.001
max_replay_size: 4194304
min_replay_size: 8192
network_factory:
q_network_layer_norm: true
normalize_observations: true
num_envs: 128
num_evals: 10
num_timesteps: 10000000
reward_scaling: 1.0

Experiment name: AcrobotSwingup-20251013-165516-test_sac_10M_continue
Logs are being stored in: /home/vrodrigues/mujoco_playground/logs/AcrobotSwingup-20251013-165516-test_sac_10M_continue
Restoring from: /home/vrodrigues/mujoco_playground/logs/AcrobotSwingup-20251013-124033-test_sac_10M/checkpoints/000010000640
Checkpoint path: /home/vrodrigues/mujoco_playground/logs/AcrobotSwingup-20251013-165516-test_sac_10M_continue/checkpoints
0: reward=139.953
Traceback (most recent call last):
File "/home/vrodrigues/mujoco_playground/learning/train_jax_sac.py", line 439, in
app.run(main)
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/absl/app.py", line 316, in run
_run_main(main, args)
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/absl/app.py", line 261, in _run_main
sys.exit(main(argv))
^^^^^^^^^^
File "/home/vrodrigues/mujoco_playground/learning/train_jax_sac.py", line 364, in main
make_inference_fn, params, _ = train_fn(
^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 571, in train
training_epoch_with_timing(
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 455, in training_epoch_with_timing
(training_state, env_state, buffer_state, metrics) = training_epoch(
^^^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 435, in training_epoch
(training_state, env_state, buffer_state, key), metrics = jax.lax.scan(
^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 432, in f
ts, es, bs, metrics = training_step(ts, es, bs, k)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 380, in training_step
(training_state, _), metrics = jax.lax.scan(
^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 277, in sgd_step
critic_loss, q_params, q_optimizer_state = critic_update(
^^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/gradients.py", line 62, in f
value, grads = loss_and_pgrad_fn(*args)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/gradients.py", line 31, in h
value, grad = g(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/agents/sac/losses.py", line 73, in critic_loss
q_old_action = q_network.apply(
^^^^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/networks.py", line 433, in apply
return q_module.apply(q_params, obs, actions)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/networks.py", line 420, in __call__
q = MLP(
^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/brax/training/networks.py", line 96, in __call__
hidden = linen.Dense(
^^^^^^^^^^^^
File "/home/vrodrigues/.venv/lib/python3.11/site-packages/flax/linen/linear.py", line 264, in __call__
kernel = self.param(
^^^^^^^^^^^
flax.errors.ScopeParamShapeError: Initializer expected to generate shape (1, 7, 256) but got shape (7, 256) instead for parameter "kernel" in "/MLP_0/hidden_0". (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ScopeParamShapeError)

For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.