Ant (positional) initialization often inevitably leads to unhealthy state / early termination with a few timesteps about brax

[btaba]

Hi @DramaCow , I'm not able to reproduce your issue. Here's what I ran:

env = envs.create(env_name=env_name, backend='positional')

jit_env_reset = jax.jit(jax.vmap(env.reset))
jit_env_step = jax.jit(jax.vmap(env.step))

rollout = []
rng = jax.random.PRNGKey(seed=1)
rng = jax.random.split(rng, 1024)
state = jit_env_reset(rng=rng)
for _ in range(1000):
  rollout.append(state.pipeline_state)
  state = jit_env_step(state, jp.zeros((1024, env.sys.act_size())))


def unhealthy(state):
    pipeline_state = state.pipeline_state
    min_z, max_z = env._healthy_z_range
    is_healthy = jp.where(pipeline_state.x.pos[0, 2] < min_z, x=0.0, y=1.0)
    is_healthy = jp.where(
        pipeline_state.x.pos[0, 2] > max_z, x=0.0, y=is_healthy
    )
    return (1 - is_healthy).astype(jp.bool_)


jax.vmap(unhealthy)(state).sum()

I get 0 at the end