Auto reset in Environment.step() means CartPole always returns reward = 1

[garymm]

CartPole has this code which only sets reward = 0 if the previous state was terminal:

gymnax/gymnax/environments/classic_control/cartpole.py

Lines 83 to 84 in aef77d5

	# Important: Reward is based on termination is previous step transition
	reward = 1.0 - prev_terminal

Environment.step has this code which means a terminal state is never passed in to CartPole.step_env:

gymnax/gymnax/environments/environment.py

Lines 45 to 50 in aef77d5

	obs_st, state_st, reward, done, info = self.step_env(key, state, action, params)
	obs_re, state_re = self.reset_env(key_reset, params)
	# Auto-reset environment based on termination
	state = jax.tree_map(
	lambda x, y: jax.lax.select(done, x, y), state_re, state_st
	)

Hence CartPole.step() always returns 1 which makes it entirely broken.

[samuelstevens]

I am experiencing the same issue.

[Michael190502]

Hi, I was looking into the same thing and I think there should not be a realistic scenario in which this zero reward is needed as it would only happen if you call step on an already terminated environment. This can also be seen in lines 212 and following of the gymnasium cartpole implementation.