Environment inverted_pendulum never returns done=True

[olivier-serris]

Context

I am trying to use the inverted_pendulum environment in Brax 0.0.12 with google colab. Even after 500 timesteps, the environnement never returns done (unless time limit is reached).

Code

try:
  import brax
except ImportError:
  from IPython.display import clear_output 
  !pip install git+https://github.com/google/brax.git@main
  clear_output()
  import brax
from brax import jumpy as jp
from brax import envs

environment = "inverted_pendulum"
env = envs.create(env_name=environment, auto_reset=True, legacy_spring=False)
state = env.reset(rng=jp.random_prngkey(seed=0))
rollout = []
for i in range(500):
    action = jp.zeros((env.action_size,))
    state = env.step(state, action)
    rollout.append(state)
print('brax version ', brax.__version__)
print('any done :', jp.sum([state.done for state in rollout]))

result :

brax version  0.0.12
any done : 0.0

[Markus28]

Some info for debugging: This problem doesn't occur with legacy_spring=True. Looking at a rollout, the hinge joint doesn't move at all. Same story for inverted_double_pendulum.

[jkterry1]

Is the default argument legacy_spring=True?

[olivier-serris]

@jkterry1 The default arg is legacy_spring=False where the hinge joint does not moves. As @Markus28 said, the problem does not happen with legacy_spring=True .

But the release notes from 0.011 made me believe that all brax environment would work with both modes: https://github.com/google/brax/releases/tag/v0.0.11

[cdfreeman-google]

Yeah this was a very silly bug. In the legacy_spring=True codepath, the joint has angle limits of 0 (i.e., a fixed joint), BUT, there's a limit_strength: 0.0 parameter which sets the strength of the restorative angle-limit force. In the new code, there's no such thing (because angle limits are exactly enforced, no spring needed). But this meant the new code was making the inverted pendulum perfectly rigid...

We just added sensible angle limits to the joint. Should make its way live shortly.

[erikfrey]

Fixed in c7735f34a48c3499516c3359d016057ed653f810.

[olivier-serris]

Thanks !

[FranzKnut]

Sorry for commenting on an old Issue but I have run into the same problem.

I am using brax 0.1.2 and inverted_pendulum never returns done=True (although it seems to be resetting). Setting legacy_spring=True did not help unfortunately. Any help is appreciated!