EpisodeWrapper only preserves the last sub‑step’s metrics when using action_repeat
- Dominant language
- Jupyter Notebook
- Stars
- 3.2k
- Forks
- 349
- PR merge metrics
- No merged PRs in 30d
Description
In brax/envs/wrappers/training.py the EpisodeWrapper.step method correctly uses lax.scan to collect and sum rewards over action_repeat sub‑steps, but it never accumulates the corresponding per‑step state.metrics. After the scan it does:
```
# only sums rewards:
state = state.replace(reward=jp.sum(rewards, axis=0))
# then aggregates episode_metrics from state.metrics, but
# state.metrics here is just the metrics dict from the *last* sub‑step
for metric_name in state.metrics.keys():
if metric_name != "reward":
state.info["episode_metrics"][metric_name] += state.metrics[metric_name]
…
```
Because state.metrics has already been overwritten by the final call to self.env.step, all earlier sub‑step metrics get dropped—so any sparse or per‑step metric (like an action‐change penalty) will only ever reflect the last sub‑step’s value. I found this out when logging the metrics on a tensorboard and verifying that sparse rewards were always zero.
**What I Expected**
Just like rewards, all per‑step metrics should be summed across the action_repeat sub‑steps before being written back into state.metrics (and into state.info["episode_metrics"]).
**Suggested fix (draft)**
Modify the lax.scan loop to return and accumulate a pair (reward, metrics):
```
def f(carry, _):
st, metrics_acc = carry
nst = self.env.step(st, action)
new_metrics_acc = tree_map(lambda acc, m: acc + m, metrics_acc, nst.metrics)
return (nst, new_metrics_acc), nst.reward
# then scan over (state, zeros_metrics) and (), sum both rewards and metrics
```
I've gotten it to work on my machine, so if you think this functionality makes sense I can create a pull request. Thank you in advance!
Contributor guide
Assessment
This issue has not been assessed yet.