Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Correctly handle non-scalar rewards for dm_env - #162

Merged
Trenza1ore merged 6 commits into
Farama-Foundation:mainfrom
teddytennant:wrangle-sh-a
Sep 12, 2026
Merged

Correctly handle non-scalar rewards for dm_env#162
Trenza1ore merged 6 commits into
Farama-Foundation:mainfrom
teddytennant:wrangle-sh-a

Conversation

@teddytennant

Copy link
Copy Markdown
Contributor

dm_env_step2gym_step used timestep.reward or 0. That turns a real 0.0 into int 0, and a vector reward raises ValueError: The truth value of an array with more than one element is ambiguous.

Swap it for is not None. FIRST timesteps still become 0; 0.0 stays a float; arrays pass through.

The three cases live in tests/test_dm_control.py next to the DiscreteArray dtype test from #157.

I ran the three conversions against dm_env.TimeStep here (no dm-control extra on this box). black, isort --profile black, and pydocstyle --convention=google are clean on the two files.

Comment thread shimmy/utils/dm_env.py Outdated
"""
obs = dm_obs2gym_obs(timestep.observation)
reward = timestep.reward or 0
reward = timestep.reward if timestep.reward is not None else 0

@Trenza1ore Trenza1ore Sep 12, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
reward = timestep.reward if timestep.reward is not None else 0
if hasattr(timestep.reward, "__getitem__") and hasattr(timestep.reward, "__len__"):
if len(timestep.reward) == 1:
reward = timestep.reward[0]
else:
raise TypeError(
f"Gymnasium only supports scalar reward, got {timestep.reward}"
)
else:
reward = timestep.reward or 0

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gymnasium doesn't support vector rewards, so we should handle them properly and reject any reward with a length other than 1.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

length-1 arrays unwrap, anything else raises TypeError. kept is not None so 0.0 stays a float.

@Trenza1ore

Copy link
Copy Markdown
Member

Well, the issue is that Gymnasium environments only support reward values of type SupportsFloat (so, no vector, only scalar), therefore it only make sense to support a NumPy array when it has a length of exactly 1. Please take a look at the review comment above.

Trenza1ore and others added 5 commits September 12, 2026 02:00
Gymnasium only accepts scalar SupportsFloat rewards. Unwrap a 1-d
length-1 array, pass through 0-d numpy values, and raise TypeError for
other vectors including shape (1, 2). Keep is not None so 0.0 stays float.
This should be a few times faster for all cases
@Trenza1ore Trenza1ore changed the title Use is not None when converting a dm_env reward Correctly handle non-scalar rewards for dm_env Sep 12, 2026
@Trenza1ore

Copy link
Copy Markdown
Member

np.ndim is quite expensive and it's not super desirable to return a 0-d np.ndarray to user, so I've made some further adjustments, should be the lowest latency we can get while somewhat handling the non-scalar reward types properly.

@Trenza1ore
Trenza1ore merged commit 8f72ed3 into Farama-Foundation:main Sep 12, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants