Spaces:
Runtime error
Runtime error
| import chex | |
| import jax | |
| from jax2d.engine import calculate_collision_matrix | |
| from kinetix.environment.env_state import EnvState, StaticEnvParams | |
| import jax.numpy as jnp | |
| from kinetix.pcg.pcg_state import PCGState | |
| def permute_state(rng: chex.PRNGKey, env_state: EnvState, static_env_params: StaticEnvParams): | |
| idxs_circles = jnp.arange(static_env_params.num_circles) | |
| idxs_polygons = jnp.arange(static_env_params.num_polygons) | |
| idxs_joints = jnp.arange(static_env_params.num_joints) | |
| idxs_thrusters = jnp.arange(static_env_params.num_thrusters) | |
| rng, *_rngs = jax.random.split(rng, 5) | |
| idxs_circles_permuted = jax.random.permutation(_rngs[0], idxs_circles, independent=True) | |
| idxs_polygons_permuted = idxs_polygons.at[static_env_params.num_static_fixated_polys :].set( | |
| jax.random.permutation(_rngs[1], idxs_polygons[static_env_params.num_static_fixated_polys :], independent=True) | |
| ) | |
| idxs_joints_permuted = jax.random.permutation(_rngs[2], idxs_joints, independent=True) | |
| idxs_thrusters_permuted = jax.random.permutation(_rngs[3], idxs_thrusters, independent=True) | |
| combined = jnp.concatenate([idxs_polygons_permuted, idxs_circles_permuted + static_env_params.num_polygons]) | |
| # Change the ordering of the shapes, and also remember to change the indices associated with the joints | |
| inverse_permutation = jnp.argsort(combined) | |
| env_state = env_state.replace( | |
| polygon_shape_roles=env_state.polygon_shape_roles[idxs_polygons_permuted], | |
| circle_shape_roles=env_state.circle_shape_roles[idxs_circles_permuted], | |
| polygon_highlighted=env_state.polygon_highlighted[idxs_polygons_permuted], | |
| circle_highlighted=env_state.circle_highlighted[idxs_circles_permuted], | |
| polygon_densities=env_state.polygon_densities[idxs_polygons_permuted], | |
| circle_densities=env_state.circle_densities[idxs_circles_permuted], | |
| polygon=jax.tree.map(lambda x: x[idxs_polygons_permuted], env_state.polygon), | |
| circle=jax.tree.map(lambda x: x[idxs_circles_permuted], env_state.circle), | |
| joint=env_state.joint.replace( | |
| a_index=inverse_permutation[env_state.joint.a_index], | |
| b_index=inverse_permutation[env_state.joint.b_index], | |
| ), | |
| thruster=env_state.thruster.replace( | |
| object_index=inverse_permutation[env_state.thruster.object_index], | |
| ), | |
| ) | |
| # And now permute the thrusters and joints | |
| env_state = env_state.replace( | |
| thruster_bindings=env_state.thruster_bindings[idxs_thrusters_permuted], | |
| motor_bindings=env_state.motor_bindings[idxs_joints_permuted], | |
| motor_auto=env_state.motor_auto[idxs_joints_permuted], | |
| joint=jax.tree.map(lambda x: x[idxs_joints_permuted], env_state.joint), | |
| thruster=jax.tree.map(lambda x: x[idxs_thrusters_permuted], env_state.thruster), | |
| ) | |
| # and collision matrix | |
| env_state = env_state.replace(collision_matrix=calculate_collision_matrix(static_env_params, env_state.joint)) | |
| return env_state | |
| def permute_pcg_state(rng: chex.PRNGKey, pcg_state: PCGState, static_env_params: StaticEnvParams): | |
| return pcg_state.replace( | |
| env_state=permute_state(rng, pcg_state.env_state, static_env_params), | |
| env_state_max=permute_state(rng, pcg_state.env_state_max, static_env_params), | |
| env_state_pcg_mask=jax.tree.map(lambda x: jnp.zeros_like(x, dtype=bool), pcg_state.env_state_pcg_mask), | |
| ) | |