google / google/brax

EpisodeWrapper only preserves the last sub‑step’s metrics when using action_repeat

Open
#610 1 comment 0 reactions 1 assignee Claimed by @btaba View on GitHub
bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.