import os
os.system(f"cd {os.getcwd()}/Kinetix")
import itertools
import time
from timeit import default_timer as tmr
import optax
from PIL import Image
from flax.serialization import to_state_dict
from flax.training.train_state import TrainState
from matplotlib import pyplot as plt
from kinetix.environment.ued.distributions import sample_kinetix_level
from kinetix.models import make_network_from_config
from kinetix.models.actor_critic import ScannedRNN
from kinetix.render.renderer_symbolic_entity import make_render_entities
ss = tmr()
import jax
jax.config.update("jax_compilation_cache_dir", ".cache-location")
import hydra
from omegaconf import OmegaConf
from kinetix.environment.ued.mutators import (
make_mutate_change_shape_rotation,
mutate_add_connected_shape,
mutate_add_shape,
make_mutate_change_shape_size,
mutate_change_shape_location,
mutate_swap_role,
mutate_remove_shape,
mutate_remove_joint,
mutate_toggle_fixture,
mutate_add_thruster,
mutate_remove_thruster,
)
from kinetix.environment.ued.ued import make_mutate_env, ALL_MUTATION_FNS
from kinetix.environment.ued.ued_state import UEDParams
from kinetix.environment.ued.util import rectangle_vertices
from kinetix.util.config import generate_params_from_config, normalise_config
import argparse
import os
import sys
sys.path.append("editor")
import tkinter
import tkinter.filedialog
from enum import Enum
from timeit import default_timer as tmr
import jax.numpy as jnp
import numpy as np
import pygame
import pygame_widgets
from pygame_widgets.slider import Slider
from pygame_widgets.textbox import TextBox
from pygame_widgets.toggle import Toggle
from jax2d.engine import (
calc_inverse_inertia_circle,
calc_inverse_inertia_polygon,
calc_inverse_mass_circle,
calc_inverse_mass_polygon,
calculate_collision_matrix,
recalculate_mass_and_inertia,
recompute_global_joint_positions,
select_shape,
)
from jax2d.maths import rmat
from jax2d.sim_state import RigidBody
from kinetix.environment.env import (
create_empty_env,
make_kinetix_env_from_name,
)
from kinetix.environment.env_state import EnvParams, EnvState, StaticEnvParams
from kinetix.environment.utils import permute_pcg_state
from kinetix.environment.wrappers import AutoResetWrapper
from kinetix.pcg.pcg import env_state_to_pcg_state, sample_pcg_state
from kinetix.pcg.pcg_state import PCGState
from kinetix.render.renderer_pixels import make_render_pixels
from kinetix.render.textures import (
CIRCLE_TEXTURE_RGBA,
EDIT_TEXTURE_RGBA,
PLAY_TEXTURE_RGBA,
RECT_TEXTURE_RGBA,
RJOINT_TEXTURE_RGBA,
SELECT_TEXTURE_RGBA,
THRUSTER_TEXTURE_RGBA,
TRIANGLE_TEXTURE_RGBA,
)
from kinetix.util.saving import (
expand_pcg_state,
export_env_state_to_json,
get_pcg_state_from_json,
load_from_json_file,
load_pcg_state_pickle,
load_world_state_pickle,
save_pickle,
load_params_from_wandb_artifact_path,
load_train_state_from_wandb_artifact_path,
)
from kinetix.util.timing import time_function
from tkinter import Tk
#root = Tk()
#root.destroy()
ee = tmr()
print(f"Imported in {ee - ss} seconds")
editor = None
outer_timer = tmr()
EMPTY_ENV = False
class ObjectType(Enum):
POLYGON = 0
CIRCLE = 1
JOINT = 2
THRUSTER = 3
class EditMode(Enum):
ADD_CIRCLE = 0
ADD_RECTANGLE = 1
ADD_JOINT = 2
SELECT = 3
ADD_TRIANGLE = 4
ADD_THRUSTER = 5
TOTAL_DUMMY_STEPS_TO_SNAP = 0
SNAPPING_DIST = 0.1
def select_object(state: EnvState, type: int, index: int):
if type is None:
type = ObjectType.POLYGON
li = {0: state.polygon, 1: state.circle, 2: state.joint, 3: state.thruster}[type.value]
return jax.tree.map(lambda x: x[index], li)
def snap_to_center(shape: RigidBody, position: jnp.ndarray):
if jnp.linalg.norm(shape.position - position) < SNAPPING_DIST:
return shape.position
return position
def snap_to_polygon_center_line(polygon: RigidBody, position: jnp.ndarray):
# Snap to the center line
r = rmat(polygon.rotation)
x = jnp.matmul(r, position - polygon.position)
if jnp.abs(x[0]) < SNAPPING_DIST:
x = x.at[0].set(0.0)
if jnp.abs(x[1]) < SNAPPING_DIST:
x = x.at[1].set(0.0)
x = jnp.matmul(r.transpose(), x)
return x + polygon.position
def snap_to_circle_center_line(circle: RigidBody, position: jnp.ndarray):
# Snap to the center line, i.e. on the edge of the circle, if the position is close enough to directly below the circle, etc., snap the position to that
x = position - circle.position
if jnp.linalg.norm(x) < SNAPPING_DIST:
return circle.position
angle = (jnp.arctan2(x[1], x[0]) + 2 * jnp.pi) % (2 * jnp.pi)
for i in range(0, 8):
if jnp.abs(angle - i * jnp.pi / 4) < jnp.radians(25): # 25 degrees
angle = i * jnp.pi / 4
break
x = jnp.array([jnp.cos(angle), jnp.sin(angle)]) * circle.radius
return x + circle.position
def prompt_file(save=False):
dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "worlds")
"""Create a Tk file dialog and cleanup when finished"""
top = tkinter.Tk()
top.withdraw() # hide window
if save:
file_name = tkinter.filedialog.asksaveasfilename(parent=top, initialdir=dir)
else:
file_name = tkinter.filedialog.askopenfilename(parent=top, initialdir=dir)
top.destroy()
return file_name
def get_numeric_key_pressed(pygame_events, is_mod=False):
for event in pygame_events:
if not is_mod:
if event.type == pygame.KEYDOWN and event.unicode.isdigit():
return int(event.unicode)
else:
if event.type == pygame.KEYDOWN:
pass
if event.type == pygame.KEYDOWN and event.key in [pygame.K_0 + i for i in range(10)]:
return int(event.key - pygame.K_0)
return None
def new_env(static_env_params):
return create_empty_env(static_env_params)
myrng = jax.random.PRNGKey(0)
def make_reset_function(static_env_params):
def reset(rng):
return env_state_to_pcg_state(create_empty_env(static_env_params))
return reset
def new_pcg_env(static_env_params):
global myrng
if EMPTY_ENV:
env_state = create_empty_env(static_env_params)
else:
return get_pcg_state_from_json("worlds/l/h0_angrybirds.json")
return env_state_to_pcg_state(env_state)
class Editor:
def __init__(self, env, env_params, config, upscale=1):
self.env = env
self.upscale = upscale
self.env_params = env_params
self.static_env_params = env.static_env_params
self.ued_params = UEDParams()
self.side_panel_width = env.static_env_params.screen_dim[0] // 3
self.rng = jax.random.PRNGKey(0)
self.config = config
self.pcg_state = new_pcg_env(self.static_env_params)
self.rng, _rng = jax.random.split(self.rng)
self.play_state = sample_pcg_state(_rng, self.pcg_state, self.env_params, self.static_env_params)
self.last_played_level = None
self.pygame_events = []
self.mutate_world = make_mutate_env(env.static_env_params, env_params, self.ued_params)
self.num_triangle_clicks = 0
self.triangle_order = jnp.array([0, 1, 2])
# Init rendering
pygame.init()
pygame.key.set_repeat(250, 75)
self.screen_surface = pygame.display.set_mode(
tuple(
(t + extra) * self.upscale
for t, extra in zip(self.static_env_params.screen_dim, (self.side_panel_width, 0))
),
display=0,
)
self.all_widgets = {}
self._setup_side_panel()
self.has_done_action = False
self._setup_rendering(self.static_env_params, env_params)
self._render_edit_overlay_fn = jax.jit(self._render_edit_overlay)
self._step_fn = jax.jit(env.step)
# Agent
if self.config["agent_taking_actions"]:
self._entity_renderer = jax.jit(make_render_entities(env_params, self.env.static_env_params))
self.network = make_network_from_config(env, env_params, config)
rng = jax.random.PRNGKey(0)
dones = jnp.zeros((config["num_train_envs"]), dtype=jnp.bool_)
rng, _rng = jax.random.split(rng)
init_hstate = ScannedRNN.initialize_carry(config["num_train_envs"])
obsv = self._entity_renderer(self.play_state)
obsv = jax.tree.map(lambda x: jnp.repeat(x[None, ...], repeats=config["num_train_envs"], axis=0), obsv)
init_x = jax.tree.map(lambda x: x[None, ...], (obsv, dones))
network_params = self.network.init(_rng, init_hstate, init_x)
tx = optax.chain(
optax.clip_by_global_norm(config["max_grad_norm"]),
optax.adam(config["lr"], eps=1e-5),
)
self.train_state = TrainState.create(
apply_fn=self.network.apply,
params=network_params,
tx=tx,
)
self.train_state = load_train_state_from_wandb_artifact_path(self.train_state, config["agent_wandb_path"])
self.apply_fn = jax.jit(self.network.apply)
# JIT Compile
def _jit_step():
rng = jax.random.PRNGKey(0)
ans = self._step_fn(
rng,
self.env.reset_env_to_level(rng, self.play_state, self.env_params)[1],
jnp.zeros(
env.static_env_params.num_motor_bindings + env.static_env_params.num_thruster_bindings, dtype=int
),
self.env_params,
)
def _jit_render():
self._render_fn(self.play_state)
self._render_fn_edit(self.play_state)
time_function(_jit_step, "_jit_step")
time_function(_jit_render, "_jit_render")
# self._step_fn(rng, self.play_state, 0, self.env_params)
# Editing
self.is_editing = True
self.edit_shape_mode = EditMode.ADD_CIRCLE
self.creating_shape = False
self.create_shape_position = jnp.array([0.0, 0.0])
self.creating_shape_index = 0
self.selected_shape_index = -1
self.selected_shape_type = ObjectType.POLYGON
self.all_selected_shapes = []
self.rng = jax.random.PRNGKey(0)
time_function(self._jit, "self._jit")
self._put_state_values_into_gui(self.pcg_state)
self.mutate_change_shape_size = make_mutate_change_shape_size(self.env_params, self.static_env_params)
self.mutate_change_shape_rotation = make_mutate_change_shape_rotation(self.env_params, self.static_env_params)
def _setup_rendering(self, static_env_params: StaticEnvParams, env_params: EnvParams):
def _make_render(should_do_edit_additions=False):
def _render(env_state):
side_panel = self._render_side_panel()
render_pixels = make_render_pixels(params=env_params, static_params=static_env_params)
pixels = render_pixels(
env_state,
)
pixels = jnp.concatenate([side_panel, pixels], axis=0)
pixels = jnp.repeat(pixels, repeats=static_env_params.downscale * self.upscale, axis=0)
pixels = jnp.repeat(pixels, repeats=static_env_params.downscale * self.upscale, axis=1)
return pixels[:, ::-1, :]
return _render
def _make_screenshot_render():
def _render(env_state):
px_upscale = 4
static_params = static_env_params.replace(screen_dim=(500 * px_upscale, 500 * px_upscale))
ss_env_params = env_params.replace(
pixels_per_unit=100 * px_upscale,
)
render_pixels = make_render_pixels(
params=ss_env_params,
static_params=static_params,
render_rjoint_sectors=False,
pixel_upscale=2 * px_upscale,
)
pixels = render_pixels(
env_state,
)
return pixels[:, ::-1, :]
return _render
self._render_fn_edit = jax.jit(_make_render(True))
self._render_fn = jax.jit(_make_render(False))
self._render_fn_screenshot = jax.jit(_make_screenshot_render())
def _jit(self):
self._get_circles_on_mouse(self.pcg_state.env_state)
self._get_polygons_on_mouse(self.pcg_state.env_state)
self._get_revolute_joints_on_mouse(self.pcg_state.env_state)
self._get_thrusters_on_mouse(self.pcg_state.env_state)
self.pygame_events = list(pygame.event.get())
self._handle_events(do_dummy=True)
state = self.play_state
for mutation_fn in ALL_MUTATION_FNS:
mutation_fn(jax.random.PRNGKey(0), state, self.env_params, self.static_env_params, self.ued_params)
def update(self, rng):
# Update pygame events
self.pygame_events = list(pygame.event.get())
for event in self.pygame_events:
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_SPACE:
self.has_done_action = False
self.is_editing = not self.is_editing
if not self.is_editing:
self.pcg_state = self._discard_shape_being_created(self.pcg_state)
self.pcg_state = self._reset_select_shape(self.pcg_state)
self.pcg_state = self.pcg_state.replace(
env_state=self.pcg_state.env_state.replace(
collision_matrix=calculate_collision_matrix(
self.static_env_params, self.pcg_state.env_state.joint
),
),
env_state_pcg_mask=self.pcg_state.env_state_pcg_mask.replace(
collision_matrix=jnp.zeros_like(self.pcg_state.env_state_pcg_mask.collision_matrix)
),
)
self.rng, _rng = jax.random.split(self.rng)
self.play_state = sample_pcg_state(
_rng, self.pcg_state, self.env_params, self.static_env_params
)
self.last_played_level = self.play_state
elif event.key == pygame.K_s and not self.is_editing:
self.take_screenshot()
if self.is_editing:
self.pcg_state = self.edit()
else:
rng, _rng = jax.random.split(rng)
# action = []
action = jnp.zeros(
self.env.static_env_params.num_motor_bindings + self.env.static_env_params.num_thruster_bindings,
dtype=jnp.int32,
)
keys = pygame.key.get_pressed()
if keys[pygame.K_LEFT]:
action = action.at[0].set(1)
if keys[pygame.K_RIGHT]:
action = action.at[0].set(2)
if keys[pygame.K_UP]:
action = action.at[1].set(1)
if keys[pygame.K_DOWN]:
action = action.at[1].set(2)
if keys[pygame.K_1]:
action = action.at[0 + self.env.static_env_params.num_motor_bindings].set(1)
if keys[pygame.K_2]:
action = action.at[1 + self.env.static_env_params.num_motor_bindings].set(1)
if keys[pygame.K_3]:
action = action.at[2 + self.env.static_env_params.num_motor_bindings].set(1)
if keys[pygame.K_4]:
action = action.at[3 + self.env.static_env_params.num_motor_bindings].set(1)
# if self.has_done_action: action = action * 0
self.has_done_action = self.has_done_action | (action != 0).any()
if self.config["agent_taking_actions"]:
obs = self._entity_renderer(self.play_state)
obs = jax.tree.map(lambda x: x[None, ...], obs)
last_done = jnp.zeros((1, 1), dtype=bool)
ac_in = (jax.tree.map(lambda x: x[np.newaxis, :], obs), last_done[np.newaxis, :])
hstate = ScannedRNN.initialize_carry(1)
hstate, pi, value = self.apply_fn(self.train_state.params, hstate, ac_in)
rng, _rng = jax.random.split(rng)
action = pi.sample(seed=_rng)
action = action[0, 0]
_rng, __rng = jax.random.split(_rng)
obs, self.play_state, reward, done, info = self._step_fn(
_rng, self.env.reset_to_level(__rng, self.play_state, self.env_params)[1], action, self.env_params
)
if done:
self.rng, _rng = jax.random.split(self.rng)
self.play_state = sample_pcg_state(_rng, self.pcg_state, self.env_params, self.static_env_params)
state_to_render = self.pcg_state.env_state if self.is_editing else self.play_state
self.render(state_to_render)
self._handle_events()
# Update screen
pygame.display.flip()
def take_screenshot(self):
print("screenshot!")
pixels = self._render_fn_screenshot(self.play_state)
mtime = round(time.time() * 1000)
pixels = pixels.transpose((1, 0, 2))
# Black border
border_thickness = 5
pixels = pixels.at[:, :border_thickness].set(0.0)
pixels = pixels.at[:, -border_thickness:].set(0.0)
pixels = pixels.at[:border_thickness, :].set(0.0)
pixels = pixels.at[-border_thickness:, :].set(0.0)
im = Image.fromarray(np.array(pixels.astype(jnp.uint8)))
im.save(f"results/screenshot_{str(mtime)}.png")
def _get_selected_shape_global_indices(self):
def _idx(idx, type):
if type == ObjectType.CIRCLE:
return idx + self.static_env_params.num_polygons
return idx
indices_to_use = jnp.array([_idx(idx, type) for idx, type in self.all_selected_shapes])
return indices_to_use
# flag1
def _handle_events(self, do_dummy=False):
pygame_widgets.update(self.pygame_events)
if do_dummy or self.selected_shape_index < 0:
gravity_main = self.all_widgets[None]["sldGravity"].getValue()
gravity_max = self.all_widgets[None]["sldMaxGravity"].getValue()
gravity_main, gravity_max = min(gravity_main, gravity_max), max(gravity_main, gravity_max)
gravity_pcg_mask = self.all_widgets[None]["pcgTglGravity"].getValue()
def _set_single_global(state, gravity):
return state.replace(
gravity=state.gravity.at[1].set(gravity),
)
env_state = _set_single_global(self.pcg_state.env_state, gravity_main)
env_state_max = _set_single_global(self.pcg_state.env_state_max, gravity_max)
env_state_pcg_mask = _set_single_global(self.pcg_state.env_state_pcg_mask, gravity_pcg_mask)
if not do_dummy:
self.pcg_state = self.pcg_state.replace(
env_state=env_state,
env_state_max=env_state_max,
env_state_pcg_mask=env_state_pcg_mask,
)
if self.edit_shape_mode == EditMode.SELECT or do_dummy: # is on the hand.
if do_dummy or len(self.all_selected_shapes) > 1:
# this processes the tying together.
indices_to_use = self._get_selected_shape_global_indices()
if len(indices_to_use) > 1:
toggle_val = self.all_widgets["TIE_TOGETHER"]["tglTieTogether"].getValue()
idxs = itertools.product(indices_to_use, indices_to_use)
idxs_a = []
idxs_b = []
for i, j in idxs:
idxs_a.append(i)
idxs_b.append(j)
idxs = jnp.array(idxs_a), jnp.array(idxs_b)
if toggle_val:
self.pcg_state = self.pcg_state.replace(
tied_together=self.pcg_state.tied_together.at[idxs[0], idxs[1]].set(True)
)
else:
self.pcg_state = self.pcg_state.replace(
tied_together=self.pcg_state.tied_together.at[idxs[0], idxs[1]].set(False)
)
if self.selected_shape_index < 0 and not do_dummy:
return
if do_dummy or self.selected_shape_type in [ObjectType.POLYGON, ObjectType.CIRCLE]: # rigidbody
shape_main = select_object(
self.pcg_state.env_state, self.selected_shape_type, self.selected_shape_index
)
shape_max = select_object(
self.pcg_state.env_state_max, self.selected_shape_type, self.selected_shape_index
)
parent_container_main = (
self.pcg_state.env_state.circle
if self.selected_shape_type == ObjectType.CIRCLE
else self.pcg_state.env_state.polygon
)
parent_container_max = (
self.pcg_state.env_state_max.circle
if self.selected_shape_type == ObjectType.CIRCLE
else self.pcg_state.env_state_max.polygon
)
parent_container_pcg_mask = (
self.pcg_state.env_state_pcg_mask.circle
if self.selected_shape_type == ObjectType.CIRCLE
else self.pcg_state.env_state_pcg_mask.polygon
)
new_density_main = self.all_widgets[self.selected_shape_type]["sldDensity"].getValue()
new_density_max = self.all_widgets[self.selected_shape_type]["sldMaxDensity"].getValue()
density_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglDensity"].getValue()
if density_pcg_mask:
new_density_main, new_density_max = min(new_density_main, new_density_max), max(
new_density_main, new_density_max
)
fixated = self.all_widgets[self.selected_shape_type]["tglFixate"].getValue()
fix_val = 0.0 if fixated else 1.0
def _density_calcs(base, new_density):
inv_mass = jax.lax.select(
self.selected_shape_type == ObjectType.CIRCLE,
calc_inverse_mass_circle(base.radius, new_density),
calc_inverse_mass_polygon(base.vertices, base.n_vertices, self.static_env_params, new_density)[
0
],
)
inv_inertia = jax.lax.select(
self.selected_shape_type == ObjectType.CIRCLE,
calc_inverse_inertia_circle(base.radius, new_density),
calc_inverse_inertia_polygon(
base.vertices, base.n_vertices, self.static_env_params, new_density
),
)
return inv_mass, inv_inertia
inv_mass_main, inv_inertia_main = _density_calcs(shape_main, new_density_main)
inv_mass_max, inv_inertia_max = _density_calcs(shape_max, new_density_max)
friction_main = self.all_widgets[self.selected_shape_type]["sldFriction"].getValue()
friction_max = self.all_widgets[self.selected_shape_type]["sldMaxFriction"].getValue()
friction_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglFriction"].getValue()
if friction_pcg_mask:
friction_main, friction_max = min(friction_main, friction_max), max(friction_main, friction_max)
restitution = self.all_widgets[self.selected_shape_type]["sldRestitution"].getValue()
position_main = jnp.array(
[
self.all_widgets[self.selected_shape_type]["sldPosition_X"].getValue(),
self.all_widgets[self.selected_shape_type]["sldPosition_Y"].getValue(),
]
)
position_max = jnp.array(
[
self.all_widgets[self.selected_shape_type]["sldMaxPosition_X"].getValue(),
self.all_widgets[self.selected_shape_type]["sldMaxPosition_Y"].getValue(),
]
)
position_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglPosition_X"].getValue()
if position_pcg_mask:
position_main, position_max = jnp.minimum(position_main, position_max), jnp.maximum(
position_main, position_max
)
rotation_main = self.all_widgets[self.selected_shape_type]["sldRotation"].getValue()
rotation_max = self.all_widgets[self.selected_shape_type]["sldMaxRotation"].getValue()
rotation_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglRotation"].getValue()
if rotation_pcg_mask:
rotation_main, rotation_max = min(rotation_main, rotation_max), max(rotation_main, rotation_max)
velocity_main = jnp.array(
[
self.all_widgets[self.selected_shape_type]["sldVelocity_X"].getValue(),
self.all_widgets[self.selected_shape_type]["sldVelocity_Y"].getValue(),
]
)
velocity_max = jnp.array(
[
self.all_widgets[self.selected_shape_type]["sldMaxVelocity_X"].getValue(),
self.all_widgets[self.selected_shape_type]["sldMaxVelocity_Y"].getValue(),
]
)
velocity_main, velocity_max = jnp.minimum(velocity_main, velocity_max), jnp.maximum(
velocity_main, velocity_max
)
velocity_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglVelocity_X"].getValue()
angular_velocity_main = self.all_widgets[self.selected_shape_type]["sldAngular_Velocity"].getValue()
angular_velocity_max = self.all_widgets[self.selected_shape_type]["sldMaxAngular_Velocity"].getValue()
angular_velocity_pcg_mask = self.all_widgets[self.selected_shape_type][
"pcgTglAngular_Velocity"
].getValue()
if angular_velocity_pcg_mask:
angular_velocity_main, angular_velocity_max = min(angular_velocity_main, angular_velocity_max), max(
angular_velocity_main, angular_velocity_max
)
# Circle stuff
radius_main, radius_max, radius_pcg_mask = None, None, None
if self.selected_shape_type == ObjectType.CIRCLE:
radius_main = self.all_widgets[self.selected_shape_type]["sldRadius"].getValue()
radius_max = self.all_widgets[self.selected_shape_type]["sldMaxRadius"].getValue()
radius_pcg_mask = self.all_widgets[self.selected_shape_type]["pcgTglRadius"].getValue()
# Poly stuff
vertices_main, vertices_max, vertices_pcg_mask = None, None, None
if self.selected_shape_type == ObjectType.POLYGON:
# Triangle
new_size_main = self.all_widgets[self.selected_shape_type]["sldSize"].getValue()
new_size_max = self.all_widgets[self.selected_shape_type]["sldMaxSize"].getValue()
current_size = jnp.abs(self.pcg_state.env_state.polygon.vertices[self.selected_shape_index]).max()
scale_main = new_size_main / current_size
scale_max = new_size_max / current_size
vertices_main = self.pcg_state.env_state.polygon.vertices[self.selected_shape_index] * scale_main
vertices_max = self.pcg_state.env_state.polygon.vertices[self.selected_shape_index] * scale_max
vertices_pcg_mask = (
jnp.ones_like(vertices_main, dtype=bool)
* self.all_widgets[self.selected_shape_type]["pcgTglSize"].getValue()
)
def _set_single_state_rbody(
state,
parent_container,
density,
inv_mass,
inv_inertia,
friction,
position,
radius,
rotation,
velocity,
angular_velocity,
vertices,
):
new = {
"friction": parent_container.friction.at[self.selected_shape_index].set(friction),
"collision_mode": parent_container.collision_mode.at[self.selected_shape_index].set(
int(self.all_widgets[self.selected_shape_type]["sldCollidability"].getValue())
),
"inverse_mass": parent_container.inverse_mass.at[self.selected_shape_index].set(
inv_mass * fix_val
),
"inverse_inertia": parent_container.inverse_inertia.at[self.selected_shape_index].set(
inv_inertia * fix_val
),
"position": parent_container.position.at[self.selected_shape_index].set(position),
"rotation": parent_container.rotation.at[self.selected_shape_index].set(rotation),
"velocity": parent_container.velocity.at[self.selected_shape_index].set(velocity),
"angular_velocity": parent_container.angular_velocity.at[self.selected_shape_index].set(
angular_velocity
),
"restitution": parent_container.restitution.at[self.selected_shape_index].set(restitution),
}
if self.selected_shape_type == ObjectType.CIRCLE:
state = state.replace(
circle=state.circle.replace(
**new,
radius=parent_container.radius.at[self.selected_shape_index].set(radius),
),
circle_shape_roles=state.circle_shape_roles.at[self.selected_shape_index].set(
int(self.all_widgets[self.selected_shape_type]["sldRole"].getValue())
),
circle_densities=state.circle_densities.at[self.selected_shape_index].set(density),
)
else:
state = state.replace(
polygon=state.polygon.replace(
**new,
vertices=parent_container.vertices.at[self.selected_shape_index].set(vertices),
),
polygon_shape_roles=state.polygon_shape_roles.at[self.selected_shape_index].set(
int(self.all_widgets[self.selected_shape_type]["sldRole"].getValue())
),
polygon_densities=state.polygon_densities.at[self.selected_shape_index].set(density),
)
return state
position_delta = position_main - shape_main.position
env_state = _set_single_state_rbody(
self.pcg_state.env_state,
parent_container_main,
new_density_main,
inv_mass_main,
inv_inertia_main,
friction_main,
position_main,
radius_main,
rotation_main,
velocity_main,
angular_velocity_main,
vertices_main,
)
env_state_max = _set_single_state_rbody(
self.pcg_state.env_state_max,
parent_container_max,
new_density_max,
inv_mass_max,
inv_inertia_max,
friction_max,
position_max,
radius_max,
rotation_max,
velocity_max,
angular_velocity_max,
vertices_max,
)
env_state_pcg_mask = _set_single_state_rbody(
self.pcg_state.env_state_pcg_mask,
parent_container_pcg_mask,
density_pcg_mask,
density_pcg_mask,
density_pcg_mask,
friction_pcg_mask,
position_pcg_mask,
radius_pcg_mask,
rotation_pcg_mask,
velocity_pcg_mask,
angular_velocity_pcg_mask,
vertices_pcg_mask,
)
# Now, we must also set all of the tied shape's positions
correct_index = self.selected_shape_index + (
self.static_env_params.num_polygons if self.selected_shape_type == ObjectType.CIRCLE else 0
)
nonzero_indices = set(
jnp.nonzero(self.pcg_state.tied_together[correct_index].at[correct_index].set(False))[0].tolist()
)
for idx in nonzero_indices:
if idx < self.static_env_params.num_polygons:
env_state = env_state.replace(
polygon=env_state.polygon.replace(
position=env_state.polygon.position.at[idx].set(
env_state.polygon.position[idx] + position_delta
)
)
)
else:
idx = idx - self.static_env_params.num_polygons
env_state = env_state.replace(
circle=env_state.circle.replace(
position=env_state.circle.position.at[idx].set(
env_state.circle.position[idx] + position_delta
)
)
)
if not do_dummy:
self.pcg_state = PCGState(
env_state=env_state,
env_state_max=env_state_max,
env_state_pcg_mask=env_state_pcg_mask,
tied_together=self.pcg_state.tied_together,
)
if do_dummy or self.selected_shape_type == ObjectType.JOINT: # rjoint
speed_main = self.all_widgets[ObjectType.JOINT]["sldSpeed"].getValue()
speed_max = self.all_widgets[ObjectType.JOINT]["sldMaxSpeed"].getValue()
speed_pcg_mask = self.all_widgets[ObjectType.JOINT]["pcgTglSpeed"].getValue()
if speed_pcg_mask:
speed_main, speed_max = min(speed_main, speed_max), max(speed_main, speed_max)
power_main = self.all_widgets[ObjectType.JOINT]["sldPower"].getValue()
power_max = self.all_widgets[ObjectType.JOINT]["sldMaxPower"].getValue()
power_pcg_mask = self.all_widgets[ObjectType.JOINT]["pcgTglPower"].getValue()
if power_pcg_mask:
power_main, power_max = min(power_main, power_max), max(power_main, power_max)
motor_binding_val_min = int(self.all_widgets[ObjectType.JOINT]["sldColour"].getValue())
motor_binding_val_max = int(self.all_widgets[ObjectType.JOINT]["sldMaxColour"].getValue())
colour_pcg_mask = self.all_widgets[ObjectType.JOINT]["pcgTglColour"].getValue()
if colour_pcg_mask:
motor_binding_val_min, motor_binding_val_max = min(
motor_binding_val_min, motor_binding_val_max
), max(motor_binding_val_min, motor_binding_val_max)
auto_motor_val = self.all_widgets[ObjectType.JOINT]["tglAutoMotor"].getValue()
joint_limits_val = self.all_widgets[ObjectType.JOINT]["tglJointLimits"].getValue()
is_fixed_val = self.all_widgets[ObjectType.JOINT]["tglIsFixedJoint"].getValue()
is_motor_on = self.all_widgets[ObjectType.JOINT]["tglIsMotorOn"].getValue()
min_rot_val = jnp.deg2rad(self.all_widgets[ObjectType.JOINT]["sldMin_Rotation"].getValue())
max_rot_val = jnp.deg2rad(self.all_widgets[ObjectType.JOINT]["sldMax_Rotation"].getValue())
# ensure the min is less than the max
min_rot_val, max_rot_val = min(min_rot_val, max_rot_val), max(min_rot_val, max_rot_val)
def _set_single_state_joint(state, speed, power, colour):
state = state.replace(
joint=state.joint.replace(
motor_speed=state.joint.motor_speed.at[self.selected_shape_index].set(speed),
motor_power=state.joint.motor_power.at[self.selected_shape_index].set(power),
motor_has_joint_limits=state.joint.motor_has_joint_limits.at[self.selected_shape_index].set(
joint_limits_val
),
min_rotation=state.joint.min_rotation.at[self.selected_shape_index].set(min_rot_val),
max_rotation=state.joint.max_rotation.at[self.selected_shape_index].set(max_rot_val),
is_fixed_joint=state.joint.is_fixed_joint.at[self.selected_shape_index].set(is_fixed_val),
motor_on=state.joint.motor_on.at[self.selected_shape_index].set(is_motor_on),
),
motor_bindings=state.motor_bindings.at[self.selected_shape_index].set(colour),
motor_auto=state.motor_auto.at[self.selected_shape_index].set(auto_motor_val),
)
return state
env_state = _set_single_state_joint(
self.pcg_state.env_state,
speed_main,
power_main,
motor_binding_val_min,
)
env_state_max = _set_single_state_joint(
self.pcg_state.env_state_max,
speed_max,
power_max,
motor_binding_val_max,
)
env_state_pcg_mask = _set_single_state_joint(
self.pcg_state.env_state_pcg_mask, speed_pcg_mask, power_pcg_mask, colour_pcg_mask
)
if not do_dummy:
self.pcg_state = self.pcg_state.replace(
env_state=env_state,
env_state_max=env_state_max,
env_state_pcg_mask=env_state_pcg_mask,
)
if do_dummy or self.selected_shape_type == ObjectType.THRUSTER: # thruster
power_main = self.all_widgets[ObjectType.THRUSTER]["sldPower"].getValue()
power_max = self.all_widgets[ObjectType.THRUSTER]["sldMaxPower"].getValue()
power_pcg_mask = self.all_widgets[ObjectType.THRUSTER]["pcgTglPower"].getValue()
if power_pcg_mask:
power_main, power_max = min(power_main, power_max), max(power_main, power_max)
def _set_single_state_thruster(state, power):
return state.replace(
thruster=state.thruster.replace(
power=state.thruster.power.at[self.selected_shape_index].set(power),
),
thruster_bindings=state.thruster_bindings.at[self.selected_shape_index].set(
int(self.all_widgets[ObjectType.THRUSTER]["sldColour"].getValue())
),
)
env_state = _set_single_state_thruster(self.pcg_state.env_state, power_main)
env_state_max = _set_single_state_thruster(self.pcg_state.env_state_max, power_max)
env_state_pcg_mask = _set_single_state_thruster(self.pcg_state.env_state_pcg_mask, power_pcg_mask)
if not do_dummy:
self.pcg_state = self.pcg_state.replace(
env_state=env_state,
env_state_max=env_state_max,
env_state_pcg_mask=env_state_pcg_mask,
)
# flag2
def _put_state_values_into_gui(self, pcg_state=None):
def _set_toggle_val(toggle, val):
if toggle.getValue() != val:
toggle.toggle()
def _enable_slider(slider):
slider.enable()
slider.colour = (200, 200, 200)
slider.handleColour = (0, 0, 0)
def _disable_slider(slider):
slider.disable()
slider.colour = (255, 0, 0)
slider.handleColour = (255, 0, 0)
if pcg_state is None:
# state = self.edit_state
raise ValueError
def pcg_text(main_val, max_val, pcg_mask):
if pcg_mask:
return f"{main_val:.2f} - {max_val:.2f}"
else:
return f"{main_val:.2f}"
# global ones
gravity_pcg_mask = pcg_state.env_state_pcg_mask.gravity[1]
self.all_widgets[None]["lblGravity"].setText(
f"Gravity: {pcg_text(pcg_state.env_state.gravity[1], pcg_state.env_state_max.gravity[1], pcg_state.env_state_pcg_mask.gravity[1])}"
)
self.all_widgets[None]["sldGravity"].setValue(pcg_state.env_state.gravity[1])
self.all_widgets[None]["sldMaxGravity"].setValue(pcg_state.env_state_max.gravity[1])
if not gravity_pcg_mask:
_disable_slider(self.all_widgets[None]["sldMaxGravity"])
else:
_enable_slider(self.all_widgets[None]["sldMaxGravity"])
if self.edit_shape_mode != EditMode.SELECT or self.selected_shape_index < 0:
return
obj_main = select_object(pcg_state.env_state, self.selected_shape_type, self.selected_shape_index)
obj_max = select_object(pcg_state.env_state_max, self.selected_shape_type, self.selected_shape_index)
obj_pcg_mask = select_object(pcg_state.env_state_pcg_mask, self.selected_shape_type, self.selected_shape_index)
if len(self.all_selected_shapes) > 1:
indices_to_use = self._get_selected_shape_global_indices()
are_all_tied = pcg_state.tied_together[indices_to_use.min(), indices_to_use].all()
_set_toggle_val(self.all_widgets["TIE_TOGETHER"]["tglTieTogether"], are_all_tied)
if self.selected_shape_type == ObjectType.JOINT:
self.all_widgets[ObjectType.JOINT]["lblSpeed"].setText(
f"Speed: {pcg_text(obj_main.motor_speed, obj_max.motor_speed, obj_pcg_mask.motor_speed)}"
)
self.all_widgets[ObjectType.JOINT]["sldSpeed"].setValue(obj_main.motor_speed)
self.all_widgets[ObjectType.JOINT]["sldMaxSpeed"].setValue(obj_max.motor_speed)
_set_toggle_val(self.all_widgets[ObjectType.JOINT]["pcgTglSpeed"], obj_pcg_mask.motor_speed)
if obj_pcg_mask.motor_speed:
_enable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxSpeed"])
else:
_disable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxSpeed"])
self.all_widgets[ObjectType.JOINT]["lblPower"].setText(
f"Power: {pcg_text(obj_main.motor_power, obj_max.motor_power, obj_pcg_mask.motor_power)}"
)
self.all_widgets[ObjectType.JOINT]["sldPower"].setValue(obj_main.motor_power)
self.all_widgets[ObjectType.JOINT]["sldMaxPower"].setValue(obj_max.motor_power)
_set_toggle_val(self.all_widgets[ObjectType.JOINT]["pcgTglPower"], obj_pcg_mask.motor_power)
if obj_pcg_mask.motor_power:
_enable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxPower"])
else:
_disable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxPower"])
self.all_widgets[ObjectType.JOINT]["lblColour"].setText(
f"Colour: {pcg_state.env_state.motor_bindings[self.selected_shape_index]}"
)
self.all_widgets[ObjectType.JOINT]["sldColour"].setValue(
pcg_state.env_state.motor_bindings[self.selected_shape_index]
)
self.all_widgets[ObjectType.JOINT]["sldMaxColour"].setValue(
pcg_state.env_state_max.motor_bindings[self.selected_shape_index]
)
colour_pcg_mask = pcg_state.env_state_pcg_mask.motor_bindings[self.selected_shape_index]
if not colour_pcg_mask:
_disable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxColour"])
else:
_enable_slider(self.all_widgets[ObjectType.JOINT]["sldMaxColour"])
self.all_widgets[ObjectType.JOINT]["lblJointLimits"].setText(
f"Joint Limits: {obj_main.motor_has_joint_limits}"
)
widget_motor_has_joint_limits = self.all_widgets[ObjectType.JOINT]["tglJointLimits"].getValue()
if obj_main.motor_has_joint_limits != widget_motor_has_joint_limits: # Update the toggle
self.all_widgets[ObjectType.JOINT]["tglJointLimits"].toggle()
mini, maxi = jnp.rad2deg(obj_main.min_rotation), jnp.rad2deg(obj_main.max_rotation)
self.all_widgets[ObjectType.JOINT]["lblMin_Rotation"].setText(f"Min Rot: {int(mini)}")
self.all_widgets[ObjectType.JOINT]["sldMin_Rotation"].setValue(mini)
self.all_widgets[ObjectType.JOINT]["lblMax_Rotation"].setText(f"Max Rot: {int(maxi)}")
self.all_widgets[ObjectType.JOINT]["sldMax_Rotation"].setValue(maxi)
if not obj_main.motor_has_joint_limits:
for k in ["min_rotation", "max_rotation"]:
self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].disable()
self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].colour = (255, 0, 0)
self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].handleColour = (255, 0, 0)
else:
for k in ["min_rotation", "max_rotation"]:
self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].enable()
self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].colour = (200, 200, 200)
self.all_widgets[self.selected_shape_type][f"sld{k.title()}"].handleColour = (0, 0, 0)
self.all_widgets[ObjectType.JOINT]["lblAutoMotor"].setText(
f"Auto: {pcg_state.env_state.motor_auto[self.selected_shape_index]}"
)
widget_is_auto_motor = self.all_widgets[ObjectType.JOINT]["tglAutoMotor"].getValue()
if pcg_state.env_state.motor_auto[self.selected_shape_index] != widget_is_auto_motor: # Update the toggle
self.all_widgets[ObjectType.JOINT]["tglAutoMotor"].toggle()
self.all_widgets[ObjectType.JOINT]["lblIsFixedJoint"].setText(f"Fixed: {obj_main.is_fixed_joint}")
widget_is_motor_on = self.all_widgets[ObjectType.JOINT]["tglIsFixedJoint"].getValue()
if obj_main.is_fixed_joint != widget_is_motor_on: # Update the toggle
self.all_widgets[ObjectType.JOINT]["tglIsFixedJoint"].toggle()
self.all_widgets[ObjectType.JOINT]["lblIsMotorOn"].setText(f"Motor On: {obj_main.motor_on}")
widget_is_motor_on = self.all_widgets[ObjectType.JOINT]["tglIsMotorOn"].getValue()
if obj_main.motor_on != widget_is_motor_on: # Update the toggle
self.all_widgets[ObjectType.JOINT]["tglIsMotorOn"].toggle()
elif self.selected_shape_type == ObjectType.THRUSTER:
# thruster
self.all_widgets[ObjectType.THRUSTER]["lblPower"].setText(
f"Power: {pcg_text(obj_main.power, obj_max.power, obj_pcg_mask.power)}"
)
self.all_widgets[ObjectType.THRUSTER]["sldPower"].setValue(obj_main.power)
self.all_widgets[ObjectType.THRUSTER]["sldMaxPower"].setValue(obj_max.power)
_set_toggle_val(self.all_widgets[ObjectType.THRUSTER]["pcgTglPower"], obj_pcg_mask.power)
if obj_pcg_mask.power:
_enable_slider(self.all_widgets[ObjectType.THRUSTER]["sldMaxPower"])
else:
_disable_slider(self.all_widgets[ObjectType.THRUSTER]["sldMaxPower"])
self.all_widgets[ObjectType.THRUSTER]["sldColour"].setValue(
pcg_state.env_state.thruster_bindings[self.selected_shape_index]
)
self.all_widgets[ObjectType.THRUSTER]["lblColour"].setText(
f"Colour: {pcg_state.env_state.thruster_bindings[self.selected_shape_index]}"
)
elif self.selected_shape_type in [ObjectType.POLYGON, ObjectType.CIRCLE]:
# rigidbody
# Position
# We use the mask for position_x for the entire position vector
self.all_widgets[self.selected_shape_type]["lblPosition_X"].setText(
f"Position X: {pcg_text(obj_main.position[0], obj_max.position[0], obj_pcg_mask.position[0])}"
)
self.all_widgets[self.selected_shape_type]["sldPosition_X"].setValue(obj_main.position[0])
self.all_widgets[self.selected_shape_type]["sldMaxPosition_X"].setValue(obj_max.position[0])
_set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglPosition_X"], obj_pcg_mask.position[0])
self.all_widgets[self.selected_shape_type]["lblPosition_Y"].setText(
f"Position Y: {pcg_text(obj_main.position[1], obj_max.position[1], obj_pcg_mask.position[0])}"
)
self.all_widgets[self.selected_shape_type]["sldPosition_Y"].setValue(obj_main.position[1])
self.all_widgets[self.selected_shape_type]["sldMaxPosition_Y"].setValue(obj_max.position[1])
if obj_pcg_mask.position[0]:
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxPosition_X"])
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxPosition_Y"])
else:
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxPosition_X"])
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxPosition_Y"])
# Velocity
# We use the mask for velocity_x for the entire velocity vector
self.all_widgets[self.selected_shape_type]["lblVelocity_X"].setText(
f"Velocity X: {pcg_text(obj_main.velocity[0], obj_max.velocity[0], obj_pcg_mask.velocity[0])}"
)
self.all_widgets[self.selected_shape_type]["sldVelocity_X"].setValue(obj_main.velocity[0])
self.all_widgets[self.selected_shape_type]["sldMaxVelocity_X"].setValue(obj_max.velocity[0])
_set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglVelocity_X"], obj_pcg_mask.velocity[0])
self.all_widgets[self.selected_shape_type]["lblVelocity_Y"].setText(
f"Velocity Y: {pcg_text(obj_main.velocity[1], obj_max.velocity[1], obj_pcg_mask.velocity[0])}"
)
self.all_widgets[self.selected_shape_type]["sldVelocity_Y"].setValue(obj_main.velocity[1])
self.all_widgets[self.selected_shape_type]["sldMaxVelocity_Y"].setValue(obj_max.velocity[1])
if obj_pcg_mask.velocity[0]:
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxVelocity_X"])
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxVelocity_Y"])
else:
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxVelocity_X"])
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxVelocity_Y"])
# Density
is_fixated = obj_main.inverse_mass == 0
def _calc_density(state):
if self.selected_shape_type == ObjectType.POLYGON:
return state.polygon_densities[self.selected_shape_index]
elif self.selected_shape_type == ObjectType.CIRCLE:
return state.circle_densities[self.selected_shape_index]
else:
raise ValueError
density_main = _calc_density(pcg_state.env_state)
density_max = _calc_density(pcg_state.env_state_max)
density_pcg_mask = obj_pcg_mask.inverse_mass
self.all_widgets[self.selected_shape_type]["lblDensity"].setText(
f"Density: {pcg_text(density_main, density_max, density_pcg_mask)}"
)
self.all_widgets[self.selected_shape_type]["sldDensity"].setValue(density_main)
self.all_widgets[self.selected_shape_type]["sldMaxDensity"].setValue(density_max)
_set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglDensity"], density_pcg_mask)
if is_fixated:
_disable_slider(self.all_widgets[self.selected_shape_type]["sldDensity"])
else:
_enable_slider(self.all_widgets[self.selected_shape_type]["sldDensity"])
if is_fixated or (not density_pcg_mask):
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxDensity"])
else:
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxDensity"])
# Friction
self.all_widgets[self.selected_shape_type]["lblFriction"].setText(
f"Friction: {pcg_text(obj_main.friction, obj_max.friction, obj_pcg_mask.friction)}"
)
self.all_widgets[self.selected_shape_type]["sldFriction"].setValue(obj_main.friction)
self.all_widgets[self.selected_shape_type]["sldMaxFriction"].setValue(obj_max.friction)
_set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglFriction"], obj_pcg_mask.friction)
if not obj_pcg_mask.friction:
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxFriction"])
else:
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxFriction"])
# Restitution
self.all_widgets[self.selected_shape_type]["lblRestitution"].setText(
f"Restitution: {obj_main.restitution:.2f}"
)
self.all_widgets[self.selected_shape_type]["sldRestitution"].setValue(obj_main.restitution)
# Rotation
self.all_widgets[self.selected_shape_type]["lblRotation"].setText(
f"Rotation: {pcg_text(obj_main.rotation, obj_max.rotation, obj_pcg_mask.rotation)}"
)
self.all_widgets[self.selected_shape_type]["sldRotation"].setValue(obj_main.rotation)
self.all_widgets[self.selected_shape_type]["sldMaxRotation"].setValue(obj_max.rotation)
_set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglRotation"], obj_pcg_mask.rotation)
if not obj_pcg_mask.rotation:
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxRotation"])
else:
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxRotation"])
# Angular_Velocity
self.all_widgets[self.selected_shape_type]["lblAngular_Velocity"].setText(
f"Angular_Velocity: {pcg_text(obj_main.angular_velocity, obj_max.angular_velocity, obj_pcg_mask.angular_velocity)}"
)
self.all_widgets[self.selected_shape_type]["sldAngular_Velocity"].setValue(obj_main.angular_velocity)
self.all_widgets[self.selected_shape_type]["sldMaxAngular_Velocity"].setValue(obj_max.angular_velocity)
_set_toggle_val(
self.all_widgets[self.selected_shape_type]["pcgTglAngular_Velocity"], obj_pcg_mask.angular_velocity
)
if not obj_pcg_mask.angular_velocity:
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxAngular_Velocity"])
else:
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxAngular_Velocity"])
# Collision mode
self.all_widgets[self.selected_shape_type]["lblCollidability"].setText(
f"Collidability: {obj_main.collision_mode}"
)
self.all_widgets[self.selected_shape_type]["sldCollidability"].setValue(obj_main.collision_mode)
# Shape role
if self.selected_shape_type == ObjectType.POLYGON:
shape_role = pcg_state.env_state.polygon_shape_roles[self.selected_shape_index]
else:
shape_role = pcg_state.env_state.circle_shape_roles[self.selected_shape_index]
self.all_widgets[self.selected_shape_type]["sldRole"].setValue(shape_role)
self.all_widgets[self.selected_shape_type]["lblRole"].setText(f"Role: {shape_role}")
# Fixate
self.all_widgets[self.selected_shape_type]["lblFixate"].setText(f"Fixate: {is_fixated}")
widget_is_fixed = self.all_widgets[self.selected_shape_type]["tglFixate"].getValue()
if is_fixated != widget_is_fixed: # Update the toggle
self.all_widgets[self.selected_shape_type]["tglFixate"].toggle()
# Radius
if self.selected_shape_type == ObjectType.CIRCLE:
self.all_widgets[self.selected_shape_type]["lblRadius"].setText(
f"Radius: {pcg_text(obj_main.radius, obj_max.radius, obj_pcg_mask.radius)}"
)
self.all_widgets[self.selected_shape_type]["sldRadius"].setValue(obj_main.radius)
self.all_widgets[self.selected_shape_type]["sldMaxRadius"].setValue(obj_max.radius)
_set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglRadius"], obj_pcg_mask.radius)
if not obj_pcg_mask.radius:
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxRadius"])
else:
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxRadius"])
elif self.selected_shape_type == ObjectType.POLYGON:
size_main = jnp.abs(obj_main.vertices).max()
size_max = jnp.abs(obj_max.vertices).max()
size_pcg_mask = obj_pcg_mask.vertices.any()
self.all_widgets[self.selected_shape_type]["lblSize"].setText(
f"Size: {pcg_text(size_main, size_max, size_pcg_mask)}"
)
self.all_widgets[self.selected_shape_type]["sldSize"].setValue(size_main)
self.all_widgets[self.selected_shape_type]["sldMaxSize"].setValue(size_max)
_set_toggle_val(self.all_widgets[self.selected_shape_type]["pcgTglSize"], size_pcg_mask)
if not size_pcg_mask:
_disable_slider(self.all_widgets[self.selected_shape_type]["sldMaxSize"])
else:
_enable_slider(self.all_widgets[self.selected_shape_type]["sldMaxSize"])
def _render_side_panel(self):
arr = jnp.ones((self.side_panel_width, self.static_env_params.screen_dim[1], 3)) * (
jnp.array([135.0, 206.0, 235.0])[None, None] + 20
)
return arr
def make_label_and_slider(
self,
start_y,
label_text,
slider_min=0.0,
slider_max=1.0,
slider_step=0.05,
font_size=18,
is_toggle=False,
is_pcg=False,
add_pcg_toggle=True,
):
Wl = round(self.W * 0.7)
pcg_lbl = None
pcg_toggle = None
widget_max = None
label = TextBox(
self.screen_surface,
self.MARGIN,
start_y,
Wl,
20,
fontSize=font_size,
margin=0,
placeholderText=label_text,
font=pygame.font.SysFont("sans-serif", font_size),
)
label.disable() # Act as label instead of textbox
if is_toggle:
widget = Toggle(self.screen_surface, self.W // 2 - 20, start_y + 23, 20, 13)
else:
widget = Slider(
self.screen_surface,
self.MARGIN,
start_y + 23,
Wl,
13,
min=slider_min,
max=slider_max,
step=slider_step,
)
if is_pcg:
widget_max = Slider(
self.screen_surface,
self.MARGIN,
start_y + 50,
Wl,
13,
min=slider_min,
max=slider_max,
step=slider_step,
)
pcg_lbl = TextBox(
self.screen_surface,
self.MARGIN + Wl + 30,
start_y,
60,
20,
fontSize=font_size,
margin=0,
placeholderText="PCG",
font=pygame.font.SysFont("sans-serif", font_size),
)
if add_pcg_toggle:
pcg_toggle = Toggle(self.screen_surface, Wl + 80, start_y + 23, 20, 13)
return label, widget, (pcg_toggle, pcg_lbl, widget_max)
def _setup_side_panel(self):
W = self.W = int(self.side_panel_width * self.upscale * 0.8)
MARGIN = self.MARGIN = (self.side_panel_width * self.upscale - W) // 2
# global values
G = {}
gravity_label, gravity_slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
150, "Gravity", -20.0, 0.0, 0.2, is_pcg=True
)
G["lblGravity"] = gravity_label
G["sldGravity"] = gravity_slider
G["sldMaxGravity"] = slider_max
G["pcgLblGravity"] = pcg_lbl
G["pcgTglGravity"] = pcg_toggle
# thruster values
T = {}
thruster_power_label, thruster_power_slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
150, "Power", slider_max=3.0, is_pcg=True
)
T["lblPower"] = thruster_power_label
T["sldPower"] = thruster_power_slider
T["sldMaxPower"] = slider_max
T["pcgLblPower"] = pcg_lbl
T["pcgTglPower"] = pcg_toggle
thruster_colour_label, thruster_colour_slider, _ = self.make_label_and_slider(
250, "Colour", 0, self.static_env_params.num_thruster_bindings - 1, 1
)
T["lblColour"] = thruster_colour_label
T["sldColour"] = thruster_colour_slider
# joints
D = {}
for i, (name, (mini, maxi, step), is_pcg) in enumerate(
zip(
["speed", "power", "colour", "min_rotation", "max_rotation"],
[
(-3, 3, 0.05),
(0, 3, 0.05),
(0, self.static_env_params.num_motor_bindings - 1, 1),
(-180, 180, 5),
(-180, 180, 5),
],
[True, True, True, False, False],
)
):
label, slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
150 + 80 * i,
f"Motor {name.title()}",
slider_min=mini,
slider_max=maxi,
slider_step=step,
is_pcg=is_pcg,
)
D["lbl" + name.title()] = label
D["sld" + name.title()] = slider
if is_pcg:
D["sldMax" + name.title()] = slider_max
D["pcgLbl" + name.title()] = pcg_lbl
D["pcgTgl" + name.title()] = pcg_toggle
label, toggle, _ = self.make_label_and_slider(150 + 80 * 5, "Joint Limits", is_toggle=True)
D["lblJointLimits"] = label
D["tglJointLimits"] = toggle
label, toggle, _ = self.make_label_and_slider(150 + 80 * 6, "Auto", is_toggle=True)
D["lblAutoMotor"] = label
D["tglAutoMotor"] = toggle
label, toggle, _ = self.make_label_and_slider(150 + 80 * 7, "Fixed", is_toggle=True)
D["lblIsFixedJoint"] = label
D["tglIsFixedJoint"] = toggle
label, toggle, _ = self.make_label_and_slider(150 + 80 * 8, "Motor On", is_toggle=True)
D["lblIsMotorOn"] = label
D["tglIsMotorOn"] = toggle
def _create_rigid_body_base_gui():
D_rigid = {}
# rigidbodies
total_toggles = 0
total_non_toggles = 0
for i, (name, bounds, is_pcg, add_pcg_toggle) in enumerate(
zip(
[
"position_x",
"position_y",
"rotation",
"velocity_x",
"velocity_y",
"angular_velocity",
"density",
"friction",
"restitution",
"collidability",
"role",
],
[
(0, 5.0, 0.01),
(0, 5.0, 0.01),
(-2 * jnp.pi, 2 * jnp.pi, 0.01),
(-10.0, 10.0, 0.1),
(-10.0, 10.0, 0.1),
(-6, 6.0, 0.01),
(0.1, 5.0, 0.1),
(0.02, 1.0, 0.02),
(0.0, 0.8, 0.02),
(0, 2, 1),
(0, 3, 1),
],
[True, True, True, True, True, True, True, True, False, False, False],
[True, False, True, True, False, True, True, True, False, False, False],
)
):
location = 50 + 80 * total_toggles + 40 * total_non_toggles
label, slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
location, name.title(), *bounds, is_pcg=is_pcg, add_pcg_toggle=add_pcg_toggle
)
total_toggles += is_pcg
total_non_toggles += not is_pcg
D_rigid["lbl" + name.title()] = label
D_rigid["sld" + name.title()] = slider
if is_pcg:
D_rigid["sldMax" + name.title()] = slider_max
if add_pcg_toggle:
D_rigid["pcgLbl" + name.title()] = pcg_lbl
D_rigid["pcgTgl" + name.title()] = pcg_toggle
location = 50 + 80 * total_toggles + 40 * total_non_toggles
# toggles:
label, toggle, _ = self.make_label_and_slider(location, "Fixate", is_toggle=True)
D_rigid["lblFixate"] = label
D_rigid["tglFixate"] = toggle
return D_rigid, location
# Circle extras
D_circle, location = _create_rigid_body_base_gui()
label, slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
location + 40, "Radius", slider_min=0.1, slider_max=1.0, slider_step=0.02, is_pcg=True, add_pcg_toggle=True
)
D_circle["lblRadius"] = label
D_circle["sldRadius"] = slider
D_circle["sldMaxRadius"] = slider_max
D_circle["pcgLblRadius"] = pcg_lbl
D_circle["pcgTglRadius"] = pcg_toggle
# Polygon extras
D_poly, location = _create_rigid_body_base_gui()
label, slider, (pcg_toggle, pcg_lbl, slider_max) = self.make_label_and_slider(
location + 40, "Size", slider_min=0.1, slider_max=2.0, slider_step=0.02, is_pcg=True, add_pcg_toggle=True
)
D_poly["lblSize"] = label
D_poly["sldSize"] = slider
D_poly["sldMaxSize"] = slider_max
D_poly["pcgLblSize"] = pcg_lbl
D_poly["pcgTglSize"] = pcg_toggle
label, toggle, _ = self.make_label_and_slider(150 + 80 * 5, "Tie Positions Together", is_toggle=True)
TIE_TOGETHER = {
"lblTieTogether": label,
"tglTieTogether": toggle,
}
self.all_widgets = {
ObjectType.THRUSTER: T,
ObjectType.JOINT: D,
"GENERAL": {
"lblGeneral": TextBox(
self.screen_surface,
MARGIN,
10,
W,
30,
fontSize=20,
margin=0,
placeholderText="General",
font=pygame.font.SysFont("sans-serif", 35),
),
},
ObjectType.POLYGON: D_poly,
ObjectType.CIRCLE: D_circle,
None: G,
"TIE_TOGETHER": TIE_TOGETHER,
}
self._hide_all_widgets()
def _render_edit_overlay(self, pixels, is_editing, edit_shape_mode):
is_editing_texture = jax.lax.select(is_editing, EDIT_TEXTURE_RGBA, PLAY_TEXTURE_RGBA)
is_editing_texture = jnp.repeat(jnp.repeat(is_editing_texture, self.upscale, axis=0), self.upscale, axis=1)
offset = self.side_panel_width * self.upscale
w = 64 * self.upscale
offset2 = int(w * 1.25)
offset_y = 16 * self.upscale
play_tex_with_background = (1 - is_editing_texture[:, :, 3:]) * pixels[
offset + 0 : offset + w, 0:w
] + is_editing_texture[:, :, 3:] * is_editing_texture[:, :, :3]
pixels = pixels.at[offset : offset + w, 0:w].set(play_tex_with_background)
edit_shape_texture = jax.lax.switch(
edit_shape_mode,
[
lambda: CIRCLE_TEXTURE_RGBA,
lambda: RECT_TEXTURE_RGBA,
lambda: RJOINT_TEXTURE_RGBA,
lambda: SELECT_TEXTURE_RGBA,
lambda: TRIANGLE_TEXTURE_RGBA,
lambda: THRUSTER_TEXTURE_RGBA,
],
)
edit_shape_texture = jnp.repeat(jnp.repeat(edit_shape_texture, self.upscale, axis=0), self.upscale, axis=1)
edit_shape_texture_alpha = edit_shape_texture[:, :, 3:] * is_editing
w = 32 * self.upscale
edit_shape_texture_with_background = (1 - edit_shape_texture_alpha) * pixels[
offset + offset2 : offset + offset2 + w, offset_y : offset_y + w
] + edit_shape_texture_alpha * edit_shape_texture[:, :, :3]
pixels = pixels.at[offset + offset2 : offset + offset2 + w, offset_y : offset_y + w].set(
edit_shape_texture_with_background
)
return pixels
def edit(self):
self.rng, _rng = jax.random.split(self.rng)
pcg_state = self.pcg_state
left_click = False
right_click = False
keys = []
keys_up_this_frame = set()
for event in self.pygame_events:
if event.type == pygame.KEYDOWN:
if (
event.key == pygame.K_s
and (pygame.key.get_mods() & pygame.KMOD_CTRL)
and (pygame.key.get_mods() & pygame.KMOD_SHIFT)
):
filename = prompt_file(save=True)
if filename:
filename += ".level.pkl"
save_pickle(filename, self.last_played_level)
print(f"Saved last sampled level to {filename}")
elif event.key == pygame.K_s and (pygame.key.get_mods() & pygame.KMOD_CTRL):
pcg_state = self._reset_select_shape(pcg_state)
filename = prompt_file(save=True)
if filename:
if filename.endswith(".json"):
export_env_state_to_json(
filename, pcg_state.env_state, self.static_env_params, self.env_params
)
elif not filename.endswith(".pcg.pkl"):
filename += ".pcg.pkl"
save_pickle(filename, pcg_state)
print(f"Saved PCG state to {filename}")
elif event.key == pygame.K_o and (pygame.key.get_mods() & pygame.KMOD_CTRL):
filename = prompt_file(save=False)
if filename:
self._reset_select_shape(pcg_state)
if filename.endswith(".pcg.pkl"):
pcg_state = load_pcg_state_pickle(filename)
pcg_state = expand_pcg_state(pcg_state, self.static_env_params)
print(f"Loaded PCG state from {filename}")
elif filename.endswith(".level.pkl"):
env_state = load_world_state_pickle(filename)
pcg_state = env_state_to_pcg_state(env_state)
print(f"Converted level state to PCG state from {filename}")
elif filename.endswith(".json"):
env_state, new_static_env_params, new_env_params = load_from_json_file(filename)
self._update_params(new_static_env_params, new_env_params)
pcg_state = env_state_to_pcg_state(env_state)
self._reset_triangles()
elif event.key == pygame.K_n and (pygame.key.get_mods() & pygame.KMOD_CTRL):
self._reset_select_shape(pcg_state)
pcg_state = new_pcg_env(self.static_env_params)
self._reset_triangles()
else:
keys.append(event.key)
elif event.type == pygame.KEYUP:
keys_up_this_frame.add(event.key)
if event.type == pygame.MOUSEBUTTONDOWN and self._get_mouse_position_world_space()[0] >= 0:
if event.button == 1:
left_click = True
if event.button == 3:
right_click = True
if event.type == pygame.MOUSEWHEEL:
pcg_state = self._handle_scroll_wheel(pcg_state, event.y)
if self.selected_shape_index == -1:
num = get_numeric_key_pressed(self.pygame_events)
if num is not None:
self.edit_shape_mode = EditMode(num % len(EditMode))
# We have to do these checks outside the loop, otherwise they get triggered multiple times per key press.
if pygame.key.get_mods() & pygame.KMOD_SHIFT:
state = pcg_state.env_state
if pygame.K_m in keys_up_this_frame:
state = self.mutate_world(_rng, state, 1)
if pygame.K_c in keys_up_this_frame:
state, _ = mutate_add_connected_shape(
_rng, state, self.env_params, self.static_env_params, self.ued_params
)
elif pygame.K_s in keys_up_this_frame:
state = mutate_add_shape(_rng, state, self.env_params, self.static_env_params, self.ued_params)
elif pygame.K_p in keys_up_this_frame:
state = mutate_swap_role(_rng, state, self.env_params, self.static_env_params, self.ued_params)
elif pygame.K_r in keys_up_this_frame:
state = mutate_remove_shape(_rng, state, self.env_params, self.static_env_params, self.ued_params)
elif pygame.K_j in keys_up_this_frame:
state = mutate_remove_joint(_rng, state, self.env_params, self.static_env_params, self.ued_params)
elif pygame.K_t in keys_up_this_frame:
state = mutate_toggle_fixture(_rng, state, self.env_params, self.static_env_params, self.ued_params)
elif pygame.K_g in keys_up_this_frame:
state = mutate_add_thruster(_rng, state, self.env_params, self.static_env_params, self.ued_params)
elif pygame.K_l in keys_up_this_frame:
state = mutate_remove_thruster(_rng, state, self.env_params, self.static_env_params, self.ued_params)
elif pygame.K_b in keys_up_this_frame:
state = self.mutate_change_shape_size(
_rng, state, self.env_params, self.static_env_params, self.ued_params
)
elif pygame.K_x in keys_up_this_frame:
state = mutate_change_shape_location(
_rng, state, self.env_params, self.static_env_params, self.ued_params
)
elif pygame.K_k in keys_up_this_frame:
state = self.mutate_change_shape_rotation(
_rng, state, self.env_params, self.static_env_params, self.ued_params
)
pcg_state = pcg_state.replace(env_state=state)
if pygame.K_p in keys_up_this_frame:
global myrng
myrng, _rng = jax.random.split(myrng)
# use the same rng
pcg_state = permute_pcg_state(_rng, pcg_state, self.static_env_params)
if self.edit_shape_mode == EditMode.SELECT: # select a shape
pcg_state = self._edit_select_shape(pcg_state, left_click, right_click, keys)
self._put_state_values_into_gui(pcg_state)
pcg_state = self._select_shape_keyboard_shortcuts(pcg_state, left_click, keys)
else:
pcg_state = self._reset_select_shape(pcg_state) # don't highlight
self._put_state_values_into_gui(pcg_state)
self._show_correct_widgets(None)
if self.edit_shape_mode != EditMode.ADD_TRIANGLE or not self.creating_shape:
self.num_triangle_clicks = 0
if self.edit_shape_mode == EditMode.ADD_CIRCLE:
pcg_state = self._edit_circle(pcg_state, left_click, right_click)
elif self.edit_shape_mode == EditMode.ADD_RECTANGLE:
pcg_state = self._edit_rect(pcg_state, left_click, right_click)
elif self.edit_shape_mode == EditMode.ADD_JOINT:
pcg_state = self._edit_joint(pcg_state, left_click, right_click)
elif self.edit_shape_mode == EditMode.ADD_TRIANGLE:
pcg_state = self._edit_triangle(pcg_state, left_click, right_click)
elif self.edit_shape_mode == EditMode.ADD_THRUSTER:
pcg_state = self._edit_thruster(pcg_state, left_click, right_click)
pcg_state = pcg_state.replace(
env_state=recompute_global_joint_positions(
pcg_state.env_state.replace(
collision_matrix=calculate_collision_matrix(self.static_env_params, pcg_state.env_state.joint),
),
self.static_env_params,
),
env_state_pcg_mask=pcg_state.env_state_pcg_mask.replace(
collision_matrix=jnp.zeros_like(pcg_state.env_state_pcg_mask.collision_matrix)
),
)
return pcg_state
def _update_params(self, new_static_env_params: StaticEnvParams, new_env_params: EnvParams):
self.static_env_params = new_static_env_params.replace(
frame_skip=self.config["frame_skip"], downscale=self.config["downscale"]
)
self.env_params = new_env_params
env = make_kinetix_env_from_name("Kinetix-Entity-MultiDiscrete-v1", static_env_params=self.static_env_params)
self.env = AutoResetWrapper(env, make_reset_function(self.static_env_params))
self._setup_rendering(self.static_env_params, self.env_params)
def _discard_shape_being_created(self, pcg_state):
env_state = pcg_state.env_state
if self.creating_shape:
if self.edit_shape_mode == EditMode.ADD_CIRCLE:
env_state = env_state.replace(
circle=env_state.circle.replace(
active=env_state.circle.active.at[self.creating_shape_index].set(False)
)
)
elif self.edit_shape_mode == EditMode.ADD_RECTANGLE:
env_state = env_state.replace(
polygon=env_state.polygon.replace(
active=env_state.polygon.active.at[self.creating_shape_index].set(False)
)
)
self.creating_shape = False
return pcg_state.replace(env_state=env_state)
def _handle_scroll_wheel(self, pcg_state, y):
if y == 0:
return pcg_state
state = self._discard_shape_being_created(pcg_state)
self.edit_shape_mode = EditMode((self.edit_shape_mode.value + y) % len(EditMode))
return state
def _get_mouse_position_world_space(self):
mouse_pos = pygame.mouse.get_pos()
return (
jnp.array(
[
mouse_pos[0] / self.upscale - self.side_panel_width,
self.static_env_params.screen_dim[1] - mouse_pos[1] / self.upscale,
]
)
/ self.env_params.pixels_per_unit
)
def _get_circles_on_mouse(self, state):
mouse_pos = self._get_mouse_position_world_space()
cis = []
for ci in jnp.arange(self.static_env_params.num_circles)[::-1]:
circle = jax.tree.map(lambda x: x[ci], state.circle)
if not circle.active:
continue
dist = jnp.linalg.norm(mouse_pos - circle.position)
if dist <= circle.radius:
cis.append(ci)
return cis
def _get_revolute_joints_on_mouse(self, state: EnvState):
mouse_pos = self._get_mouse_position_world_space()
ris = []
for ri in jnp.arange(self.static_env_params.num_joints)[::-1]:
joint = jax.tree.map(lambda x: x[ri], state.joint)
if not joint.active:
continue
dist = jnp.linalg.norm(mouse_pos - joint.global_position)
if dist <= 10 / 100: # arbitrary
ris.append(ri)
return ris
def _get_thrusters_on_mouse(self, state: EnvState):
mouse_pos = self._get_mouse_position_world_space()
ris = []
for ri in jnp.arange(self.static_env_params.num_thrusters)[::-1]:
thruster = jax.tree.map(lambda x: x[ri], state.thruster)
if not thruster.active:
continue
dist = jnp.linalg.norm(mouse_pos - thruster.global_position)
if dist <= 16 / 100: # arbitrary
ris.append(ri)
return ris
def _get_joints_attached_to_shape(self, state, shape_index):
r_a = jnp.arange(self.static_env_params.num_joints)[state.joint.a_index == shape_index]
r_b = jnp.arange(self.static_env_params.num_joints)[state.joint.b_index == shape_index]
t = jnp.arange(self.static_env_params.num_thrusters)[state.thruster.object_index == shape_index]
return jnp.concatenate([r_a, r_b], axis=0), t
def _edit_thruster(self, pcg_state: PCGState, left_click: bool, right_click: bool):
if not self.creating_shape and (1 - pcg_state.env_state.thruster.active.astype(int)).sum() == 0:
if not right_click:
return pcg_state
thruster_pos = self._get_mouse_position_world_space()
idx = -1
for ri in self._get_polygons_on_mouse(pcg_state.env_state):
r = jax.tree.map(lambda x: x[ri], pcg_state.env_state.polygon)
thruster_pos = snap_to_polygon_center_line(r, thruster_pos)
relative_pos = jnp.matmul(rmat(r.rotation).transpose((1, 0)), thruster_pos - r.position)
idx = ri
break
if idx == -1:
for ci in self._get_circles_on_mouse(pcg_state.env_state):
c = jax.tree.map(lambda x: x[ci], pcg_state.env_state.circle)
thruster_pos = snap_to_center(c, thruster_pos)
thruster_pos = snap_to_circle_center_line(c, thruster_pos)
relative_pos = thruster_pos - c.position
idx = ci + self.static_env_params.num_polygons
break
if left_click:
if self.creating_shape:
self.creating_shape = False
else:
if idx >= 0:
self.creating_shape = True
self.creating_shape_position = thruster_pos
self.creating_shape_index = jnp.argmin(pcg_state.env_state.thruster.active)
shape = select_shape(pcg_state.env_state, idx, self.static_env_params)
def _add_thruster_to_state(state):
state = state.replace(
thruster=state.thruster.replace(
object_index=state.thruster.object_index.at[self.creating_shape_index].set(idx),
relative_position=state.thruster.relative_position.at[self.creating_shape_index].set(
relative_pos
),
power=state.thruster.power.at[self.creating_shape_index].set(
1.0 / jax.lax.select(shape.inverse_mass == 0, 1.0, shape.inverse_mass)
),
active=state.thruster.active.at[self.creating_shape_index].set(True),
global_position=state.thruster.global_position.at[self.creating_shape_index].set(
thruster_pos
),
rotation=state.thruster.rotation.at[self.creating_shape_index].set(0.0),
),
thruster_bindings=state.thruster_bindings.at[self.creating_shape_index].set(0),
)
return state
pcg_state = pcg_state.replace(
env_state=_add_thruster_to_state(pcg_state.env_state),
env_state_max=_add_thruster_to_state(pcg_state.env_state_max),
)
elif right_click:
for ti in self._get_thrusters_on_mouse(pcg_state.env_state):
def _remove_thruster_from_state(state):
return state.replace(
thruster=state.thruster.replace(active=state.thruster.active.at[ti].set(False))
)
return pcg_state.replace(
env_state=_remove_thruster_from_state(pcg_state.env_state),
env_state_max=_remove_thruster_from_state(pcg_state.env_state_max),
)
else:
if self.creating_shape:
curr_pos = self._get_mouse_position_world_space()
normal = pcg_state.env_state.thruster.relative_position[self.creating_shape_index]
angle = jnp.arctan2(normal[1], normal[0])
relative_pos = curr_pos - self.creating_shape_position
# rotation = jnp.arctan2(relative_pos[1], relative_pos[0])
rotation = jnp.pi + jnp.arctan2(relative_pos[1], relative_pos[0]) + angle
angle_round = jnp.round(rotation / (jnp.pi / 2))
angle_norm = rotation / (jnp.pi / 2)
if jnp.abs(angle_round - angle_norm) < 0.3:
rotation = angle_round * (jnp.pi / 2)
def _update_thruster_rotation(state):
return state.replace(
thruster=state.thruster.replace(
rotation=state.thruster.rotation.at[self.creating_shape_index].set(rotation - angle),
)
)
pcg_state = pcg_state.replace(
env_state=_update_thruster_rotation(pcg_state.env_state),
env_state_max=_update_thruster_rotation(pcg_state.env_state_max),
)
else:
pass
return pcg_state
def _edit_circle(self, pcg_state: PCGState, left_click: bool, right_click: bool):
if right_click:
for ci in self._get_circles_on_mouse(pcg_state.env_state):
attached_j, attached_t = self._get_joints_attached_to_shape(
pcg_state.env_state, ci + self.static_env_params.num_polygons
)
def _remove_circle_from_state(state):
return state.replace(
circle=state.circle.replace(active=state.circle.active.at[ci].set(False)),
joint=state.joint.replace(active=state.joint.active.at[attached_j].set(False)),
thruster=state.thruster.replace(active=state.thruster.active.at[attached_t].set(False)),
)
env_state = _remove_circle_from_state(pcg_state.env_state)
env_state_pcg_mask = _remove_circle_from_state(pcg_state.env_state_pcg_mask)
env_state_max = _remove_circle_from_state(pcg_state.env_state_max)
env_state = env_state.replace(
collision_matrix=calculate_collision_matrix(self.static_env_params, env_state.joint)
)
return PCGState(
env_state=env_state,
env_state_pcg_mask=env_state_pcg_mask,
env_state_max=env_state_max,
tied_together=pcg_state.tied_together,
)
if not self.creating_shape and (1 - pcg_state.env_state.circle.active.astype(int)).sum() == 0:
return pcg_state
radius = jnp.linalg.norm(self._get_mouse_position_world_space() - self.create_shape_position)
radius = jnp.clip(radius, 5.0 / self.env_params.pixels_per_unit, self.static_env_params.max_shape_size / 2)
def _add_circle(state, highlight):
state = state.replace(
circle=state.circle.replace(
position=state.circle.position.at[self.creating_shape_index].set(self.create_shape_position),
velocity=state.circle.velocity.at[self.creating_shape_index].set(jnp.array([0.0, 0.0])),
radius=state.circle.radius.at[self.creating_shape_index].set(radius),
inverse_mass=state.circle.inverse_mass.at[self.creating_shape_index].set(1.0),
inverse_inertia=state.circle.inverse_inertia.at[self.creating_shape_index].set(1.0),
active=state.circle.active.at[self.creating_shape_index].set(True),
collision_mode=state.circle.collision_mode.at[self.creating_shape_index].set(1),
),
circle_shape_roles=state.circle_shape_roles.at[self.creating_shape_index].set(0),
circle_highlighted=state.circle_highlighted.at[self.creating_shape_index].set(highlight),
circle_densities=state.circle_densities.at[self.creating_shape_index].set(1.0),
)
return state
if left_click:
if self.creating_shape:
env_state = _add_circle(pcg_state.env_state, False)
env_state_max = _add_circle(pcg_state.env_state_max, False)
env_state = recalculate_mass_and_inertia(
env_state,
self.static_env_params,
env_state.polygon_densities,
env_state.circle_densities,
)
env_state_max = recalculate_mass_and_inertia(
env_state_max,
self.static_env_params,
env_state.polygon_densities,
env_state.circle_densities,
)
pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)
self.creating_shape = False
else:
self.creating_shape_index = jnp.argmin(pcg_state.env_state.circle.active)
self.create_shape_position = self._get_mouse_position_world_space()
self.creating_shape = True
else:
if self.creating_shape:
env_state = _add_circle(pcg_state.env_state, True)
env_state_max = _add_circle(pcg_state.env_state_max, True)
pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)
else:
pass
return pcg_state
def _get_polygons_on_mouse(self, state, n_vertices=None):
# n_vertices=None selects both triangles and quads
ris = []
mouse_pos = self._get_mouse_position_world_space()
for ri in jnp.arange(self.static_env_params.num_polygons)[::-1]:
polygon = jax.tree.map(lambda x: x[ri], state.polygon)
if (not polygon.active) or ((n_vertices is not None) and polygon.n_vertices != n_vertices):
continue
mpos = rmat(-polygon.rotation) @ (mouse_pos - polygon.position)
def _signed_line_distance(a, b, c):
return (b[0] - a[0]) * (c[1] - a[1]) - (b[1] - a[1]) * (c[0] - a[0])
inside = True
for fi in range(polygon.n_vertices):
v1 = polygon.vertices[fi]
v2 = polygon.vertices[(fi + 1) % polygon.n_vertices]
if _signed_line_distance(mpos, v1, v2) > 0:
inside = False
if inside:
ris.append(ri)
return ris
def _edit_rect(self, pcg_state: PCGState, left_click: bool, right_click: bool):
if right_click:
for ri in self._get_polygons_on_mouse(pcg_state.env_state, n_vertices=4):
attached_j, attached_t = self._get_joints_attached_to_shape(pcg_state.env_state, ri)
def _remove_rect_from_state(state):
state = state.replace(
polygon=state.polygon.replace(
active=state.polygon.active.at[ri].set(False),
rotation=state.polygon.rotation.at[ri].set(0.0),
),
joint=state.joint.replace(active=state.joint.active.at[attached_j].set(False)),
thruster=state.thruster.replace(active=state.thruster.active.at[attached_t].set(False)),
)
return state
env_state = _remove_rect_from_state(pcg_state.env_state)
env_state_max = _remove_rect_from_state(pcg_state.env_state_max)
env_state_pcg_mask = _remove_rect_from_state(pcg_state.env_state_pcg_mask)
env_state = env_state.replace(
collision_matrix=calculate_collision_matrix(self.static_env_params, env_state.joint)
)
return PCGState(
env_state=env_state,
env_state_max=env_state_max,
env_state_pcg_mask=env_state_pcg_mask,
tied_together=pcg_state.tied_together,
)
if not self.creating_shape and (1 - pcg_state.env_state.polygon.active.astype(int)).sum() == 0:
return pcg_state
diff = (self._get_mouse_position_world_space() - self.create_shape_position) / 2
diff = jnp.clip(
diff,
-(self.static_env_params.max_shape_size / 2) / jnp.sqrt(2),
(self.static_env_params.max_shape_size / 2) / jnp.sqrt(2),
)
half_dim = jnp.abs(diff)
half_dim = jnp.clip(half_dim, a_min=5.0 / self.env_params.pixels_per_unit)
vertices = rectangle_vertices(half_dim)
def _add_rect_to_state(state, highlight):
state = state.replace(
polygon=state.polygon.replace(
position=state.polygon.position.at[self.creating_shape_index].set(
self.create_shape_position + diff
),
velocity=state.polygon.velocity.at[self.creating_shape_index].set(jnp.array([0.0, 0.0])),
vertices=state.polygon.vertices.at[self.creating_shape_index].set(vertices),
inverse_mass=state.polygon.inverse_mass.at[self.creating_shape_index].set(1.0),
inverse_inertia=state.polygon.inverse_inertia.at[self.creating_shape_index].set(1.0),
active=state.polygon.active.at[self.creating_shape_index].set(True),
collision_mode=state.polygon.collision_mode.at[self.creating_shape_index].set(1),
n_vertices=state.polygon.n_vertices.at[self.creating_shape_index].set(4),
),
polygon_shape_roles=state.polygon_shape_roles.at[self.creating_shape_index].set(0),
polygon_highlighted=state.polygon_highlighted.at[self.creating_shape_index].set(highlight),
polygon_densities=state.polygon_densities.at[self.creating_shape_index].set(1.0),
)
return state
if left_click:
if self.creating_shape:
env_state = _add_rect_to_state(pcg_state.env_state, False)
env_state_max = _add_rect_to_state(pcg_state.env_state_max, False)
env_state = recalculate_mass_and_inertia(
env_state,
self.static_env_params,
env_state.polygon_densities,
env_state.circle_densities,
)
env_state_max = recalculate_mass_and_inertia(
env_state_max,
self.static_env_params,
env_state.polygon_densities,
env_state.circle_densities,
)
pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)
self.creating_shape = False
else:
self.creating_shape_index = jnp.argmin(pcg_state.env_state.polygon.active)
self.create_shape_position = self._get_mouse_position_world_space()
self.creating_shape = True
else:
if self.creating_shape:
env_state = _add_rect_to_state(pcg_state.env_state, True)
env_state_max = _add_rect_to_state(pcg_state.env_state_max, True)
pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)
else:
pass
return pcg_state
def _reset_triangles(self):
self.triangle_order = jnp.array([0, 1, 2])
self.num_triangle_clicks = 0
self.creating_shape = False
def _edit_triangle(self, pcg_state: PCGState, left_click: bool, right_click: bool):
if right_click:
self.num_triangle_clicks = 0
for ri in self._get_polygons_on_mouse(pcg_state.env_state, n_vertices=3):
attached_r, attached_f = self._get_joints_attached_to_shape(pcg_state.env_state, ri)
def _remove_triangle_from_state(state):
state = state.replace(
polygon=state.polygon.replace(
active=state.polygon.active.at[ri].set(False),
rotation=state.polygon.rotation.at[ri].set(0.0),
),
joint=state.joint.replace(active=state.joint.active.at[attached_r].set(False)),
)
return state
env_state = _remove_triangle_from_state(pcg_state.env_state)
env_state_max = _remove_triangle_from_state(pcg_state.env_state_max)
env_state_pcg_mask = _remove_triangle_from_state(pcg_state.env_state_pcg_mask)
env_state = env_state.replace(
collision_matrix=calculate_collision_matrix(self.static_env_params, env_state.joint)
)
pcg_state = PCGState(
env_state=env_state,
env_state_max=env_state_max,
env_state_pcg_mask=env_state_pcg_mask,
tied_together=pcg_state.tied_together,
)
return pcg_state
if not self.creating_shape and (1 - pcg_state.env_state.polygon.active.astype(int)).sum() == 0:
return pcg_state
def get_correct_center_two_verts(verts):
return (jnp.max(verts, axis=0) + jnp.min(verts, axis=0)) / 2
def order_clockwise(verts, loose_ordering=False):
# verts has shape (3, 2), order them clockwise.
# https://stackoverflow.com/questions/51074984/sort-vertices-in-clockwise-order
# Calculate centroid
centroid = jnp.mean(verts, axis=0)
# Calculate angles
angles = jnp.round(jnp.arctan2(verts[:, 1] - centroid[1], verts[:, 0] - centroid[0]), 2)
# Order vertices
order = jnp.argsort(-angles, stable=True)
if loose_ordering:
order = jnp.arange(len(order))
ans = verts[order]
# order is of shape (2, ) or (3, ). I want it to always be of shape 3
if len(order) < 3:
order = jnp.concatenate([order, jnp.array([2])])
return ans, order
def do_triangle_n_click(pcg_state, how_many_clicks, is_on_a_click=False):
n = how_many_clicks
# if we must keep them clockwise all the time, then the one we edit / move around may have varying indices.
current_index_to_change = self.triangle_order[n]
sign = 1
idxs = jnp.arange(n + 1)
idxs_to_allow = idxs[~(idxs == current_index_to_change)]
# Get the new vertex and clip its position
new_tentative_vert = (
self._get_mouse_position_world_space() - pcg_state.env_state.polygon.position[self.creating_shape_index]
)
new_tentative_vert = jnp.clip(
new_tentative_vert,
jnp.max(pcg_state.env_state.polygon.vertices[self.creating_shape_index, idxs_to_allow], axis=0)
- self.static_env_params.max_shape_size * 0.8,
jnp.min(pcg_state.env_state.polygon.vertices[self.creating_shape_index, idxs_to_allow], axis=0)
+ self.static_env_params.max_shape_size * 0.8,
)
new_verts = pcg_state.env_state.polygon.vertices.at[self.creating_shape_index, current_index_to_change].set(
new_tentative_vert
)
new_center_two = get_correct_center_two_verts(new_verts[self.creating_shape_index, : n + 1])
_, new_center_three = calc_inverse_mass_polygon(
new_verts[self.creating_shape_index],
3,
self.static_env_params,
1.0,
)
new_center = jax.lax.select(n == 1, new_center_two, new_center_three)
new_verts = new_verts.at[self.creating_shape_index].add(-sign * new_center)
vvs = new_verts[self.creating_shape_index, : n + 1]
ordered_vertices, new_permutation = order_clockwise(vvs, loose_ordering=not is_on_a_click)
self.triangle_order = self.triangle_order[new_permutation]
new_verts = new_verts.at[self.creating_shape_index, : n + 1].set(ordered_vertices)
env_state = pcg_state.env_state.replace(
polygon=pcg_state.env_state.polygon.replace(
vertices=new_verts,
position=pcg_state.env_state.polygon.position.at[self.creating_shape_index].add(sign * new_center),
n_vertices=pcg_state.env_state.polygon.n_vertices.at[self.creating_shape_index].set(n + 1),
),
)
env_state_max = pcg_state.env_state_max.replace(
polygon=pcg_state.env_state_max.polygon.replace(
vertices=new_verts,
position=pcg_state.env_state_max.polygon.position.at[self.creating_shape_index].add(
sign * new_center
),
n_vertices=pcg_state.env_state_max.polygon.n_vertices.at[self.creating_shape_index].set(n + 1),
),
)
pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)
return pcg_state
if left_click:
if self.creating_shape:
assert 3 > self.num_triangle_clicks > 0
if self.num_triangle_clicks == 1:
pcg_state = do_triangle_n_click(pcg_state, 1, is_on_a_click=True)
self.num_triangle_clicks += 1
else: # this finishes it
pcg_state = do_triangle_n_click(pcg_state, 2, is_on_a_click=True)
self.creating_shape = False
self.num_triangle_clicks = 0
pcg_state = pcg_state.replace(
env_state=recalculate_mass_and_inertia(
pcg_state.env_state,
self.static_env_params,
pcg_state.env_state.polygon_densities,
pcg_state.env_state.circle_densities,
)
)
else:
self.triangle_order = jnp.array([0, 1, 2])
self.creating_shape_index = jnp.argmin(pcg_state.env_state.polygon.active)
self.create_shape_position = self._get_mouse_position_world_space()
self.creating_shape = True
self.num_triangle_clicks = 1
vertices = jnp.zeros((self.static_env_params.max_polygon_vertices, 2), dtype=jnp.float32)
def _add_triangle_to_state(state):
state = state.replace(
polygon=state.polygon.replace(
position=state.polygon.position.at[self.creating_shape_index].set(
self.create_shape_position
),
velocity=state.polygon.velocity.at[self.creating_shape_index].set(jnp.array([0.0, 0.0])),
vertices=state.polygon.vertices.at[self.creating_shape_index].set(vertices),
inverse_mass=state.polygon.inverse_mass.at[self.creating_shape_index].set(1.0),
inverse_inertia=state.polygon.inverse_inertia.at[self.creating_shape_index].set(1.0),
active=state.polygon.active.at[self.creating_shape_index].set(True),
n_vertices=state.polygon.n_vertices.at[self.creating_shape_index].set(1),
),
polygon_shape_roles=state.polygon_shape_roles.at[self.creating_shape_index].set(0),
polygon_highlighted=state.polygon_highlighted.at[self.creating_shape_index].set(False),
polygon_densities=state.polygon_densities.at[self.creating_shape_index].set(1.0),
)
return state
pcg_state = pcg_state.replace(
env_state=_add_triangle_to_state(pcg_state.env_state),
env_state_max=_add_triangle_to_state(pcg_state.env_state_max),
)
elif self.creating_shape:
assert 1 <= self.num_triangle_clicks <= 2
pcg_state = do_triangle_n_click(
pcg_state, self.num_triangle_clicks, is_on_a_click=self.num_triangle_clicks == 1
)
return pcg_state
def _edit_joint(self, pcg_state: PCGState, left_click: bool, right_click: bool):
if left_click and pcg_state.env_state.joint.active.all():
return pcg_state
if left_click:
joint_index = jnp.argmin(pcg_state.env_state.joint.active)
joint_position = self._get_mouse_position_world_space()
# reverse them so that the joint order and rendering order remains the same.
# We want the first shape to have a lower index than the second shape, with circles always having higher indices compared to rectangles.
circles = self._get_circles_on_mouse(pcg_state.env_state)[::-1]
rects = self._get_polygons_on_mouse(pcg_state.env_state)[::-1]
if len(rects) + len(circles) >= 2:
r1 = len(rects) >= 1
r2 = len(rects) >= 2
a_index = rects[0] if r1 else circles[0] # + self.static_env_params.num_polygons
b_index = rects[r1 * 1] if r2 else circles[1 - 1 * r1] # + self.static_env_params.num_polygons
a_shape = pcg_state.env_state.polygon if r1 else pcg_state.env_state.circle
b_shape = pcg_state.env_state.polygon if r2 else pcg_state.env_state.circle
a = jax.tree.map(lambda x: x[a_index], a_shape)
b = jax.tree.map(lambda x: x[b_index], b_shape)
a_index += (not r1) * self.static_env_params.num_polygons
b_index += (not r2) * self.static_env_params.num_polygons
joint_position = snap_to_center(a, joint_position)
joint_position = snap_to_center(b, joint_position)
a_relative_pos = jnp.matmul(rmat(a.rotation).transpose((1, 0)), joint_position - a.position)
b_relative_pos = jnp.matmul(rmat(b.rotation).transpose((1, 0)), joint_position - b.position)
def _add_joint_to_state(state):
state = state.replace(
joint=state.joint.replace(
a_index=state.joint.a_index.at[joint_index].set(a_index),
b_index=state.joint.b_index.at[joint_index].set(b_index),
a_relative_pos=state.joint.a_relative_pos.at[joint_index].set(a_relative_pos),
b_relative_pos=state.joint.b_relative_pos.at[joint_index].set(b_relative_pos),
active=state.joint.active.at[joint_index].set(True),
global_position=state.joint.global_position.at[joint_index].set(joint_position),
motor_on=state.joint.motor_on.at[joint_index].set(True),
motor_speed=state.joint.motor_speed.at[joint_index].set(1.0),
motor_power=state.joint.motor_power.at[joint_index].set(1.0),
rotation=state.joint.rotation.at[joint_index].set(b.rotation - a.rotation),
)
)
return state
env_state = _add_joint_to_state(pcg_state.env_state)
env_state_max = _add_joint_to_state(pcg_state.env_state_max)
env_state = env_state.replace(
collision_matrix=calculate_collision_matrix(self.static_env_params, env_state.joint)
)
pcg_state = pcg_state.replace(env_state=env_state, env_state_max=env_state_max)
return pcg_state
def _reset_select_shape(self, pcg_state):
pcg_state = pcg_state.replace(
env_state=pcg_state.env_state.replace(
polygon_highlighted=jnp.zeros_like(pcg_state.env_state.polygon_highlighted),
circle_highlighted=jnp.zeros_like(pcg_state.env_state.circle_highlighted),
)
)
self.selected_shape_index = -1
self.selected_shape_type = ObjectType.POLYGON
self._hide_all_widgets()
return pcg_state
def _hide_all_widgets(self):
for widget in self.all_widgets.values():
for w in widget.values():
w.hide()
def _show_correct_widgets(self, type: ObjectType | None, do_tie_ui: bool = False):
for widget in self.all_widgets["GENERAL"].values():
widget.show()
if do_tie_ui:
for widget in self.all_widgets["TIE_TOGETHER"].values():
widget.show()
n = len(self.all_selected_shapes)
# {[int(i) for (i, t) in self.all_selected_shapes]}
self.all_widgets["GENERAL"]["lblGeneral"].setText(f"Selected {n} Objects")
return
for widget in self.all_widgets[type].values():
widget.show()
if type is None:
self.all_widgets["GENERAL"]["lblGeneral"].setText(f"Global")
else:
self.all_widgets["GENERAL"]["lblGeneral"].setText(f"{type.name} (idx {self.selected_shape_index})")
def _select_shape_keyboard_shortcuts(self, pcg_state: PCGState, left_click: bool, keys: list[int]):
if left_click:
return pcg_state
if len(keys) != 0 and self.selected_shape_index != -1:
s = 1.0
ang_s = 0.1
vel = jnp.array([0.0, 0.0])
angular_vel = 0.0
should_toggle_fixed = False
should_toggle_collidable = False
change_angle = 0
def add_step(widget_name, direction, speed=10, overwrite_amount=None):
widget = self.all_widgets[self.selected_shape_type][widget_name]
val = widget.getValue()
step = widget.step
amount_to_add = overwrite_amount or step * direction * speed
widget.setValue(jnp.clip(val + amount_to_add, widget.min, widget.max))
if pygame.K_w in keys:
add_step("sldPosition_Y", 1)
if pygame.K_s in keys:
add_step("sldPosition_Y", -1)
if pygame.K_a in keys:
add_step("sldPosition_X", -1)
if pygame.K_d in keys:
add_step("sldPosition_X", 1)
if pygame.K_q in keys:
add_step("sldRotation", 1)
if pygame.K_e in keys:
add_step("sldRotation", -1)
if pygame.K_f in keys:
self.all_widgets[self.selected_shape_type]["tglFixate"].toggle()
if pygame.K_c in keys and not (pygame.key.get_mods() & pygame.KMOD_CTRL):
widget = self.all_widgets[self.selected_shape_type]["sldCollidability"]
curr_val = int(widget.getValue())
widget.setValue((curr_val + 1) % (widget.max + 1))
if pygame.K_r in keys and not (pygame.key.get_mods() & pygame.KMOD_CTRL):
widget = self.all_widgets[self.selected_shape_type]["sldRole"]
curr_val = int(widget.getValue())
widget.setValue((curr_val + 1) % (widget.max + 1))
if pygame.K_LEFTBRACKET in keys:
add_step("sldRotation", 1, 10, jnp.pi / 4)
if pygame.K_RIGHTBRACKET in keys:
add_step("sldRotation", -1, 10, -jnp.pi / 4)
if pygame.K_c in keys and (pygame.key.get_mods() & pygame.KMOD_CTRL):
# copy
if self.selected_shape_type == ObjectType.POLYGON: # rect
if not self.pcg_state.env_state.polygon.active.all():
where_to_add = jnp.argmin(pcg_state.env_state.polygon.active)
if where_to_add < self.static_env_params.num_polygons:
def _copy_polygon(state, shift):
state = state.replace(
polygon=jax.tree.map(
lambda x: x.at[where_to_add].set(x[self.selected_shape_index]), state.polygon
)
)
if shift:
state = state.replace(
polygon=state.polygon.replace(
position=state.polygon.position.at[where_to_add].add(0.1),
),
polygon_highlighted=state.polygon_highlighted.at[where_to_add].set(False),
)
return state
pcg_state = pcg_state.replace(
env_state=_copy_polygon(pcg_state.env_state, shift=True),
env_state_max=_copy_polygon(pcg_state.env_state_max, shift=True),
env_state_pcg_mask=_copy_polygon(pcg_state.env_state_pcg_mask, shift=False),
)
elif self.selected_shape_type == ObjectType.CIRCLE: # circle
if not self.pcg_state.env_state.circle.active.all():
where_to_add = jnp.argmin(pcg_state.env_state.circle.active)
if where_to_add < self.static_env_params.num_circles:
def _copy_circle(state, shift=True):
state = state.replace(
circle=jax.tree.map(
lambda x: x.at[where_to_add].set(x[self.selected_shape_index]), state.circle
)
)
if shift:
state = state.replace(
circle=state.circle.replace(
position=state.circle.position.at[where_to_add].add(0.1),
),
circle_highlighted=state.circle_highlighted.at[where_to_add].set(False),
)
return state
pcg_state = pcg_state.replace(
env_state=_copy_circle(pcg_state.env_state),
env_state_max=_copy_circle(pcg_state.env_state_max),
env_state_pcg_mask=_copy_circle(pcg_state.env_state_pcg_mask, shift=False),
)
if self.selected_shape_index >= 0:
num = get_numeric_key_pressed(self.pygame_events)
if num is not None:
if self.selected_shape_type in [ObjectType.CIRCLE, ObjectType.POLYGON]:
self.all_widgets[self.selected_shape_type]["sldRole"].setValue(num % 4)
elif self.selected_shape_type == ObjectType.JOINT:
self.all_widgets[self.selected_shape_type]["sldColour"].setValue(
num % self.static_env_params.num_motor_bindings
)
elif self.selected_shape_type == ObjectType.THRUSTER:
self.all_widgets[self.selected_shape_type]["sldColour"].setValue(
num % self.static_env_params.num_thruster_bindings
)
return pcg_state
def _edit_select_shape(self, pcg_state: PCGState, left_click: bool, right_click: bool, keys: list[int]):
def _find_shape(pcg_state):
found_shape = False
selected_shape_index, selected_shape_type = -1, ObjectType.POLYGON
for ri in self._get_revolute_joints_on_mouse(pcg_state.env_state):
selected_shape_index = ri
selected_shape_type = ObjectType.JOINT
found_shape = True
break
if not found_shape:
for ti in self._get_thrusters_on_mouse(pcg_state.env_state):
selected_shape_index = ti
selected_shape_type = ObjectType.THRUSTER
found_shape = True
break
if not found_shape:
for ri in self._get_polygons_on_mouse(pcg_state.env_state):
pcg_state = pcg_state.replace(
env_state=pcg_state.env_state.replace(
polygon_highlighted=pcg_state.env_state.polygon_highlighted.at[ri].set(True),
)
)
selected_shape_index = ri
selected_shape_type = ObjectType.POLYGON
found_shape = True
break
if not found_shape:
for ci in self._get_circles_on_mouse(pcg_state.env_state):
pcg_state = pcg_state.replace(
env_state=pcg_state.env_state.replace(
circle_highlighted=pcg_state.env_state.circle_highlighted.at[ci].set(True),
)
)
selected_shape_index = ci
selected_shape_type = ObjectType.CIRCLE
found_shape = True
break
return selected_shape_index, selected_shape_type, found_shape, pcg_state
if found_shape and self.selected_shape_type in self.all_widgets:
self._show_correct_widgets(self.selected_shape_type)
# if left and shift
if left_click and (pygame.key.get_mods() & pygame.KMOD_SHIFT):
# This is trying to select multiple things.
idx, type, found, pcg_state = _find_shape(pcg_state)
if found:
t = (idx, type)
if t in self.all_selected_shapes:
self.all_selected_shapes.remove(t)
else:
self.all_selected_shapes.append(t)
self._hide_all_widgets()
self._show_correct_widgets(None, do_tie_ui=True)
elif left_click:
self.all_selected_shapes = []
self._hide_all_widgets()
pcg_state = self._reset_select_shape(pcg_state)
self.selected_shape_index, self.selected_shape_type, found_shape, pcg_state = _find_shape(pcg_state)
if found_shape:
self.all_selected_shapes = [(self.selected_shape_index, self.selected_shape_type)]
if self.selected_shape_type in self.all_widgets:
self._show_correct_widgets(self.selected_shape_type)
if self.selected_shape_index < 0:
self._show_correct_widgets(None)
return pcg_state
def render(self, env_state):
# Clear
self.screen_surface.fill((0, 0, 0))
if self.is_editing:
pixels = self._render_fn_edit(env_state)
else:
pixels = self._render_fn(env_state)
pixels = self._render_edit_overlay_fn(pixels, self.is_editing, self.edit_shape_mode.value)
surface = pygame.surfarray.make_surface(np.array(pixels))
self.screen_surface.blit(surface, (0, 0))
def is_quit_requested(self):
for event in self.pygame_events:
if event.type == pygame.QUIT:
return True
return False
@hydra.main(version_base=None, config_path="../configs", config_name="editor")
def main(config):
config = normalise_config(OmegaConf.to_container(config), "EDITOR", editor_config=True)
env_params, static_env_params = generate_params_from_config(config)
static_env_params = static_env_params.replace(frame_skip=config["frame_skip"], downscale=config["downscale"])
config["env_params"] = to_state_dict(env_params)
config["static_env_params"] = to_state_dict(static_env_params)
env = make_kinetix_env_from_name("Kinetix-Entity-MultiDiscrete-v1", static_env_params=static_env_params)
env = AutoResetWrapper(env, make_reset_function(static_env_params))
seed = config["seed"]
print("seed", seed)
rng = jax.random.PRNGKey(seed)
outer_timer = tmr()
editor = Editor(env, env_params, config, upscale=config["upscale"])
time_e = tmr()
print("Took {:2f}s to create editor".format(time_e - outer_timer))
clock = pygame.time.Clock()
while not editor.is_quit_requested():
rng, _rng = jax.random.split(rng)
editor.update(_rng)
clock.tick(config["fps"])
if __name__ == "__main__":
main()