Some bugs in model inference

[hdadong]

Hi, I noticed a potential parameter order mismatch when calling reward_fn in the training logic.

In ssrl/brax/training/agents/ssrl/networks.py#L133 the current implementation is:

reward = c.reward_fn(obs, obs_next, jp.mean(us), action)

However, the reward_fn definition in go1_go_fast.py` shows the expected parameter order as:

def reward_fn(obs_next, obs, us, action):

The parameters obs and obs_next should be swapped in the function call.

Additionally, I think it should use the clean obs (obs_next_mean) to compute rewards instead of using the noise version (obs_next), so it should be like this:
Proposed Fix:

reward = c.reward_fn(obs_next_mean, obs, jp.mean(us), action)

[hdadong]

And I noticed a bug with the reward computation related to the torque:

• The low_level_control function returns the scaled action, not the actual torques:
  https://github.com/CLeARoboticsLab/ssrl/blob/main/ssrl/brax/envs/go1_go_fast.py#L1041

• During model inference:
  https://github.com/CLeARoboticsLab/ssrl/blob/main/ssrl/brax/training/agents/ssrl/networks.py#L122C23-L122C40
  the variable u refers to the scaled action, not the torque.

• However, in the reward function, u is directly passed as input to _reward_normalized:
  https://github.com/CLeARoboticsLab/ssrl/blob/main/ssrl/brax/envs/go1_go_fast.py#L877
  This suggests that the reward is computed using scaled actions rather than the actual torques, which seems to be a bug.