Spaces:
Runtime error
Runtime error
File size: 3,433 Bytes
581eeac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
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),
)
|