Kinetix/kinetix/environment/ued/util.py

import math
from functools import partial

import jax
import jax.numpy as jnp

from jax2d.engine import PhysicsEngine, calculate_collision_matrix, recalculate_mass_and_inertia, select_shape
from jax2d.sim_state import RigidBody, Thruster
from kinetix.environment.env_state import EnvParams, EnvState, StaticEnvParams


def sample_dimensions(rng, static_env_params: StaticEnvParams, is_rect: bool, ued_params, max_shape_size=None):
    if max_shape_size is None:
        max_shape_size = static_env_params.max_shape_size
    # Returns (half_dimensions, radius)

    rng, _rng = jax.random.split(rng)
    # Don't want overly small shapes
    min_rect_size = 0.05
    min_circle_size = 0.1
    cap_rect = max_shape_size / 2.0 / jnp.sqrt(2.0)
    cap_circ = max_shape_size / 2.0 * ued_params.circle_max_size_coeff
    half_dimensions = (
        jax.lax.select(is_rect, jax.random.uniform(_rng, shape=(2,)), jnp.zeros(2, dtype=jnp.float32))
        * (cap_rect - min_rect_size)
        + min_rect_size
    )

    rng, _rng, __rng = jax.random.split(rng, 3)
    dim_scale = (
        jnp.ones(2)
        .at[jax.random.randint(_rng, shape=(), minval=0, maxval=2)]
        .set(
            jax.lax.select(
                jax.random.uniform(__rng) < ued_params.large_rect_dim_chance, ued_params.large_rect_dim_scale, 1.0
            )
        )
    )
    half_dimensions *= dim_scale

    vertices = jnp.array(
        [
            half_dimensions * jnp.array([1, 1]),
            half_dimensions * jnp.array([1, -1]),
            half_dimensions * jnp.array([-1, -1]),
            half_dimensions * jnp.array([-1, 1]),
        ]
    )

    rng, _rng = jax.random.split(rng)
    radius = (
        jax.lax.select(is_rect, jnp.zeros((), dtype=jnp.float32), jax.random.uniform(_rng, shape=()))
        * (cap_circ - min_circle_size)
        + min_circle_size
    )
    return vertices, half_dimensions, radius


def count_roles(state: EnvState, static_env_params: StaticEnvParams, role: int, include_static_polys=True) -> int:
    active_to_use = state.polygon.active
    if not include_static_polys:
        active_to_use = active_to_use.at[: static_env_params.num_static_fixated_polys].set(False)
    return ((state.polygon_shape_roles == role) * active_to_use).sum() + (
        (state.circle_shape_roles == role) * state.circle.active
    ).sum()


def random_position_on_triangle(rng, vertices):
    verts = vertices[:3]
    rng, _rng, _rng2 = jax.random.split(rng, 3)
    f1 = jax.random.uniform(_rng)
    f2 = jax.random.uniform(_rng2)
    # https://www.reddit.com/r/godot/comments/mqp29g/how_do_i_get_a_random_position_inside_a_collision/
    return verts[0] + jnp.sqrt(f1) * (-verts[0] + verts[1] + f2 * (verts[2] - verts[1]))


def random_position_on_rectangle(rng, vertices):
    verts = vertices[:4]
    rng, _rng, _rng2 = jax.random.split(rng, 3)
    f1 = jax.random.uniform(_rng)
    f2 = jax.random.uniform(_rng2)

    min_x, max_x = jnp.min(verts[:, 0]), jnp.max(verts[:, 0])
    min_y, max_y = jnp.min(verts[:, 1]), jnp.max(verts[:, 1])
    random_x_pos = min_x + f1 * (max_x - min_x)
    random_y_pos = min_y + f2 * (max_y - min_y)

    return jnp.array([random_x_pos, random_y_pos])


def random_position_on_polygon(rng, vertices, n_vertices, static_env_params: StaticEnvParams):
    assert static_env_params.max_polygon_vertices <= 4, "Only supports up to 4 vertices"
    return jax.lax.select(
        n_vertices <= 3, random_position_on_triangle(rng, vertices), random_position_on_rectangle(rng, vertices)
    )


def random_position_on_circle(rng, radius, on_centre_chance):
    rngs = jax.random.split(rng, 3)

    on_centre = jax.random.uniform(rngs[0]) < on_centre_chance

    local_joint_position_circle_theta = jax.random.uniform(rngs[1], shape=()) * 2 * math.pi
    local_joint_position_circle_r = jax.random.uniform(rngs[2], shape=()) * radius
    local_joint_position_circle = jnp.array(
        [
            local_joint_position_circle_r * jnp.cos(local_joint_position_circle_theta),
            local_joint_position_circle_r * jnp.sin(local_joint_position_circle_theta),
        ]
    )

    return jax.lax.select(on_centre, jnp.array([0.0, 0.0]), local_joint_position_circle)


def get_role(rng, state: EnvState, static_env_params: StaticEnvParams, initial_p=None) -> int:

    if initial_p is None:
        initial_p = jnp.array([1.0, 1.0, 1.0, 1.0])

    needs_ball = count_roles(state, static_env_params, 1) == 0
    needs_goal = count_roles(state, static_env_params, 2) == 0
    needs_lava = count_roles(state, static_env_params, 3) == 0

    # always put goal/ball first.
    prob_of_something_else = (needs_ball == 0) & (needs_goal == 0)
    p = initial_p * jnp.array(
        [prob_of_something_else, needs_ball, needs_goal, prob_of_something_else * needs_lava / 3]
    )  # This ensures we cannot more than one ball or goal.
    return jax.random.choice(rng, jnp.array([0, 1, 2, 3]), p=p)


def is_space_for_shape(state: EnvState):
    return jnp.logical_not(jnp.concatenate([state.polygon.active, state.circle.active])).sum() > 0


def is_space_for_joint(state: EnvState):
    return jnp.logical_not(state.joint.active).sum() > 0


def are_there_shapes_present(state: EnvState, static_env_params: StaticEnvParams):
    m = (
        jnp.concatenate([state.polygon.active, state.circle.active])
        .at[: static_env_params.num_static_fixated_polys]
        .set(False)
    )
    return m.sum() > 0


@partial(jax.jit, static_argnums=(2, 9))
def add_rigidbody_to_state(
    state: EnvState,
    env_params: EnvParams,
    static_env_params: StaticEnvParams,
    position: jnp.ndarray,
    vertices: jnp.ndarray,
    n_vertices: int,
    radius: float,
    shape_role: int,
    density: float = 1,
    is_circle: bool = False,
):

    new_rigid_body = RigidBody(
        position=position,
        velocity=jnp.array([0.0, 0.0]),
        inverse_mass=1.0,
        inverse_inertia=1.0,
        rotation=0.0,
        angular_velocity=0.0,
        radius=radius,
        active=True,
        friction=1.0,
        vertices=vertices,
        n_vertices=n_vertices,
        collision_mode=1,
        restitution=0.0,
    )

    if is_circle:
        actives = state.circle.active
    else:
        actives = state.polygon.active

    idx = jnp.argmin(actives)

    def noop(state):
        return state

    def replace(state):
        add_func = lambda all, new: all.at[idx].set(new)
        if is_circle:
            state = state.replace(
                circle=jax.tree.map(add_func, state.circle, new_rigid_body),
                circle_densities=state.circle_densities.at[idx].set(density),
                circle_shape_roles=state.circle_shape_roles.at[idx].set(shape_role),
            )
        else:
            state = state.replace(
                polygon=jax.tree.map(add_func, state.polygon, new_rigid_body),
                polygon_densities=state.polygon_densities.at[idx].set(density),
                polygon_shape_roles=state.polygon_shape_roles.at[idx].set(shape_role),
            )

        state = state.replace(
            collision_matrix=calculate_collision_matrix(static_env_params, state.joint),
        )

        state = recalculate_mass_and_inertia(state, static_env_params, state.polygon_densities, state.circle_densities)
        return state

    return jax.lax.cond(jnp.logical_not(actives).sum() > 0, replace, noop, state)


def rectangle_vertices(half_dim):
    return jnp.array(
        [
            half_dim * jnp.array([1, 1]),
            half_dim * jnp.array([1, -1]),
            half_dim * jnp.array([-1, -1]),
            half_dim * jnp.array([-1, 1]),
        ]
    )


# More Manual Control
@partial(jax.jit, static_argnums=(2,))
def add_rectangle_to_state(
    state: EnvState,
    env_params: EnvParams,
    static_env_params: StaticEnvParams,
    position: jnp.ndarray,
    width: float,
    height: float,
    shape_role: int,
    density: float = 1,
):

    return add_rigidbody_to_state(
        state,
        env_params,
        static_env_params,
        position,
        rectangle_vertices(jnp.array([width, height]) / 2),
        4,
        0.0,
        shape_role,
        density,
        is_circle=False,
    )


@partial(jax.jit, static_argnums=(2,))
def add_circle_to_state(
    state: EnvState,
    env_params: EnvParams,
    static_env_params: StaticEnvParams,
    position: jnp.ndarray,
    radius: float,
    shape_role: int,
    density: float = 1,
):
    return add_rigidbody_to_state(
        state,
        env_params,
        static_env_params,
        position,
        jnp.array([0.0, 0.0]),
        0,
        radius,
        shape_role,
        density,
        is_circle=True,
    )


@partial(jax.jit, static_argnums=(2,))
def add_thruster_to_object(
    state: EnvState,
    env_params: EnvParams,
    static_env_params: StaticEnvParams,
    shape_index: int,
    rotation: float,
    colour: int,
    thruster_power_multiplier: float,
):
    def dummy(state):
        return state

    def do_add(state: EnvState):
        thruster_idx = jnp.argmin(state.thruster.active)

        shape = select_shape(state, shape_index, static_env_params)

        thruster = Thruster(
            object_index=shape_index,
            active=True,
            relative_position=jnp.array([0.0, 0.0]),  # a bit of a hack but reasonable.
            rotation=rotation,
            power=1.0 / jax.lax.select(shape.inverse_mass == 0, 1.0, shape.inverse_mass) * thruster_power_multiplier,
            global_position=select_shape(state, shape_index, static_env_params).position,
        )

        state = state.replace(
            thruster=jax.tree_map(lambda y, x: y.at[thruster_idx].set(x), state.thruster, thruster),
            thruster_bindings=state.thruster_bindings.at[thruster_idx].set(colour),
        )

        return state

    return jax.lax.cond(
        (select_shape(state, shape_index, static_env_params).active)
        & (jnp.logical_not(state.thruster.active).sum() > 0),
        do_add,
        dummy,
        state,
    )


def make_velocities_zero(state: EnvState):
    def inner(state):
        return state.replace(
            polygon=state.polygon.replace(
                angular_velocity=state.polygon.angular_velocity * 0,
                velocity=state.polygon.velocity * 0,
            ),
            circle=state.circle.replace(
                angular_velocity=state.circle.angular_velocity * 0,
                velocity=state.circle.velocity * 0,
            ),
        )

    return inner(state)


def make_do_dummy_step(
    params: EnvParams, static_sim_params: StaticEnvParams, zero_collisions=True, zero_velocities=True
):
    env = PhysicsEngine(static_sim_params)

    @jax.jit
    def _step_fn(state):
        state, _ = env.step(state, params, jnp.zeros((static_sim_params.num_joints + static_sim_params.num_thrusters,)))
        return state

    def do_dummy_step(state: EnvState) -> EnvState:
        rng = jax.random.PRNGKey(0)
        og_col = state.collision_matrix
        g = state.gravity
        state = state.replace(
            collision_matrix=state.collision_matrix & (not zero_collisions), gravity=state.gravity * 0
        )
        state = _step_fn(state)
        state = state.replace(gravity=g, collision_matrix=og_col)
        if zero_velocities:
            state = make_velocities_zero(state)
        return state

    return do_dummy_step