diff --git a/self_organising_systems/biomakerca/agent_logic.py b/self_organising_systems/biomakerca/agent_logic.py index f5f66a6..a9d8238 100644 --- a/self_organising_systems/biomakerca/agent_logic.py +++ b/self_organising_systems/biomakerca/agent_logic.py @@ -82,18 +82,26 @@ def excl_f(self, key: KeyType, perc: PerceivedData, params: AgentProgramType Return a ExclusiveInterface. """ pass - + @abstractmethod def repr_f(self, key: KeyType, perc: PerceivedData, params: AgentProgramType ) -> ReproduceInterface: """Perform a reproduce function. - + params must be only the reproduce params, not all of them. - + Return a ReproduceInterface. """ pass + @abstractmethod + def get_sex(self, params: AgentProgramType): + """Extract the sex parameter from params. + + The sex parameter is a jax int32. + """ + pass + def __str__(self): return stringify_class(self) @@ -127,17 +135,39 @@ class BasicAgentLogic(AgentLogic): The cells can optionally perceive agent_ids. If so, they will never give nutrients to cells with different agent_ids. + + sex_sensitivity is a parameter that is used to initialize and detect the sex + of a DNA. Check get_sex to see how it is used. + I recommend having a sex_sensitivity equal to 1/sd of the mutator used. """ - def __init__(self, config: EnvConfig, perceive_ids=True, minimal_net=False): + def __init__(self, config: EnvConfig, perceive_ids=True, minimal_net=False, + make_asexual_flowers_likely=True, + make_sexual_flowers_likely=True, + sex_sensitivity=1000., init_noise=None): + """Constructor. + + make_asexual_flowers_likely and make_sexual_flowers_likely are checked to + initialize the parameters to encourage a certain kind of flowers to spawn. + + If init_noise is not None, the initialization for !minimal_net will + initialize several parameters with a certain noise, with it being + init_noise * glorot_initialization of weight matrices. + if init_noise is None, the initial network is a homomorphism of a + minimal_nel (same outputs). + """ self.config = config # the types are perceived as one-hot vectors. self.n_types = len(config.etd.types.keys()) - self.n_spec = 4 # specializations of agents + self.n_spec = len(config.etd.agent_types) # specializations of agents # Whether agent ids are perceivable by the agent. # if set to true, agents do not give energy to agents with different ids. self.perceive_ids = perceive_ids self.minimal_net = minimal_net + self.make_asexual_flowers_likely = make_asexual_flowers_likely + self.make_sexual_flowers_likely = make_sexual_flowers_likely + self.sex_sensitivity = sex_sensitivity + self.init_noise = init_noise self.state_clip_val = 3. ## Parallel op: @@ -207,11 +237,16 @@ def split_params_f( params, (n_par_params, n_par_params + self.excl_num_params), axis=-1) + def _cond_init(self, key, shape): + if self.init_noise is not None: + return glorot_normal(batch_axis=0)(key, shape) * self.init_noise + return jp.zeros(shape) + def dsm_init(self, key): if self.minimal_net: # in this case, the agent cannot modify its internal state. return (jp.empty(0),) - + # Set initial effect on state to zero. # a 2 layer NN. # # We look at proportions of neighboring cells and your internal state. @@ -225,8 +260,8 @@ def dsm_init(self, key): ku, key = jr.split(key) dw0 = glorot_normal(batch_axis=0)(ku, (self.n_spec, insize, hsize)) db0 = jp.zeros((self.n_spec, hsize)) - # output is defaulted to zero. - dw1 = jp.zeros((self.n_spec, hsize, outsize)) + key, ku = jr.split(key) + dw1 = self._cond_init(ku, (self.n_spec, hsize, outsize)) db1 = jp.zeros((self.n_spec, outsize)) return (dw0, db0, dw1, db1) @@ -281,18 +316,30 @@ def nsl_init(self, key): w = w.at[:, etd.types.AGENT_FLOWER, spec_idxs.AGENT_UNSPECIALIZED].set( div / 8.0 ) + w = w.at[:, etd.types.AGENT_FLOWER_SEXUAL, + spec_idxs.AGENT_UNSPECIALIZED].set(div / 8.0) w = w.at[:, etd.types.EARTH, spec_idxs.AGENT_UNSPECIALIZED].set( -div / 8.0 ) w = w.at[:, etd.types.AIR, spec_idxs.AGENT_UNSPECIALIZED].set( -div / 8.0 ) - # Flowers only grow if they are surrounded by leaves and some air. - w = w.at[:, etd.types.AGENT_LEAF, spec_idxs.AGENT_FLOWER].set(div / 4.0) - w = w.at[:, etd.types.AIR, spec_idxs.AGENT_FLOWER].set(div / 2.0) + + if self.make_asexual_flowers_likely: + # Flowers only grow if they are surrounded by leaves and some air. + w = w.at[:, etd.types.AGENT_LEAF, spec_idxs.AGENT_FLOWER].set(div / 4.0) + w = w.at[:, etd.types.AIR, spec_idxs.AGENT_FLOWER].set(div / 2.0) + + if self.make_sexual_flowers_likely: + # Sexual flowers start with the same chance of being spawned. + w = w.at[:, etd.types.AGENT_LEAF, spec_idxs.AGENT_FLOWER_SEXUAL].set( + div / 4.0) + w = w.at[:, etd.types.AIR, spec_idxs.AGENT_FLOWER_SEXUAL].set(div / 2.0) # If you are a flower, never change! w = w.at[spec_idxs.AGENT_FLOWER, :, spec_idxs.AGENT_FLOWER].set(div) + w = w.at[spec_idxs.AGENT_FLOWER_SEXUAL, :, + spec_idxs.AGENT_FLOWER_SEXUAL].set(div) if self.minimal_net: return w, b @@ -308,8 +355,8 @@ def nsl_init(self, key): ku, key = jr.split(key) dw0 = glorot_normal(batch_axis=0)(ku, (self.n_spec, insize, hsize)) db0 = jp.zeros((self.n_spec, hsize)) - # output is defaulted to zero. - dw1 = jp.zeros((self.n_spec, hsize, outsize)) + key, ku = jr.split(key) + dw1 = self._cond_init(ku, (self.n_spec, hsize, outsize)) db1 = jp.zeros((self.n_spec, outsize)) return (w, b), (dw0, db0, dw1, db1) @@ -374,8 +421,8 @@ def denm_init(self, key): ku, key = jr.split(key) dw0 = glorot_normal(batch_axis=0)(ku, (self.n_spec, insize, hsize)) db0 = jp.zeros((self.n_spec, hsize)) - # output is defaulted to zero. - dw1 = jp.zeros((self.n_spec, hsize, outsize)) + key, ku = jr.split(key) + dw1 = self._cond_init(ku, (self.n_spec, hsize, outsize)) db1 = jp.zeros((self.n_spec, outsize)) return (b, keep_en, (dw0, db0, dw1, db1)) @@ -407,7 +454,7 @@ def denm_f(self, params, norm_neigh_state, i, neigh_type, neigh_id, self_en): # defaults to 0 if it is not an agent. It doesn't matter since nonagents # are later filtered out. neigh_spec = jax.nn.one_hot( - self.config.etd.get_agent_specialization_idx(neigh_type), 4) + self.config.etd.get_agent_specialization_idx(neigh_type), self.n_spec) def compute_logits_f(t_state, t_spec): inputs = jp.concatenate([norm_self_state, t_state, t_spec], -1) @@ -504,8 +551,8 @@ def excl_init(self, key): ku, key = jr.split(key) dw0 = glorot_normal(batch_axis=0)(ku, (self.n_spec, insize, hsize)) db0 = jp.zeros((self.n_spec, hsize)) - # output is defaulted to zero. - dw1 = jp.zeros((self.n_spec, hsize, outsize)) + key, ku = jr.split(key) + dw1 = self._cond_init(ku, (self.n_spec, hsize, outsize)) db1 = jp.zeros((self.n_spec, outsize)) return minimal_params, (dw0, db0, dw1, db1) @@ -513,11 +560,28 @@ def excl_init(self, key): def repr_init(self, key): """Initialization for repr_f. - Simply create a default requirement of nutrients for triggering the op. + Create a default requirement of nutrients for triggering the op. Note that this value is normalized by the material_nutrient_cap. + This value is different for FLOWER and FLOWER_SEXUAL. + Also, create a sex. + A sex is inferred by doing floor(sex_p * self.sex_sensitivity). """ - return (self.config.reproduce_cost + (self.config.dissipation_per_step * 4) - + self.config.specialize_cost * 2) / self.config.material_nutrient_cap + min_repr_en = ( + self.config.reproduce_cost + (self.config.dissipation_per_step * 4) + + self.config.specialize_cost * 2) / self.config.material_nutrient_cap + # sexual reproduction combines energy of two plants. + min_repr_en_sex = min_repr_en / 2. + + # format params fn has issues with zero dimensional parameters. So as a hack + # I make this a one dimensional array. + sex_m = (jr.uniform(key, (1,)) < 0.5).astype(jp.float32) + sex_p = (0.5 + sex_m) / self.sex_sensitivity + return (min_repr_en, min_repr_en_sex, sex_p) + + def get_sex(self, params: AgentProgramType): + _, _, repr_p = self.split_params_f(params) + _,_, sex_p = self._repr_format_params_fn(repr_p) + return jp.floor(sex_p[..., 0] * self.sex_sensitivity).astype(jp.int32) def initialize(self, key): k1, k2, k3, k4, k5 = jr.split(key, 5) @@ -597,7 +661,6 @@ def par_f(self, key, perc, params): denergy_neigh = self.denm_f( denm_params, norm_neigh_state, spec_idx, neigh_type, neigh_id, self_en) - return ParallelInterface(denergy_neigh, dstate, new_spec_logit) def excl_f(self, key, perc, params): @@ -646,7 +709,8 @@ def excl_f(self, key, perc, params): # This is a probability depending on the average number of neighbors. avg_neigh_types = self.get_avg_neigh_types(neigh_type) avg_agents = avg_neigh_types[ - etd.types.AGENT_UNSPECIALIZED : etd.types.AGENT_UNSPECIALIZED + 4 + etd.types.AGENT_UNSPECIALIZED : + etd.types.AGENT_UNSPECIALIZED + self.n_spec ].sum(-1) ku, key = jax.random.split(key) rand_prob_sp = (jr.uniform(ku) < jax.nn.sigmoid( @@ -712,18 +776,26 @@ def excl_f(self, key, perc, params): def repr_f(self, key, perc, params): """Implementation of repr_f. - - A simple mask that entirely depends on how much energy we have. + + A simple mask that depends on how much energy we have, and whether the + flower is sexual or asexual. """ + min_repr_en, min_repr_en_sex, _ = self._repr_format_params_fn(params) + neigh_type, _, _ = perc + self_type = neigh_type[4] + norm_self_en = (perc.neigh_state[4, evm.EN_ST : evm.EN_ST + 2] / self.config.material_nutrient_cap) - # we could also check whether we are flowers, but it doesn't matter since - # only flowers can reproduce and it gets masked afterwards. + is_flower = self_type == self.config.etd.types.AGENT_FLOWER + is_flower_sex = self_type == self.config.etd.types.AGENT_FLOWER_SEXUAL + flower_mask_logit = is_flower * (norm_self_en > min_repr_en).all().astype( + jp.float32) + flower_sex_mask_logit = is_flower_sex * ( + norm_self_en > min_repr_en_sex).all().astype(jp.float32) # The switch checks whether the output logit is >0. - mask_logit = (norm_self_en > params).all().astype(jp.float32) - return ReproduceInterface(mask_logit) + return ReproduceInterface(flower_mask_logit + flower_sex_mask_logit) def adapt_dna_to_different_basic_logic( diff --git a/self_organising_systems/biomakerca/env_logic.py b/self_organising_systems/biomakerca/env_logic.py index 8b06aff..873149a 100644 --- a/self_organising_systems/biomakerca/env_logic.py +++ b/self_organising_systems/biomakerca/env_logic.py @@ -30,9 +30,10 @@ from collections import namedtuple from functools import partial from typing import Callable, Iterable +import copy import flax -from jax import vmap +from jax import vmap, jit import jax.numpy as jp import jax.random as jr import jax.scipy @@ -152,20 +153,20 @@ def make_empty_parallel_op_cell(config): # ReproduceOp. # flowers can create ReproduceInterfaces that are converted to ReproduceOps. -# Nutrients from the flower get converted into stored_en and we record the +# Nutrients from the flower get converted into stored_en and we record the # original position and agent id. A new seed is then tentatively created # following the logic of the environment decided upon. if "ReproduceOp" not in globals(): - ReproduceOp = namedtuple("ReproduceOp", "mask pos stored_en aid") + ReproduceOp = namedtuple("ReproduceOp", "mask pos stored_en aid is_sexual") # Helpers for making ReproduceOps. EMPTY_REPRODUCE_OP = ReproduceOp( jp.array(0.), jp.zeros([2], dtype=jp.int32), jp.zeros([2]), - jp.array(0, dtype=jp.uint32)) + jp.array(0, dtype=jp.uint32), jp.array(0, dtype=jp.int32)) ### ReproduceInterface # Interface for reproduction of agents. -# If a flower triggers reproduction, all of its energy is converted for the +# If a flower triggers reproduction, all of its energy is converted for the # new seed. However, reproduction has a cost and it may fail. # Arguments: # mask_logit: whether or not to perform reproduction. True if > 0. @@ -180,7 +181,7 @@ def make_empty_parallel_op_cell(config): # No matter the kind of cell, be it a material or an agent, they can only # perceive a limited amount of data. This is the 3x3 neighborhood of the # environment. The difference from Environment is that each cell has 9 values -# per grid. That is, neigh_type will be [w,h,9] as opposed to type_grid: [w,h]. +# per grid. That is, neigh_type will be [h,w,9] as opposed to type_grid: [h,w]. # Also note that perceived data is intended to be passed using vmap2 so that # each cell only perceives their neighbors, that is, their input is of size: # neigh_type:[9], neigh_state:[9,env_state_size], neigh_id:[9]. @@ -191,7 +192,7 @@ def make_empty_parallel_op_cell(config): def perceive_neighbors(env: Environment, etd: EnvTypeDef) -> PerceivedData: - """Return PerceivedData (gridwise, with leading axes of size [w,h]). + """Return PerceivedData (gridwise, with leading axes of size [h,w]). Cells can only perceive their neighbors. Of the neighbors, they can perceive all: type, state and agent_id. @@ -211,7 +212,7 @@ def perceive_neighbors(env: Environment, etd: EnvTypeDef) -> PerceivedData: neigh_state = jax.lax.conv_general_dilated_patches( env.state_grid[None,:], (3, 3), (1, 1), "SAME", dimension_numbers=("NHWC", "OIHW", "NHWC"))[0] - # We want to have [w,h,9,c] so that the indexing is intuitive and consistent + # We want to have [h,w,9,c] so that the indexing is intuitive and consistent # for all neigh vectors. env_state_size = env.state_grid.shape[-1] neigh_state = neigh_state.reshape( @@ -229,21 +230,69 @@ def perceive_neighbors(env: Environment, etd: EnvTypeDef) -> PerceivedData: # These functions make ExclusiveOps and execute them into the environment. def vectorize_cell_exclusive_f( - cell_f: Callable[[KeyType, PerceivedData], ExclusiveOp], w, h): + cell_f: Callable[[KeyType, PerceivedData], ExclusiveOp], h, w): """Vectorizes a cell's exclusive_f to work from 0d to 2d.""" - return lambda key, perc: vmap2(cell_f)(split_2d(key, w, h), perc) + return lambda key, perc: vmap2(cell_f)(split_2d(key, h, w), perc) def vectorize_agent_cell_f( cell_f: Callable[[KeyType, PerceivedData, AgentProgramType], - ExclusiveOp|ParallelOp], w, h): + ExclusiveOp|ParallelOp], h, w): """Vectorizes an agent cell_f to work from 0d to 2d. This works for both ExclusiveOp and ParallelOp. Note that the cell_f *requires* to have the proper argument name 'programs', so this is an informal interface. """ return lambda key, perc, progs: vmap2(partial(cell_f, programs=progs))( - split_2d(key, w, h), perc) + split_2d(key, h, w), perc) + + +def compute_sparse_agent_cell_f( + key: KeyType, + cell_f: (Callable[[KeyType, PerceivedData, AgentProgramType], + ExclusiveOp|ParallelOp] | + Callable[[ + KeyType, PerceivedData, CellPositionType, AgentProgramType], + ReproduceOp]), + perc: PerceivedData, + env_type_grid, programs: AgentProgramType, etd: EnvTypeDef, + n_sparse_max: int, + b_pos=None): + """Compute a sparse agent cell_f. + + This works for ExclusiveOp, ParallelOp and ReproduceOp. + Note that the cell_f *requires* to have the proper argument name 'programs', + so this is an informal interface. + + if it is used for a ReproduceOp, b_pos needs to be set, being a flat list of + (y,x) positions. + """ + # get the args of alive cells. + # note that cell_f will not work by itself if the cell is not an agent. + is_agent_flat = etd.is_agent_fn(env_type_grid.flatten()) + # leftmost are not agents, rightmost are agents. + sorted_agent_idx = jp.argsort(is_agent_flat) + sparse_idx = sorted_agent_idx[-n_sparse_max:] + + # compute, sparsely, cell_f + v_part_cell_f = vmap(partial(cell_f, programs=programs)) + v_keys = jr.split(key, n_sparse_max) + sparse_perc_flat = jax.tree_util.tree_map( + lambda x: x.reshape((-1,)+ x.shape[2:])[sparse_idx], perc) + if b_pos is None: + sparse_output_tree = v_part_cell_f(v_keys, sparse_perc_flat) + else: + # this is a reproduce op, and therefore requires pos too. + sparse_output_tree = v_part_cell_f( + v_keys, sparse_perc_flat, b_pos[sparse_idx]) + + # scatter the result. + # we also need to reshape so that it is [h, w, ...] shape. + return jax.tree_util.tree_map( + lambda x: jax.ops.segment_sum( + x, sparse_idx, num_segments=is_agent_flat.shape[0]).reshape( + env_type_grid.shape + x.shape[1:]), + sparse_output_tree) def make_material_exclusive_interface( @@ -396,33 +445,43 @@ def execute_and_aggregate_exclusive_ops( excl_fs: Iterable[tuple[EnvTypeType, Callable[[KeyType, PerceivedData], ExclusiveOp]]], agent_excl_f: Callable[[ - KeyType, PerceivedData, AgentProgramType], ExclusiveInterface] + KeyType, PerceivedData, AgentProgramType], ExclusiveInterface], + n_sparse_max: int|None = None ) -> ExclusiveOp: """Execute all exclusive functions and aggregate them all into a single ExclusiveOp for each cell. - + This function constructs sanitized interfaces for the input excl_fs and agent_excl_f, making sure that no laws of physics are broken. - It also then executes these operation for each cell in the grid and then - aggregates the resulting ExclusiveOp for each cell. + if n_sparse_max is None, it then executes these operation for each cell in the + grid. If n_sparse_max is an integer, it instead performs a sparse computation + for agent operations, only for alive agent cells. The computation is capped, + so some agents may be skipped if the cap is too low. The agents being chosen + in that case are NOT ensured to be random. Consider changing the code to allow + for a random subset of cells to be run if you want that behavior. + It then aggregates the resulting ExclusiveOp for each cell. Aggregation can be done because *only one function*, at most, will be allowed to output nonzero values for each cell. """ etd = config.etd perc = perceive_neighbors(env, etd) - w, h = env.type_grid.shape + h, w = env.type_grid.shape v_excl_fs = [vectorize_cell_exclusive_f( - make_material_exclusive_interface(t, f, config), w, h) for (t, f) + make_material_exclusive_interface(t, f, config), h, w) for (t, f) in excl_fs] - v_agent_excl_f = vectorize_agent_cell_f( - make_agent_exclusive_interface(agent_excl_f, config), w, h) - - k1, k2, key = jr.split(key, 3) - w, h = env.type_grid.shape - excl_ops = [f(k, perc) for k, f in - zip(jr.split(k1, len(v_excl_fs)), v_excl_fs)] + [ - v_agent_excl_f(k2, perc, programs) - ] + + agent_excl_intf_f = make_agent_exclusive_interface(agent_excl_f, config) + k1, key = jr.split(key) + if n_sparse_max is None: + v_agent_excl_f = vectorize_agent_cell_f(agent_excl_intf_f, h, w) + agent_excl_op = v_agent_excl_f(k1, perc, programs) + else: + agent_excl_op = compute_sparse_agent_cell_f( + k1, agent_excl_intf_f, perc, env.type_grid, programs, etd, n_sparse_max) + + k1, key = jr.split(key) + excl_ops = [f(k, perc) for k, f in + zip(jr.split(k1, len(v_excl_fs)), v_excl_fs)] + [agent_excl_op] excl_op = tree_map_sum_ops(excl_ops) return excl_op @@ -446,7 +505,7 @@ def env_exclusive_decision( key: KeyType, env: Environment, excl_op: ExclusiveOp): """Choose up to one excl_op for each cell and execute them, updating the env. """ - w, h, chn = env.state_grid.shape + h, w, chn = env.state_grid.shape def extract_patches_that_target_cell(x): """Extract all ExclusiveOps subarray of neighbors that target each cell. @@ -456,7 +515,7 @@ def extract_patches_that_target_cell(x): ExclusiveOps that target a given cell, and return them, so that one of them can be chosen to be executed. """ - # input is either (w,h,9) or (w,h,9,c) + # input is either (h,w,9) or (h,w,9,c) ndim = x.ndim x_shape = x.shape old_dtype = x.dtype @@ -468,7 +527,7 @@ def extract_patches_that_target_cell(x): neigh_state = jax.lax.conv_general_dilated_patches( x[None,:], (3, 3), (1, 1), "SAME", dimension_numbers=("NHWC", "OIHW", "NHWC"))[0] - # make it (w, h, k, 9) + # make it (h, w, k, 9) # where k is either 9 or 9*c neigh_state = neigh_state.reshape( neigh_state.shape[:-1] + (x.shape[-1], 9)) @@ -477,7 +536,7 @@ def extract_patches_that_target_cell(x): neigh_state = neigh_state.reshape( neigh_state.shape[:-2] + x_shape[-2:] + (9,)).transpose( (0, 1, 2, 4, 3)) - # now it's either (w,h,9,9) or (w,h,9,9,c). + # now it's either (h,w,9,9) or (h,w,9,9,c). # The leftmost '9' refers to the position of the cell's neighbors. # the rightmost '9' indicates the exclusive op slice of the neighbor, which # in turn targets 9 neighbors. @@ -498,7 +557,7 @@ def extract_patches_that_target_cell(x): # choose a random one from them. rnd_neigh_idx = vmap2(_cell_choose_random_action)( - split_2d(key, w, h), excl_op_neighs) + split_2d(key, h, w), excl_op_neighs) # Do so by zeroing out everything except the chosen update. action_mask = jax.nn.one_hot(rnd_neigh_idx, 9) @@ -531,12 +590,12 @@ def extract_patches_that_target_cell(x): ) def map_to_actor(x): pad_x = jp.pad(x, ((1, 1), (1, 1), (0, 0))) - return jp.stack([jax.lax.dynamic_slice(pad_x, s, (w, h, 9))[:,:,n] for n, s + return jp.stack([jax.lax.dynamic_slice(pad_x, s, (h, w, 9))[:,:,n] for n, s in enumerate(MAP_TO_ACTOR_SLICING)], 2) def map_to_actor_e(x): pad_x = jp.pad(x, ((1, 1), (1, 1), (0, 0), (0, 0))) - return jp.stack([jax.lax.dynamic_slice(pad_x, s+(0,), (w, h, 9, chn))[:,:,n] + return jp.stack([jax.lax.dynamic_slice(pad_x, s+(0,), (h, w, 9, chn))[:,:,n] for n, s in enumerate(MAP_TO_ACTOR_SLICING)], 2) a_upd_mask = map_to_actor(a_upd_mask_from_t).sum(-1) @@ -568,28 +627,33 @@ def env_perform_exclusive_update( excl_fs: Iterable[tuple[EnvTypeType, Callable[[KeyType, PerceivedData], ExclusiveOp]]], agent_excl_f: Callable[[ - KeyType, PerceivedData, AgentProgramType], ExclusiveInterface] + KeyType, PerceivedData, AgentProgramType], ExclusiveInterface], + n_sparse_max: int|None = None ) -> Environment: """Perform exclusive operations in the environment. - + This is the function that should be used for high level step_env design. - + Arguments: key: a jax random number generator. env: the input environment to be modified. programs: params of agents that govern their agent_excl_f. They should be one line for each agent_id allowed in the environment. config: EnvConfig describing the physics of the environment. - excl_fs: list of pairs of (env_type, func) where env_type is the type of + excl_fs: list of pairs of (env_type, func) where env_type is the type of material that triggers the exclusive func. - agent_excl_f: The exclusive function that agents perform. It takes as + agent_excl_f: The exclusive function that agents perform. It takes as input a exclusive program and outputs a ExclusiveInterface. + n_sparse_max: if None, agent_excl_f are performed for the entire grid. If it + is an integer, instead, we perform a sparse computation masked by actual + agent cells. Note that this is capped and if more agents are alive, an + undefined subset of agent cells will be run. Returns: an updated environment. """ k1, key = jr.split(key) excl_op = execute_and_aggregate_exclusive_ops( - k1, env, programs, config, excl_fs, agent_excl_f) + k1, env, programs, config, excl_fs, agent_excl_f, n_sparse_max) key, key1 = jr.split(key) env = env_exclusive_decision(key1, env, excl_op) @@ -685,7 +749,8 @@ def env_perform_parallel_update( key: KeyType, env: Environment, programs: AgentProgramType, config: EnvConfig, par_f: Callable[[ - KeyType, PerceivedData, AgentProgramType], ParallelInterface] + KeyType, PerceivedData, AgentProgramType], ParallelInterface], + n_sparse_max: int|None = None ) -> Environment: """Perform parallel operations in the environment. @@ -699,23 +764,29 @@ def env_perform_parallel_update( config: EnvConfig describing the physics of the environment. par_f: The parallel function that agents perform. It takes as input a parallel program and outputs a ParallelInterface. + n_sparse_max: either an int or None. If set to int, we will use a budget for + the amounts of agent operations allowed at each step. Returns: an updated environment. """ # First compute the ParallelOp for each cell. - w, h = env.type_grid.shape - v_par_f = vectorize_agent_cell_f(make_agent_parallel_interface( - par_f, config), w, h) + h, w = env.type_grid.shape etd = config.etd perc = perceive_neighbors(env, etd) + par_interface_f = make_agent_parallel_interface( par_f, config) k1, key = jr.split(key) - par_op = v_par_f(k1, perc, programs) + if n_sparse_max is None: + v_par_interface_f = vectorize_agent_cell_f(par_interface_f, h, w) + par_op = v_par_interface_f(k1, perc, programs) + else: + par_op = compute_sparse_agent_cell_f( + k1, par_interface_f, perc, env.type_grid, programs, etd, n_sparse_max) # Then process them. mask, denergy_neigh, dstate, new_type = par_op - w, h = env.type_grid.shape + h, w = env.type_grid.shape MAP_TO_NEIGH_SLICING = ( (2, 2, 0), # indexing (-1, -1) @@ -730,7 +801,7 @@ def env_perform_parallel_update( ) def map_to_neigh(x): pad_x = jp.pad(x, ((1, 1), (1, 1), (0, 0), (0, 0))) - return jp.stack([jax.lax.dynamic_slice(pad_x, s+(0,), (w, h, 9, 2))[:, :, n] + return jp.stack([jax.lax.dynamic_slice(pad_x, s+(0,), (h, w, 9, 2))[:, :, n] for n, s in enumerate(MAP_TO_NEIGH_SLICING)], 2) denergy = map_to_neigh(denergy_neigh).sum(2) @@ -757,9 +828,9 @@ def map_to_neigh(x): # Functions for making and processing ReproduceOps. -def vectorize_reproduce_f(repr_f, w, h): +def vectorize_reproduce_f(repr_f, h, w): return lambda key, perc, pos, progs: vmap2(partial(repr_f, programs=progs))( - split_2d(key, w, h), perc, pos) + split_2d(key, h, w), perc, pos) def _convert_to_reproduce_op( @@ -781,20 +852,27 @@ def _convert_to_reproduce_op( # must want to reproduce want_to_repr_m = (mask_logit > 0.0).astype(jp.float32) - # must be a flower - is_flower_m = (self_type == env_config.etd.types.AGENT_FLOWER).astype( - jp.float32) + # must be a flower. + # depending on the kind, flag it as either sexual or asexual reproduction. + is_flower_m = self_type == env_config.etd.types.AGENT_FLOWER + is_flower_sexual_m = self_type == env_config.etd.types.AGENT_FLOWER_SEXUAL + + is_flower_type = is_flower_m | is_flower_sexual_m # must have enough energy. + min_repr_cost = env_config.reproduce_cost * ( + 1 *is_flower_m + 0.5 * is_flower_sexual_m) has_enough_en_m = ( - (self_en >= env_config.reproduce_cost).all().astype(jp.float32) + (self_en >= min_repr_cost).all().astype(jp.float32) ) - mask = want_to_repr_m * is_flower_m * has_enough_en_m + mask = want_to_repr_m * is_flower_type * has_enough_en_m stored_en = (self_en - env_config.reproduce_cost) * mask aid = (self_id * mask).astype(jp.uint32) - return ReproduceOp(mask, pos, stored_en, aid) + # don't use booleans. + return ReproduceOp(mask, pos, stored_en, aid, + is_flower_sexual_m.astype(jp.int32)) def make_agent_reproduce_interface( @@ -816,8 +894,11 @@ def make_agent_reproduce_interface( # Note that we need to pass the extra information of the position of the cell. # this is used in the conversion to ReproduceOp, not by the cell. def f(key, perc, pos, programs): - # it has to be a flower! - is_correct_type = config.etd.types.AGENT_FLOWER == perc.neigh_type[4] + # it has to be a flower type (sexual or asexual) + is_correct_type = ( + perc.neigh_type[4] == jp.array( + [config.etd.types.AGENT_FLOWER, + config.etd.types.AGENT_FLOWER_SEXUAL], dtype=jp.uint32)).any(-1) curr_agent_id = perc.neigh_id[4] program = programs[curr_agent_id] @@ -874,23 +955,23 @@ def _select_random_position_for_seed_within_range( return chosen_idx, n_available > 0 -def env_perform_one_reproduce_op( - key: KeyType, env: Environment, repr_op: ReproduceOp, config: EnvConfig): - """Perform one single ReproduceOp. +def env_try_place_one_seed( + key: KeyType, env: Environment, op_info, config: EnvConfig): + """Try to place one seed in the environment. - For a ReproduceOp to be successful, fertile soil in the neighborhood must be + For this op to be successful, fertile soil in the neighborhood must be found. If it is, a new seed (two unspecialized cells) are placed in the environment. Their age is reset to zero, and they may have a different agent_id than their parent, if mutation was set to true. """ - mask, pos, stored_en, aid = repr_op + mask, pos, stored_en, aid = op_info etd = config.etd def true_fn(env): best_idx_per_column, column_m = find_fertile_soil(env.type_grid, etd) t_column, column_valid = _select_random_position_for_seed_within_range( - key, pos[1], config.reproduce_min_dist, config.reproduce_max_dist, + key, pos[1], config.reproduce_min_dist, config.reproduce_max_dist, column_m) def true_fn2(env): @@ -904,28 +985,34 @@ def true_fn2(env): return jax.lax.cond(mask, true_fn, lambda env: env, env) -def env_reproduce_operations(key, env, b_repr_op, config): - """Perform a batch of ReproduceOps. +def env_try_place_seeds(key, env, b_op_info, config): + """Try to place seeds in the environment. These are performed sequentially. Note that some ops may be masked and therefore be noops. """ - def body_f(carry, repr_op): + def body_f(carry, op_info): env, key = carry key, ku = jr.split(key) - env = env_perform_one_reproduce_op(ku, env, repr_op, config) + env = env_try_place_one_seed(ku, env, op_info, config) return (env, key), 0 - (env, key), _ = jax.lax.scan(body_f, (env, key), b_repr_op) + (env, key), _ = jax.lax.scan(body_f, (env, key), b_op_info) return env -def _select_subset_of_reproduce_ops(key, b_repr_op, neigh_type, config): +def _select_subset_of_reproduce_ops( + key, b_repr_op, neigh_type, config, select_sexual_repr): # Only a small subset of possible ReproduceOps are selected at each step. - # This is config dependent (config.n_reproduce_per_step). b_pos = b_repr_op.pos mask_flat = b_repr_op.mask.flatten() b_pos_flat = b_pos.reshape((-1, 2)) + n_repr_ops = (config.n_sexual_reproduce_per_step * 2 if select_sexual_repr + else config.n_reproduce_per_step) + + # Because sexual and asexual reproductions are fundamentally different, use + # select_sexual_repr to decide which kinds of reproduce_op to extract. + mask_flat *= (select_sexual_repr == b_repr_op.is_sexual).flatten() p_logits = mask_flat # Moreover, flowers can only reproduce if in contact with air. The more air, @@ -939,7 +1026,7 @@ def _select_subset_of_reproduce_ops(key, b_repr_op, neigh_type, config): k1, key = jr.split(key) selected_pos = jr.choice( k1, b_pos_flat, p=p_logits/p_logits.sum().clip(1), - shape=(config.n_reproduce_per_step,), replace=False) + shape=(n_repr_ops,), replace=False) return selected_pos @@ -952,11 +1039,21 @@ def env_perform_reproduce_update( mutate_programs=False, programs: (AgentProgramType | None) = None, mutate_f: (Callable[[KeyType, AgentProgramType], AgentProgramType] | None - ) = None): + ) = None, + enable_asexual_reproduction=True, + enable_sexual_reproduction=True, + does_sex_matter=True, + sexual_mutate_f: (Callable[[KeyType, AgentProgramType, AgentProgramType], + AgentProgramType] | None + ) = None, + split_mutator_params_f = None, + get_sex_f: (Callable[[AgentProgramType], AgentProgramType] | None) = None, + n_sparse_max: int|None = None, + return_metrics=False): """Perform reproduce operations in the environment. - + This is the function that should be used for high level step_env design. - + Arguments: key: a jax random number generator. env: the input environment to be modified. @@ -974,37 +1071,102 @@ def env_perform_reproduce_update( if mutate_programs is set to True. mutate_f: The mutation function to mutate 'programs'. This is intended to be the method 'mutate' of a Mutator class. + enable_asexual_reproduction: if set to True, asexual reproduction is + enabled. At least enable_asexual_reproduction or + enable_sexual_reproduction must be set to True. + enable_sexual_reproduction: if set to True, sexual reproduction is enabled. + does_sex_matter: if set to True, and enable_sexual_reproduction is True, + then sexual reproduction can only occur between different sex entities. + sexual_mutate_f: The sexual mutation function to mutate 'programs'. This is + intended to be the method 'mutate' of a SexualMutator class. + split_mutator_params_f: a mutator function that explains how to split the + parameters of the mutator. + get_sex_f: The function that extracts the sex of the agent from its params. + n_sparse_max: either an int or None. If set to int, we will use a budget for + the amounts of agent operations allowed at each step. + return_metrics: if True, return metrics about whether reproduction occurred, + and who are the parents and children. Returns: an updated environment. if mutate_programs is True, it also returns the updated programs. """ + assert enable_asexual_reproduction or enable_sexual_reproduction etd = config.etd perc = perceive_neighbors(env, etd) - k1, key = jr.split(key) - w, h = env.type_grid.shape + h, w = env.type_grid.shape - v_repr_f = vectorize_reproduce_f( - make_agent_reproduce_interface(repr_f, config), w, h) - - b_pos = jp.stack(jp.meshgrid(jp.arange(w), jp.arange(h), indexing="ij"), -1) - b_repr_op = v_repr_f(k1, perc, b_pos, repr_programs) + b_pos = jp.stack(jp.meshgrid(jp.arange(h), jp.arange(w), indexing="ij"), -1) + repr_interface_f = make_agent_reproduce_interface(repr_f, config) + k1, key = jr.split(key) + if n_sparse_max is None: + v_repr_f = vectorize_reproduce_f(repr_interface_f, h, w) + b_repr_op = v_repr_f(k1, perc, b_pos, repr_programs) + else: + b_repr_op = compute_sparse_agent_cell_f( + k1, repr_interface_f, perc, env.type_grid, repr_programs, etd, + n_sparse_max, b_pos.reshape((h*w, 2))) # Only a small subset of possible ReproduceOps are selected at each step. - # This is config dependent (config.n_reproduce_per_step). - selected_pos = _select_subset_of_reproduce_ops( - k1, b_repr_op, perc.neigh_type, config) - spx, spy = selected_pos[:, 0], selected_pos[:, 1] - - selected_mask = b_repr_op.mask[spx, spy] - selected_stored_en = b_repr_op.stored_en[spx, spy] - selected_aid = b_repr_op.aid[spx, spy] - - if mutate_programs: + # sexual and asexual reproductions are treated independently. + if enable_asexual_reproduction: + ## First, asexual reproduction. + # This is config dependent (config.n_reproduce_per_step). + k1, key = jr.split(key) + selected_pos = _select_subset_of_reproduce_ops( + k1, b_repr_op, perc.neigh_type, config, 0) + spx, spy = selected_pos[:, 0], selected_pos[:, 1] + + selected_mask = b_repr_op.mask[spx, spy] + selected_stored_en = b_repr_op.stored_en[spx, spy] + selected_aid = b_repr_op.aid[spx, spy] + + if enable_sexual_reproduction: + ## Sexual reproduction + # This is config dependent (2 * config.n_sexual_reproduce_per_step). + # twice, because we need 2 flowers per sexual operation. + k1, key = jr.split(key) + selected_pos_sx = _select_subset_of_reproduce_ops( + k1, b_repr_op, perc.neigh_type, config, 1) + spx_sx, spy_sx = selected_pos_sx[:, 0], selected_pos_sx[:, 1] + selected_aid_sx = b_repr_op.aid[spx_sx, spy_sx] + selected_aid_sx_1 = selected_aid_sx[::2] + selected_aid_sx_2 = selected_aid_sx[1::2] + + selected_mask_sx = b_repr_op.mask[spx_sx, spy_sx] + # to reproduce, selected_mask_sx of both parent cells to be True. Representing + # whether there actually are sexual flowers that triggered a reproduce op. + pair_repr_mask_sx = selected_mask_sx[::2] * selected_mask_sx[1::2] + if does_sex_matter: + # moreover, their sex has to be different. Done to prevent selfing. + # we could consider moving this to the sexual mutator as a responsibility. + selected_programs_sx = programs[selected_aid_sx] + logic_params, _ = vmap(split_mutator_params_f)(selected_programs_sx) + sex_of_selected = vmap(get_sex_f)(logic_params) + pair_repr_mask_sx = pair_repr_mask_sx * ( + sex_of_selected[::2] != sex_of_selected[1::2]) + + selected_stored_en_sx = b_repr_op.stored_en[spx_sx, spy_sx] + # the total energy is the sum of the two flowers. + pair_stored_en_sx = selected_stored_en_sx[::2] + selected_stored_en_sx[1::2] + + if not mutate_programs: + if enable_asexual_reproduction: + repr_aid = selected_aid + # with sexual reproduction, we assume that programs are mutated. + # but, as a failsafe, we randomly pick a parent's ID otherwise. + if enable_sexual_reproduction: + k1, key = jr.split(key) + parent_m = ( + jr.uniform(k1, [config.n_sexual_reproduce_per_step]) < 0.5).astype( + jp.uint32)[..., None] + repr_aid_sx = (selected_aid_sx_1 * parent_m + + selected_aid_sx_2 * (1 - parent_m)) + else: # Logic: # look into the pool of programs and see if some of them are not used ( # there are no agents alive with such program). - # if that is the case, create a new program with mutate_f, then modify the - # corresponding 'selected_aid'. + # if that is the case, create a new program with mutate_f or + # sexual_mutate_f, then modify the corresponding 'selected_aid'. # If there is no space, set the mask to zero instead. # get n agents per id. @@ -1015,51 +1177,136 @@ def env_perform_reproduce_update( is_agent_flat, env_aid_flat, num_segments=programs.shape[0]) sorted_na_idx = jp.argsort(n_agents_in_env).astype(jp.uint32) - # we only care about the first few indexes - sorted_na_chosen_idx = sorted_na_idx[:config.n_reproduce_per_step] - sorted_na_chosen_mask = (n_agents_in_env[sorted_na_chosen_idx] == 0 - ).astype(jp.float32) - - # assume that the number of selected reproductions is LESS than the total - # number of programs. - to_mutate_programs = programs[selected_aid] - mutated_programs = vmap(mutate_f)( - jr.split(key, config.n_reproduce_per_step), to_mutate_programs) - - mutation_mask = (selected_mask * sorted_na_chosen_mask) - mutation_mask_e = mutation_mask[:, None] - n_mutation_mask_e = 1. - mutation_mask_e - - # substitute the programs - programs = programs.at[sorted_na_chosen_idx].set( - mutation_mask_e * mutated_programs - + n_mutation_mask_e * programs[sorted_na_chosen_idx]) - # update the aid and the mask - selected_mask = mutation_mask - selected_aid = sorted_na_chosen_idx - - # these positions (if mask says yes) are then selected to reproduce. - # A seed is spawned if possible. - env = env_reproduce_operations( - key, env, - ReproduceOp(selected_mask, selected_pos, selected_stored_en, - selected_aid), - config) - # The flower is destroyed, regardless of whether the operation succeeds. - n_selected_mask = 1 - selected_mask - n_selected_mask_uint = n_selected_mask.astype(jp.uint32) - env = Environment( - env.type_grid.at[spx, spy].set( - n_selected_mask_uint * env.type_grid[spx, spy]), # 0 is VOID - env.state_grid.at[spx, spy].set( - n_selected_mask[..., None] * env.state_grid[spx, spy] - ), # set everything to zero - env.agent_id_grid.at[spx, spy].set( - n_selected_mask_uint * env.agent_id_grid[spx, spy]) # default id 0. - ) - if mutate_programs: - return env, programs - return env + + if enable_asexual_reproduction: + ## Asexual reproduction + # we only care about the first few indexes + sorted_na_chosen_idx = sorted_na_idx[:config.n_reproduce_per_step] + sorted_na_chosen_mask = (n_agents_in_env[sorted_na_chosen_idx] == 0 + ).astype(jp.float32) + + # assume that the number of selected reproductions is LESS than the total + # number of programs. + to_mutate_programs = programs[selected_aid] + ku, key = jr.split(key) + mutated_programs = vmap(mutate_f)( + jr.split(ku, config.n_reproduce_per_step), to_mutate_programs) + + mutation_mask = (selected_mask * sorted_na_chosen_mask) + mutation_mask_e = mutation_mask[:, None] + n_mutation_mask_e = 1. - mutation_mask_e + + # substitute the programs + programs = programs.at[sorted_na_chosen_idx].set( + mutation_mask_e * mutated_programs + + n_mutation_mask_e * programs[sorted_na_chosen_idx]) + # update the aid and the mask + selected_mask = mutation_mask + repr_aid = sorted_na_chosen_idx + + if enable_sexual_reproduction: + ## Sexual reproduction + # we only care about some few indexes + sorted_na_chosen_idx_sx = sorted_na_idx[ + config.n_reproduce_per_step: + config.n_reproduce_per_step+config.n_sexual_reproduce_per_step] + sorted_na_chosen_mask_sx = ( + n_agents_in_env[sorted_na_chosen_idx_sx] == 0).astype(jp.float32) + + # assume that the number of selected reproductions is LESS than the total + # number of programs. + to_mutate_programs_sx = programs[selected_aid_sx] + ku, key = jr.split(key) + mutated_programs_sx = vmap(sexual_mutate_f)( + jr.split(ku, config.n_sexual_reproduce_per_step), + to_mutate_programs_sx[::2], to_mutate_programs_sx[1::2]) + # To assess whether we can perform a sexual reproduction, we need, for + # each chosen position: + # 1. sorted_na_chosen_mask_sx to be True, representing whether there is + # space for new programs. + # 2. pair_repr_mask_sx to be True. + pair_repr_mask_sx = (pair_repr_mask_sx * sorted_na_chosen_mask_sx) + + pair_repr_mask_sx_e = pair_repr_mask_sx[:, None] + n_pair_repr_mask_sx_e = 1. - pair_repr_mask_sx_e + + # substitute the programs + programs = programs.at[sorted_na_chosen_idx_sx].set( + pair_repr_mask_sx_e * mutated_programs_sx + + n_pair_repr_mask_sx_e * programs[sorted_na_chosen_idx_sx]) + + # save the new ids for the children. + repr_aid_sx = sorted_na_chosen_idx_sx + + if return_metrics: + metrics = {} + + if enable_asexual_reproduction: + ## Asexual reproduction + # these positions (if mask says yes) are then selected to reproduce. + # A seed is spawned if possible. + k1, key = jr.split(key) + env = env_try_place_seeds( + k1, env, + (selected_mask, selected_pos, selected_stored_en, repr_aid), + config) + # The flower is destroyed, regardless of whether the operation succeeds. + n_selected_mask = 1 - selected_mask + n_selected_mask_uint = n_selected_mask.astype(jp.uint32) + env = Environment( + env.type_grid.at[spx, spy].set( + n_selected_mask_uint * env.type_grid[spx, spy]), # 0 is VOID + env.state_grid.at[spx, spy].set( + n_selected_mask[..., None] * env.state_grid[spx, spy] + ), # set everything to zero + env.agent_id_grid.at[spx, spy].set( + n_selected_mask_uint * env.agent_id_grid[spx, spy]) # default id 0. + ) + + if return_metrics: + metrics["asexual_reproduction"] = (selected_mask, selected_aid, repr_aid) + + if enable_sexual_reproduction: + ## Sexual reproduction + # seeds will spawn in a neighborhood of one of the two parents, chosen at + # random. + k1, key = jr.split(key) + pos_m = (jr.uniform(k1, [config.n_sexual_reproduce_per_step]) < 0.5).astype( + jp.int32)[..., None] + repr_pos_sx = (selected_pos_sx[::2] * pos_m + + selected_pos_sx[1::2] * (1 - pos_m)) + + k1, key = jr.split(key) + env = env_try_place_seeds( + k1, env, + (pair_repr_mask_sx, repr_pos_sx, pair_stored_en_sx, repr_aid_sx), + config) + # The flower is destroyed, regardless of whether the operation succeeds. + # update the selected_mask_sx, since some flowers may not have been actually + # selected. + # mutation_mask_sx is the outcome of a pair, so you need to repeat it twice. + destroy_mask_sx = jp.repeat(pair_repr_mask_sx, 2, axis=-1) + n_destroy_mask_sx = 1 - destroy_mask_sx + n_destroy_mask_sx_uint = n_destroy_mask_sx.astype(jp.uint32) + env = Environment( + env.type_grid.at[spx_sx, spy_sx].set( + n_destroy_mask_sx_uint * env.type_grid[spx_sx, spy_sx]), # 0 is VOID + env.state_grid.at[spx_sx, spy_sx].set( + n_destroy_mask_sx[..., None] * env.state_grid[spx_sx, spy_sx] + ), # set everything to zero + env.agent_id_grid.at[spx_sx, spy_sx].set( # default id is 0 + n_destroy_mask_sx_uint * env.agent_id_grid[spx_sx, spy_sx]) + ) + + if return_metrics: + metrics["sexual_reproduction"] = ( + pair_repr_mask_sx, selected_aid_sx_1, selected_aid_sx_2, + repr_aid_sx) + + result = (env, programs) if mutate_programs else env + if return_metrics: + return result, metrics + return result def intercept_reproduce_ops( @@ -1077,6 +1324,9 @@ def intercept_reproduce_ops( env_perform_reproduce_update). I chose to do that to avoid making the latter too complex. But I might refactor it eventually. + TODO: this was not modified for sexual reproduction! Double check it can be + used anyway. + Arguments: key: a jax random number generator. env: the input environment to be modified. @@ -1096,17 +1346,17 @@ def intercept_reproduce_ops( etd = config.etd perc = perceive_neighbors(env, etd) k1, key = jr.split(key) - w, h = env.type_grid.shape + h, w = env.type_grid.shape v_repr_f = vectorize_reproduce_f( - make_agent_reproduce_interface(repr_f, config), w, h) + make_agent_reproduce_interface(repr_f, config), h, w) - b_pos = jp.stack(jp.meshgrid(jp.arange(w), jp.arange(h), indexing="ij"), -1) + b_pos = jp.stack(jp.meshgrid(jp.arange(h), jp.arange(w), indexing="ij"), -1) b_repr_op = v_repr_f(k1, perc, b_pos, repr_programs) # Only a small subset of possible ReproduceOps are selected at each step. # This is config dependent (config.n_reproduce_per_step). selected_pos = _select_subset_of_reproduce_ops( - k1, b_repr_op, perc.neigh_type, config) + k1, b_repr_op, perc.neigh_type, config, False) spx, spy = selected_pos[:, 0], selected_pos[:, 1] selected_mask = b_repr_op.mask[spx, spy] @@ -1136,7 +1386,7 @@ def intercept_reproduce_ops( ### Gravity logic. -def _line_gravity(env, x, h, etd): +def _line_gravity(env, x, w, etd): type_grid, state_grid, agent_id_grid = env env_state_size = state_grid.shape[-1] # self needs to be affected by gravity: @@ -1152,25 +1402,25 @@ def _line_gravity(env, x, h, etd): vmap(lambda ctype: (ctype != etd.structural_mats).all())(type_grid[x])) swap_mask = (is_gravity_mat & is_down_intangible_mat & is_crumbling ).astype(jp.float32) - # [h] -> [1,h] + # [w] -> [1,w] swap_mask_e = swap_mask[None,:] swap_mask_uint_e = swap_mask_e.astype(jp.uint32) - idx_swap_x = jp.repeat(jp.array([x, x+1]), h) - idx_swap_y = jp.concatenate([jp.arange(0, h, dtype=jp.int32), - jp.arange(0, h, dtype=jp.int32)], 0) - type_slice = jax.lax.dynamic_slice(type_grid, (x, 0), (2, h)) + idx_swap_x = jp.repeat(jp.array([x, x+1]), w) + idx_swap_y = jp.concatenate([jp.arange(0, w, dtype=jp.int32), + jp.arange(0, w, dtype=jp.int32)], 0) + type_slice = jax.lax.dynamic_slice(type_grid, (x, 0), (2, w)) type_upd = (type_slice * (1 - swap_mask_uint_e) + type_slice[::-1] * swap_mask_uint_e).reshape(-1) new_type_grid = type_grid.at[idx_swap_x, idx_swap_y].set(type_upd) swap_mask_ee = swap_mask_e[..., None] state_slice = jax.lax.dynamic_slice( - state_grid, (x, 0, 0), (2, h, env_state_size)) + state_grid, (x, 0, 0), (2, w, env_state_size)) state_upd = (state_slice * (1. - swap_mask_ee) + state_slice[::-1] * swap_mask_ee).reshape([-1, env_state_size]) new_state_grid = state_grid.at[idx_swap_x, idx_swap_y].set(state_upd) # agent ids - id_slice = jax.lax.dynamic_slice(agent_id_grid, (x, 0), (2, h)) + id_slice = jax.lax.dynamic_slice(agent_id_grid, (x, 0), (2, w)) id_upd = (id_slice * (1 - swap_mask_uint_e) + id_slice[::-1] * swap_mask_uint_e).reshape(-1) new_agent_id_grid = agent_id_grid.at[idx_swap_x, idx_swap_y].set(id_upd) @@ -1188,11 +1438,11 @@ def env_process_gravity(env: Environment, etd: EnvTypeDef) -> Environment: Create a new env by applying gravity on every line, from bottom to top. Nit: right now, you can't fall off, so we start from the second to bottom. """ - w, h = env.type_grid.shape + h, w = env.type_grid.shape env, _ = jax.lax.scan( - partial(_line_gravity, h=h, etd=etd), + partial(_line_gravity, w=w, etd=etd), env, - jp.arange(w-2, -1, -1)) + jp.arange(h-2, -1, -1)) return env @@ -1349,7 +1599,7 @@ def process_energy(env: Environment, config: EnvConfig) -> Environment: neigh_asking_nutrients = jax.lax.conv_general_dilated_patches( asking_nutrients[None,:], (3, 3), (1, 1), "SAME", dimension_numbers=("NHWC", "OIHW", "NHWC"))[0] - # we want to have [w,h,9,c] so that the indexing is intuitive and consistent + # we want to have [h,w,9,c] so that the indexing is intuitive and consistent # for both neigh vectors. neigh_asking_nutrients = neigh_asking_nutrients.reshape( neigh_asking_nutrients.shape[:2] + (2, 9)).transpose((0, 1, 3, 2)) @@ -1363,7 +1613,7 @@ def process_energy(env: Environment, config: EnvConfig) -> Environment: absorb_perc = jax.lax.conv_general_dilated_patches( perc_to_give[None,:], (3, 3), (1, 1), "SAME", dimension_numbers=("NHWC", "OIHW", "NHWC"))[0] - # we want to have [w,h,9,c] so that the indexing is intuitive and consistent + # we want to have [h,w,9,c] so that the indexing is intuitive and consistent # for both neigh vectors. absorb_perc = absorb_perc.reshape( absorb_perc.shape[:2] + (2, 9)).transpose((0, 1, 3, 2)) @@ -1376,7 +1626,7 @@ def process_energy(env: Environment, config: EnvConfig) -> Environment: is_agent_grid = etd.is_agent_fn(env.type_grid) is_agent_grid_e_f = is_agent_grid.astype(jp.float32)[..., None] ag_spec_idx = etd.get_agent_specialization_idx(env.type_grid) - # [w,h,3] @ [3,2] -> [w,h,2] + # [h,w,3] @ [3,2] -> [h,w,2] agent_dissipation_rate = etd.dissipation_rate_per_spec[ag_spec_idx] dissipated_energy = (config.dissipation_per_step * agent_dissipation_rate * is_agent_grid_e_f) diff --git a/self_organising_systems/biomakerca/environments.py b/self_organising_systems/biomakerca/environments.py index db082b7..17852d5 100644 --- a/self_organising_systems/biomakerca/environments.py +++ b/self_organising_systems/biomakerca/environments.py @@ -247,9 +247,12 @@ def __str__(self): "AGENT_ROOT", # Leaf: Capable of absorbing air nutrients. "AGENT_LEAF", - # Flower: Capable of performing a reproduce operation. They *tend* to - # consume more nutrients. + # Flower: Capable of performing an asexual reproduce operation. They *tend* + # to consume more nutrients. "AGENT_FLOWER", + # Sexual Flower: Capable of performing a sexual reproduce operation. They + # *tend* to consume more nutrients. + "AGENT_FLOWER_SEXUAL", ] # indexed by the type, it tells how much structure decays. @@ -265,16 +268,18 @@ def __str__(self): "AGENT_ROOT": 5, "AGENT_LEAF": 5, "AGENT_FLOWER": 5, + "AGENT_FLOWER_SEXUAL": 5, } # A modifier of the dissipation based on the agent specialization. -# the first element is for the earth nutrients, the second element is for the +# the first element is for the earth nutrients, the second element is for the # air nutrients. DEFAULT_DISSIPATION_RATE_PER_SPEC_DICT = { "AGENT_UNSPECIALIZED": jp.array([0.5, 0.5]), "AGENT_ROOT": jp.array([1.0, 1.0]), "AGENT_LEAF": jp.array([1.0, 1.0]), "AGENT_FLOWER": jp.array([1.2, 1.2]), + "AGENT_FLOWER_SEXUAL": jp.array([1.2, 1.2]), } # Colors for visualising the default types. @@ -289,6 +294,7 @@ def __str__(self): "AGENT_ROOT": jp.array([0.52, 0.39, 0.14]), # RGB: 133,99,36 "AGENT_LEAF": jp.array([0.16, 0.49, 0.10]), # RGB: 41,125,26 "AGENT_FLOWER": jp.array([1., 0.42, 0.71]), # RGB: 255,107,181 + "AGENT_FLOWER_SEXUAL": jp.array([1., 0.749, 0.]), # RGB: 255,191,0 } def convert_string_dict_to_type_array(d, types): @@ -339,6 +345,7 @@ def __init__( types = self.types # setup material specific properties. If you are subclassing this, consider # changing these values manually. + #TODO make these uint32 self.intangible_mats = jp.array([types.VOID, types.AIR], dtype=jp.int32) self.gravity_mats = jp.concatenate([ jp.array([types.EARTH], dtype=jp.int32), self.agent_types], 0) @@ -411,11 +418,18 @@ class EnvConfig: seed can be placed. reproduce_max_dist: the maximum distance from a reproducing flower where a seed can be placed. - n_reproduce_per_step: how many reproduce ops can be selected per step to be - executed. In effect, this means that flowers may not execute reproduce ops - as soon as they desire, but they may wait, depending on how many other - flowers are asking the same in the environment. See that competition over - a scarse external resource (bees, for instance). + n_reproduce_per_step: how many asexual reproduce ops can be selected per + step to be executed. In effect, this means that flowers may not execute + reproduce ops as soon as they desire, but they may wait, depending on how + many other flowers are asking the same in the environment. See that + competition over a scarse external resource (bees, for instance). + n_sexual_reproduce_per_step: how many sexual reproduce ops can be selected + per step to be executed. Note that since 2 ops are needed per each sexual + reproduction, we will perform n_sexual_reproduce_per_step by getting twice + as many reproduce ops. In effect, this means that flowers may not execute + reproduce ops as soon as they desire, but they may wait, depending on how + many other flowers are asking the same in the environment. See that + competition over a scarse external resource (bees, for instance). nutrient_cap: maximum value of nutrients for agents. material_nutrient_cap: maximum value of nutrients for other materials. max_lifetime: the maximum lifetime of an organism. Agents age at every step. @@ -446,6 +460,7 @@ def __init__(self, reproduce_min_dist=5, reproduce_max_dist=15, n_reproduce_per_step=2, + n_sexual_reproduce_per_step=2, nutrient_cap=DEFAULT_NUTRIENT_CAP, material_nutrient_cap=DEFAULT_MATERIAL_NUTRIENT_CAP, max_lifetime=int(1e6), @@ -462,6 +477,7 @@ def __init__(self, self.reproduce_min_dist = reproduce_min_dist self.reproduce_max_dist = reproduce_max_dist self.n_reproduce_per_step = n_reproduce_per_step + self.n_sexual_reproduce_per_step = n_sexual_reproduce_per_step self.nutrient_cap = nutrient_cap self.material_nutrient_cap = material_nutrient_cap self.max_lifetime = max_lifetime @@ -596,6 +612,21 @@ def create_default_environment(config, h, w, with_earth=True, return env +def create_multiseed_environment(h, w, config, with_earth=True, + init_nutrient_perc=0.2, place_every=10): + """Create a default environment that contains several different seeds. + + Each seed has a different agent id. + """ + env = create_default_environment(config, h, w, with_earth, init_nutrient_perc) + + aid = 0 + for i in range(1, w//place_every): + env = place_seed(env, place_every*i, config, aid=aid) + aid += 1 + + return env + def infer_width(h, width_type): """Infer the width of the environment. @@ -791,9 +822,7 @@ def hsl_to_rgb(h,s,l): """ l_lt_05 = jp.array(l < 0.5).astype(jp.float32) q = l_lt_05 * l * (1 + s) + (1. - l_lt_05) * (l + s - l * s) - print(q) p = 2 * l - q - print(p) return jp.stack( [hue_to_rgb(p,q,h+1/3), hue_to_rgb(p,q,h), hue_to_rgb(p,q,h-1/3)], -1) diff --git a/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb b/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb index 0b382bb..3042e6e 100644 --- a/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb +++ b/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb @@ -73,8 +73,6 @@ "from self_organising_systems.biomakerca.step_maker import step_env\n", "from self_organising_systems.biomakerca.display_utils import zoom\n", "from self_organising_systems.biomakerca.custom_ipython_display import display\n", - "from self_organising_systems.biomakerca.env_logic import ReproduceOp\n", - "from self_organising_systems.biomakerca.env_logic import env_perform_one_reproduce_op\n", "\n", "import cv2\n", "import numpy as np\n", @@ -222,10 +220,6 @@ "# How many unique programs (organisms) are allowed in the simulation.\n", "N_MAX_PROGRAMS = 128\n", "\n", - "# if True, every 50 steps we check whether the agents go extinct. If they did,\n", - "# we replace a seed in the environment.\n", - "replace_if_extinct = False\n", - "\n", "# The number of frames of the video. This is NOT the number of steps.\n", "# The total number of steps depend on the number of steps per frame, which can\n", "# vary over time.\n", @@ -276,25 +270,6 @@ " env, programs = step_env(\n", " ku, env, config, agent_logic, programs, do_reproduction=True,\n", " mutate_programs=True, mutator=mutator)\n", - " if replace_if_extinct and step % 50 == 0:\n", - " # check if there is no alive cell.\n", - " any_alive = jit(lambda type_grid: evm.is_agent_fn(type_grid).sum() \u003e 0)(env.type_grid)\n", - " if not any_alive:\n", - " # Then place a new seed.\n", - " agent_init_nutrients = (config.dissipation_per_step * 4 +\n", - " config.specialize_cost)\n", - " ku, key = jr.split(key)\n", - " rpos = jp.stack([0, jr.randint(ku, (),minval=0, maxval=env.type_grid.shape[1])],0)\n", - " ku, key = jr.split(key)\n", - " raid = jr.randint(ku, (),minval=0, maxval=N_MAX_PROGRAMS).astype(jp.uint32)\n", - " repr_op = ReproduceOp(1., rpos, agent_init_nutrients*2, raid)\n", - " ku, key = jr.split(key)\n", - " env = jit(partial(env_perform_one_reproduce_op, config=config))(ku, env, repr_op)\n", - "\n", - " # show it, though\n", - " frame = make_frame(env, step, steps_per_frame)\n", - " for stop_i in range(10):\n", - " video.add_image(frame)\n", "\n", " video.add_image(make_frame(env, step, steps_per_frame))\n", "\n", @@ -375,7 +350,7 @@ "else:\n", " # Extract a living program from the final environment.\n", " aid_flat = env.agent_id_grid.flatten()\n", - " is_agent_flat = evm.is_agent_fn(env.type_grid).flatten().astype(jp.float32)\n", + " is_agent_flat = config.etd.is_agent_fn(env.type_grid).flatten().astype(jp.float32)\n", " n_alive_per_id = jax.ops.segment_sum(is_agent_flat, aid_flat, num_segments=N_MAX_PROGRAMS)\n", " alive_programs = programs[n_alive_per_id\u003e0]\n", " print(\"Extracted {} programs.\".format(alive_programs.shape[0]))\n", diff --git a/self_organising_systems/biomakerca/examples/notebooks/run_sexual_configuration.ipynb b/self_organising_systems/biomakerca/examples/notebooks/run_sexual_configuration.ipynb new file mode 100644 index 0000000..fd93f06 --- /dev/null +++ b/self_organising_systems/biomakerca/examples/notebooks/run_sexual_configuration.ipynb @@ -0,0 +1,523 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "LpQvs2DVfkVy" + }, + "source": [ + "# Biomaker CA: performing advanced runs on a configuration\n", + "\n", + "In this colab we show how to run models on a configuration and how to evaluate them.\n", + "\n", + "This colab allows to choose whether to perform sexual and/or asexual reproduction.\n", + "It also allows for sparse computations of agent logics.\n", + "\n", + "Copyright 2023 Google LLC\n", + "\n", + "Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "you may not use this file except in compliance with the License.\n", + "You may obtain a copy of the License at\n", + "\n", + " https://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "Unless required by applicable law or agreed to in writing, software\n", + "distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "See the License for the specific language governing permissions and\n", + "limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "iuhJcrwNgIEF" + }, + "source": [ + "## Imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "287RU0sup5J2" + }, + "outputs": [], + "source": [ + "#@title install selforg package\n", + "# install the package locally\n", + "!pip install --upgrade -e git+https://github.com/google-research/self-organising-systems.git#egg=self_organising_systems\u0026subdirectory=biomakerca\n", + "# activate the locally installed package (otherwise a runtime restart is required)\n", + "import pkg_resources\n", + "import importlib\n", + "# Reload the resources because we uninstalled and reinstalled some packages.\n", + "importlib.reload(pkg_resources)\n", + "pkg_resources.get_distribution(\"self_organising_systems\").activate()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "DY7J0cz7NUCV" + }, + "outputs": [], + "source": [ + "#@title imports \u0026 notebook utilities\n", + "\n", + "from self_organising_systems.biomakerca import environments as evm\n", + "from self_organising_systems.biomakerca.agent_logic import BasicAgentLogic\n", + "from self_organising_systems.biomakerca.mutators import BasicMutator\n", + "from self_organising_systems.biomakerca.mutators import RandomlyAdaptiveMutator\n", + "from self_organising_systems.biomakerca.mutators import CrossOverSexualMutator\n", + "from self_organising_systems.biomakerca.step_maker import step_env\n", + "from self_organising_systems.biomakerca.display_utils import zoom, tile2d, add_text_to_img, imshow\n", + "from self_organising_systems.biomakerca.custom_ipython_display import display\n", + "from self_organising_systems.biomakerca.env_logic import env_perform_multi_world_reproduce_update\n", + "\n", + "import cv2\n", + "import numpy as np\n", + "import jax.random as jr\n", + "import jax.numpy as jp\n", + "from jax import vmap\n", + "from jax import jit\n", + "import jax\n", + "import time\n", + "import matplotlib.pyplot as plt\n", + "\n", + "import tqdm\n", + "import mediapy as media\n", + "from functools import partial\n", + "\n", + "\n", + "def pad_text(img, text):\n", + " font = cv2.FONT_HERSHEY_SIMPLEX\n", + " orgin = (5, 15)\n", + " fontScale = 0.5\n", + " color = (0, 0, 0)\n", + " thickness = 1\n", + "\n", + " # ensure to preserve even size (assumes the input size was even.\n", + " new_h = img.shape[0]//15\n", + " new_h = new_h if new_h % 2 == 0 else new_h + 1\n", + " img = np.concatenate([np.ones([new_h, img.shape[1], img.shape[2]]), img], 0)\n", + " img = cv2.putText(img, text, orgin, font, fontScale, color, thickness, cv2.LINE_AA)\n", + " return img" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aRRMQ1aNhqv6" + }, + "source": [ + "## Select the configuration, the agent logic and the mutator\n", + "\n", + "Set soil_unbalance_limit to 0 to reproduce the original environment. Set it to 1/3 for having self-balancing environments (recommended)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "dtc32MAfOBTK" + }, + "outputs": [], + "source": [ + "ec_id = \"persistence\" #@param ['persistence', 'pestilence', 'collaboration', 'sideways']\n", + "env_width_type = \"wide\" #@param ['wide', 'landscape', 'square', 'petri', '10x', '20x']\n", + "soil_unbalance_limit = 1/3 #@param [0, \"1/3\"] {type:\"raw\"}\n", + "\n", + "h = 72\n", + "if env_width_type == \"10x\":\n", + " env_width_type = h * 10\n", + "if env_width_type == \"20x\":\n", + " env_width_type = h * 20\n", + "else:\n", + " env_width_type = evm.infer_width(h, env_width_type)\n", + "\n", + "env_and_config = evm.get_env_and_config(ec_id, width_type=env_width_type, h=h)\n", + "_, config = env_and_config\n", + "\n", + "st_env = evm.create_multiseed_environment(h, env_width_type, config)\n", + "\n", + "config.soil_unbalance_limit = soil_unbalance_limit\n", + "reproduction_type = \"asexual\" #@param ['both', 'asexual', 'sexual']\n", + "does_sex_matter = True #@param ['False', 'True'] {type:\"raw\"}\n", + "\n", + "enable_asexual_reproduction = reproduction_type != \"sexual\"\n", + "enable_sexual_reproduction = reproduction_type != \"asexual\"\n", + "\n", + "sex_sensitivity = 1000 #@param [1, 10, 100, 1000] {type:\"raw\"}\n", + "\n", + "agent_model = \"extended\" #@param ['minimal', 'extended']\n", + "agent_logic = BasicAgentLogic(config, minimal_net=agent_model==\"minimal\",\n", + " make_asexual_flowers_likely=enable_asexual_reproduction,\n", + " make_sexual_flowers_likely=enable_sexual_reproduction,\n", + " init_noise=0.001, sex_sensitivity=sex_sensitivity)\n", + "\n", + "n_sparse_max = 2**13 #@param ['None', '2**13', '2**12', '2**11', '2**10'] {type:\"raw\"}\n", + "\n", + "mutator_type = \"randomly_adaptive\" #@param ['basic', 'randomly_adaptive']\n", + "sd = 1e-3\n", + "mutator = (BasicMutator(sd=sd, change_perc=0.2) if mutator_type == \"basic\"\n", + " else RandomlyAdaptiveMutator(init_sd=sd, change_perc=0.2))\n", + "sexual_mutator = CrossOverSexualMutator(mutator, n_frequencies=64)\n", + "\n", + "exp_id = \"{}_sm^{}_sens^{}_w^{}\".format(reproduction_type, does_sex_matter, sex_sensitivity, env_width_type)\n", + "print(exp_id)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "y3p9hXhKuSk2" + }, + "source": [ + "## Optionally, modify the config for custom configurations" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "cKshhtYDvE81" + }, + "outputs": [], + "source": [ + "print(\"Current config:\")\n", + "print('\\n'.join(\"%s: %s\" % item for item in vars(config).items()))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "H1wR-EQium-X" + }, + "outputs": [], + "source": [ + "## Examples for modifying the config\n", + "## Uncomment relevant lines or do like them.\n", + "\n", + "## Regardless, to trigger the recomputation of step_env and similar,\n", + "## config needs to be a new object! So, first, we create a new copy.\n", + "import copy\n", + "config = copy.copy(config)\n", + "\n", + "## Change simple isolated parameters (most of them)\n", + "# config.struct_integrity_cap = 100\n", + "# config.max_lifetime = 500\n", + "## Vectors can be modified either by writing new vectors:\n", + "# config.dissipation_per_step = jp.array([0.02, 0.02])\n", + "## Or by multiplying previous values. Note that they are immutable!\n", + "# config.dissipation_per_step = config.dissipation_per_step * 2\n", + "\n", + "## agent_state_size is trickier, because it influences env_state_size.\n", + "## So you can either create a new config:\n", + "## Note that you would have to insert all values that you don't want to take\n", + "## default initializations.\n", + "# config = evm.EnvConfig(agent_state_size=4)\n", + "## Or you can just modify env_state_size as well.\n", + "## (env_state_size = agent_state_size + 4) for now.\n", + "# config.agent_state_size = 4\n", + "# config.env_state_size = config.agent_state_size + 4\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Y-xLUg_Uh14k" + }, + "source": [ + "## Perform a simulation\n", + "\n", + "Consider modifying the code to vary the extent of the simulation and video configs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8tj0U8naJDC3" + }, + "outputs": [], + "source": [ + "# create auxiliary frames that show interesting counters.\n", + "\n", + "def make_aux_frame(n_asex, n_sex):\n", + " img = np.ones([130 if reproduction_type == \"both\" else 65, 550, 3])\n", + " yorigin = 50\n", + " if enable_asexual_reproduction:\n", + " img = add_text_to_img(\n", + " img, \"Asexual reproductions: {}\".format(n_asex),\n", + " origin=(20, yorigin), color=\"black\")\n", + " yorigin = 100\n", + " if enable_sexual_reproduction:\n", + " img = add_text_to_img(\n", + " img, \"Sexual reproductions: {}\".format(n_sex),\n", + " origin=(20, yorigin), color=\"black\")\n", + " return img\n", + "\n", + "imshow(make_aux_frame(100000, 212122))\n", + "\n", + "def make_nsexes_frame(n_sexes):\n", + " img = np.ones([65, 300, 3])\n", + " yorigin = 50\n", + " img = add_text_to_img(\n", + " img, \"Num sexes: {}\".format(n_sexes),\n", + " origin=(20, yorigin), color=\"black\")\n", + " return img\n", + "\n", + "imshow(make_nsexes_frame(200))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6OMyOYUUHodo" + }, + "outputs": [], + "source": [ + "\n", + "@partial(jit, static_argnames=[\"config\", \"n_max_programs\"])\n", + "def get_alive_programs_mask(env, config, n_max_programs):\n", + " aid_flat = env.agent_id_grid.flatten()\n", + " is_agent_flat = config.etd.is_agent_fn(env.type_grid).flatten().astype(jp.float32)\n", + " n_alive_per_id = jax.ops.segment_sum(is_agent_flat, aid_flat, num_segments=n_max_programs)\n", + "\n", + " has_alive = n_alive_per_id \u003e 0\n", + " return has_alive\n", + "\n", + "@jit\n", + "def get_alive_and_sexes(env, programs):\n", + " has_alive = get_alive_programs_mask(env, config, N_MAX_PROGRAMS)\n", + " all_sexes = vmap(agent_logic.get_sex)(mutator.split_params(programs)[0])\n", + " return has_alive, all_sexes\n", + "\n", + "def get_num_sexes(env, programs):\n", + " has_alive, all_sexes = get_alive_and_sexes(env, programs)\n", + " has_alive = np.array(has_alive)\n", + " all_sexes = np.array(all_sexes)\n", + "\n", + " sexes = all_sexes[has_alive]\n", + " return len(np.unique(sexes))\n", + "\n", + "def run_env(\n", + " key, programs, env, n_steps, step_f,\n", + " curr_asexual_repr = 0, curr_sexual_repr = 0,\n", + " zoom_sz=12,\n", + " steps_per_frame=2, when_to_double_speed=[100, 500, 1000, 2000, 5000]):\n", + "\n", + " fps = 20\n", + " def make_frame(env):\n", + " return zoom(evm.grab_image_from_env(env, config),zoom_sz)\n", + "\n", + " frame = make_frame(env)\n", + "\n", + " # remember that the metrics are per step, right now, and that are 'inbetween'\n", + " # frames, at best.\n", + " n_asexual_repr_log = []\n", + " n_sexual_repr_log = []\n", + " n_sexes_log = []\n", + "\n", + " aux_frames = [make_aux_frame(0, 0)]\n", + " num_sexes_frames = [make_nsexes_frame(get_num_sexes(env, programs))]\n", + "\n", + " out_file = \"video.mp4\"\n", + " with media.VideoWriter(out_file, shape=frame.shape[:2], fps=fps, crf=18\n", + " ) as video:\n", + " video.add_image(frame)\n", + " for i in tqdm.trange(n_steps):\n", + " if i in when_to_double_speed:\n", + " steps_per_frame *= 2\n", + "\n", + " key, ku = jr.split(key)\n", + " (env, programs), metrics = step_f(ku, env, programs=programs)\n", + "\n", + " if enable_asexual_reproduction:\n", + " step_asex_repr = int(metrics[\"asexual_reproduction\"][0].sum())\n", + " n_asexual_repr_log.append(step_asex_repr)\n", + " curr_asexual_repr += step_asex_repr\n", + " if enable_sexual_reproduction:\n", + " step_sexual_repr = int(metrics[\"sexual_reproduction\"][0].sum())\n", + " n_sexual_repr_log.append(step_sexual_repr)\n", + " curr_sexual_repr += step_sexual_repr\n", + "\n", + " # get sexes alive\n", + " num_sexes = get_num_sexes(env, programs)\n", + " n_sexes_log.append(num_sexes)\n", + " if i % steps_per_frame == 0:\n", + " video.add_image(make_frame(env))\n", + " aux_frames.append(make_aux_frame(curr_asexual_repr, curr_sexual_repr))\n", + " num_sexes_frames.append(make_nsexes_frame(num_sexes))\n", + "\n", + "\n", + " media.show_video(media.read_video(out_file))\n", + " media.show_video(aux_frames, fps=fps)\n", + " media.show_video(num_sexes_frames, fps=fps)\n", + " ret_metrics = {'n_asexual_repr_log': n_asexual_repr_log,\n", + " 'n_sexual_repr_log': n_sexual_repr_log,\n", + " 'n_sexes_log': n_sexes_log,\n", + " 'curr_asexual_repr': curr_asexual_repr,\n", + " 'curr_sexual_repr': curr_sexual_repr}\n", + " return programs, env, ret_metrics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "65WC-fp1aI1Y" + }, + "outputs": [], + "source": [ + "key = jr.PRNGKey(42)\n", + "\n", + "# How many unique programs (organisms) are allowed in the simulation.\n", + "N_MAX_PROGRAMS = 512\n", + "\n", + "# for 20x environments, you need shorter videos.\n", + "n_steps = 15000\n", + "\n", + "# on what FRAME to double speed.\n", + "when_to_double_speed = [100, 200, 300, 400, 500]\n", + "# on what FRAME to reset speed.\n", + "when_to_reset_speed = []\n", + "fps = 20\n", + "# this affects the size of the image. If this number is not even, the resulting\n", + "# video *may* not be supported by all renderers.\n", + "zoom_sz = 4\n", + "\n", + "ku, key = jr.split(key)\n", + "programs = vmap(agent_logic.initialize)(jr.split(ku, N_MAX_PROGRAMS))\n", + "ku, key = jr.split(key)\n", + "programs = vmap(mutator.initialize)(jr.split(ku, programs.shape[0]), programs)\n", + "\n", + "env = st_env\n", + "\n", + "step_f = partial(step_env, config=config, agent_logic=agent_logic, do_reproduction=True,\n", + " enable_asexual_reproduction=enable_asexual_reproduction,\n", + " enable_sexual_reproduction=enable_sexual_reproduction,\n", + " does_sex_matter=does_sex_matter,\n", + " mutate_programs=True, mutator=mutator, sexual_mutator=sexual_mutator,\n", + " n_sparse_max=n_sparse_max, return_metrics=True)\n", + "\n", + "step = 0\n", + "# how many steps per frame we start with. This gets usually doubled many times\n", + "# during the simulation.\n", + "# In the article, we usually use 2 or 4 as the starting value, sometimes 1.\n", + "steps_per_frame = 2\n", + "\n", + "ku, key = jr.split(key)\n", + "programs, env, metrics = run_env(\n", + " ku, programs, env, n_steps, step_f, zoom_sz=6)\n", + "\n", + "\n", + "\n", + "def running_average(a, n):\n", + " a = np.concatenate([np.full([n], a[0]), a], axis=0)\n", + " return np.convolve(a, np.ones(n)/n, mode=\"valid\")\n", + "\n", + "if enable_asexual_reproduction:\n", + " plt.plot(running_average(metrics['n_asexual_repr_log'], 100), label=\"n_asexual_repr_log\")\n", + "if enable_sexual_reproduction:\n", + " plt.plot(running_average(metrics['n_sexual_repr_log'], 100), label=\"n_sexual_repr_log\")\n", + "plt.legend()\n", + "plt.show()\n", + "\n", + "plt.plot(running_average(metrics['n_sexes_log'], 100), label=\"n_sexes_log\")\n", + "plt.legend()\n", + "plt.show()\n", + "\n", + "aid_flat = env.agent_id_grid.flatten()\n", + "is_agent_flat = config.etd.is_agent_fn(env.type_grid).flatten().astype(jp.float32)\n", + "n_alive_per_id = jax.ops.segment_sum(is_agent_flat, aid_flat, num_segments=N_MAX_PROGRAMS)\n", + "alive_programs = programs[n_alive_per_id\u003e0]\n", + "print(\"Extracted {} programs.\".format(alive_programs.shape[0]))\n", + "print(\"sexes:\", vmap(agent_logic.get_sex)(mutator.split_params(alive_programs)[0]))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oBggxzAZm4lM" + }, + "outputs": [], + "source": [ + "# continue...\n", + "n_steps = 25000\n", + "ku, key = jr.split(key)\n", + "\n", + "programs, env, metrics = run_env(\n", + " ku, programs, env, n_steps, step_f, zoom_sz=6,\n", + " steps_per_frame=64, when_to_double_speed=[],\n", + " curr_asexual_repr=metrics['curr_asexual_repr'],\n", + " curr_sexual_repr=metrics['curr_sexual_repr'])\n", + "\n", + "\n", + "if enable_asexual_reproduction:\n", + " plt.plot(running_average(metrics['n_asexual_repr_log'], 100), label=\"n_asexual_repr_log\")\n", + "if enable_sexual_reproduction:\n", + " plt.plot(running_average(metrics['n_sexual_repr_log'], 100), label=\"n_sexual_repr_log\")\n", + "plt.legend()\n", + "plt.show()\n", + "\n", + "plt.plot(running_average(metrics['n_sexes_log'], 100), label=\"n_sexes_log\")\n", + "plt.legend()\n", + "plt.show()\n", + "\n", + "aid_flat = env.agent_id_grid.flatten()\n", + "is_agent_flat = config.etd.is_agent_fn(env.type_grid).flatten().astype(jp.float32)\n", + "n_alive_per_id = jax.ops.segment_sum(is_agent_flat, aid_flat, num_segments=N_MAX_PROGRAMS)\n", + "alive_programs = programs[n_alive_per_id\u003e0]\n", + "print(\"Extracted {} programs.\".format(alive_programs.shape[0]))\n", + "print(\"sexes:\", vmap(agent_logic.get_sex)(mutator.split_params(alive_programs)[0]))\n" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "V100", + "last_runtime": { + "build_target": "//learning/grp/tools/ml_python:ml_notebook", + "kind": "private" + }, + "machine_shape": "hm", + "name": "run_sexual_configuration.ipynb", + "private_outputs": true, + "provenance": [ + { + "file_id": "1rgSP9D1gT3S4Aljk-jLys8dFbla2InCu", + "timestamp": 1708355883667 + }, + { + "file_id": "/piper/depot/google3/third_party/py/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb?workspaceId=etr:biomaker::citc", + "timestamp": 1701254768987 + }, + { + "file_id": "1ADfcMRj-JmfN6VUIcuqU-3bTMGrWSkj_", + "timestamp": 1688723295778 + }, + { + "file_id": "1XY102qIEc9MY9hd-Jb6Oirmyw7ga2LZL", + "timestamp": 1688637712371 + } + ], + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/self_organising_systems/biomakerca/extensions/eruption.py b/self_organising_systems/biomakerca/extensions/eruption.py index b1e029f..dac3e7f 100644 --- a/self_organising_systems/biomakerca/extensions/eruption.py +++ b/self_organising_systems/biomakerca/extensions/eruption.py @@ -79,7 +79,7 @@ def is_burnable_fn(t, etd): """Return True if t is any of (LEAF,FLOWER,UNSPECIALIZED) agent types.""" burnable_types = jp.array([ etd.types.AGENT_LEAF, etd.types.AGENT_FLOWER, - etd.types.AGENT_UNSPECIALIZED]) + etd.types.AGENT_FLOWER_SEXUAL, etd.types.AGENT_UNSPECIALIZED]) return (t == burnable_types).any(axis=-1) @@ -255,7 +255,7 @@ def get_eruption_config(): return evm.get_env_and_config( "persistence", width_type="petri", h=10, etd=EruptionTypeDef()).config -### Eruption environment +### Original Eruption environment def create_eruption_env(h, config): """Create the Eruption environment. @@ -408,7 +408,7 @@ def test_freq_lava(key, st_env, config, init_program, agent_logic, mutator, def make_frame(env, lava_perc): return add_text_to_img( zoom(evm.grab_image_from_env(env, config), zoom_sz), - "LAVA FREQUENCY: {:.3f}%".format(lava_perc*100), + "LAVA FREQUENCY: {:.3f}%".format(lava_perc*100), origin=(5, 35), fontScale=1.) diff --git a/self_organising_systems/biomakerca/mutators.py b/self_organising_systems/biomakerca/mutators.py index d82bf25..3705f08 100644 --- a/self_organising_systems/biomakerca/mutators.py +++ b/self_organising_systems/biomakerca/mutators.py @@ -21,14 +21,15 @@ limitations under the License. """ from abc import ABC, abstractmethod +import math import jax.numpy as jp import jax.random as jr from self_organising_systems.biomakerca.utils import stringify_class class Mutator(ABC): - """Interface of all mutators. - + """Interface of all asexual mutators. + The abstract methods need to be implemented in order to allow for in-environment mutations, through the method step_maker.step_env. """ @@ -36,9 +37,11 @@ class Mutator(ABC): @abstractmethod def initialize(self, key, p): """Initialize mutation parameters. - + p must be one-dimensional. Return a concatenation of p and the mutation parameters. + Advice: it is better if related params are contiguous to one another. This + is because SexualMutators likely use crossing over as a mechanism. """ pass @@ -124,14 +127,19 @@ def __init__(self, init_sd, change_perc: float | None, meta_sd_perc=0.1, self.min_sd = min_sd self.max_sd = max_sd + def join_params(self, p, sd): + return jp.ravel(jp.column_stack((p, sd))) + def initialize(self, key, p): - # works also if batched. + # Does not work if batched. (use vmap) sd = jp.full_like(p, self.init_sd) - return jp.concatenate([p, sd], -1) + return self.join_params(p, sd) def split_params(self, p): # works also if batched. - return jp.split(p, 2, axis=-1) + p_e, sd_e = jp.split( + p.reshape(p.shape[:-1]+ (p.shape[-1]//2, 2)), 2, axis=-1) + return p_e[..., 0], sd_e[..., 0] def mutate(self, key, p): # not batched. @@ -154,4 +162,81 @@ def mutate(self, key, p): new_mu = mu + dmu new_sd = (sd + dsd).clip(self.min_sd, self.max_sd) - return jp.concatenate([new_mu, new_sd], -1) + return self.join_params(new_mu, new_sd) + +### Sexual mutators + + +class SexualMutator(ABC): + """Interface of all sexual mutators. + + The abstract methods need to be implemented in order to allow for + in-environment sexual mutations, through the method step_maker.step_env. + + For now, a SexualMutator does not accept variable parameters. + This is because it uses a Mutator for extra variation. + Ideally, the Mutator used is the same of asexual reproduction, but you can + put custom ones here as well. + """ + + def __init__(self, mutator: Mutator): + """Set the mutator for the general variation of parameters. + """ + self.mutator = mutator + + @abstractmethod + def mutate(self, key, p1, p2): + """Mutate params. + + The inputs p1 and p2 must be all params, as generated by 'initialize'. + This method can mutate mutation params too. + p must be one-dimensional. + """ + pass + + def __str__(self): + return stringify_class(self) + + +def valuenoise1d(key, f, n, interpolation="linear"): + """Create continous noise, useful for make a XOver mask.""" + ku, key = jr.split(key) + grads = jr.uniform(ku, (f,), minval=-1, maxval=1) + n_repeats = int(math.ceil(n/(f-1))) + prev = jp.repeat(grads[:-1], n_repeats, axis=0)[:n] + next = jp.repeat(grads[1:], n_repeats, axis=0)[:n] + # fraction of where you are + t = jp.tile(jp.linspace(0, 1, n_repeats+1)[:-1], f-1)[:n] + + # linear interpolation because we don't need anything more complex. + # if you want it cubic, do this: + if interpolation == "cubic": + t = t*t*t*(t*(t*6 - 15) + 10) + return prev + t * (next - prev) + + +class CrossOverSexualMutator(SexualMutator): + """Sexual mutator that performs a classic Crossing over recombination mutation. + + This mutator also performs an asexual mutation through the mutator input. + """ + + def __init__(self, mutator: Mutator, n_frequencies=16): + super().__init__(mutator) + self.n_frequencies = n_frequencies + + def mutate(self, key, p1, p2): + """Mutate params. + + The inputs p1 and p2 must be all params, as generated by 'initialize'. + This method can mutate mutation params too. + p must be one-dimensional. + """ + ku, key = jr.split(key) + xo_mask = valuenoise1d(ku, self.n_frequencies, p1.shape[0], "linear") > 0 + + new_p = xo_mask * p1 + (1. - xo_mask) * p2 + + ku, key = jr.split(key) + return self.mutator.mutate(ku, new_p) + diff --git a/self_organising_systems/biomakerca/step_maker.py b/self_organising_systems/biomakerca/step_maker.py index 0ce57c2..e25da18 100644 --- a/self_organising_systems/biomakerca/step_maker.py +++ b/self_organising_systems/biomakerca/step_maker.py @@ -44,13 +44,15 @@ from self_organising_systems.biomakerca.env_logic import process_structural_integrity_n_times from self_organising_systems.biomakerca.environments import EnvConfig from self_organising_systems.biomakerca.environments import Environment -from self_organising_systems.biomakerca.mutators import Mutator +from self_organising_systems.biomakerca.mutators import Mutator, SexualMutator @partial(jit, static_argnames=[ - "config", "agent_logic", "excl_fs", "do_reproduction", "mutate_programs", - "mutator", "intercept_reproduction"]) + "config", "agent_logic", "excl_fs", "do_reproduction", + "enable_asexual_reproduction", "enable_sexual_reproduction", + "does_sex_matter", "mutate_programs", "mutator", "sexual_mutator", + "intercept_reproduction", "n_sparse_max", "return_metrics"]) def step_env( key: KeyType, env: Environment, config: EnvConfig, agent_logic: AgentLogic, @@ -59,10 +61,16 @@ def step_env( tuple[EnvTypeType, Callable[[KeyType, PerceivedData], ExclusiveOp]]] = None, do_reproduction=True, + enable_asexual_reproduction=True, + enable_sexual_reproduction=False, + does_sex_matter=True, mutate_programs=False, mutator: (Mutator | None) = None, + sexual_mutator: (SexualMutator | None) = None, intercept_reproduction=False, - min_repr_energy_requirement=None): + min_repr_energy_requirement=None, + n_sparse_max: (int | None ) = None, + return_metrics=False): """Perform one step for the environment. There are several different settings for performing a step. The most important @@ -80,6 +88,12 @@ def step_env( excl_fs: the exclusive logic of materials. Defaults to AIR spreading through VOID, and EARTH acting like falling-sand. do_reproduction: whether reproduction is enabled. + enable_asexual_reproduction: if set to True, asexual reproduction is + enabled. if do_reproduction is enabled, at least one of + enable_asexual_reproduction and enable_sexual_reproduction must be True. + enable_sexual_reproduction: if set to True, sexual reproduction is enabled. + does_sex_matter: if set to True, and enable_sexual_reproduction is True, + then sexual reproduction can only occur between different sex entities. mutate_programs: relevant only if do_reproduction==True. In that case, determines whether reproduction is performed with or without mutation. Beware! Reproduction *without* mutation creates agents with identical @@ -88,6 +102,8 @@ def step_env( If set to true, 'mutator' must be a valid Mutator class. mutator: relevant only if we reproduce with mutation. In that case, mutator determines how to extract parameters and how to modify them. + sexual_mutator: relevant only if we reproduce with mutation. In that case, + it determines how to perform sexual reproduction. intercept_reproduction: useful for petri-dish-like experiments. If set to true, whenever a ReproduceOp triggers, instead of creating a new seed, we simply destroy the flower and record its occurrence. We consider it a @@ -98,14 +114,20 @@ def step_env( min_repr_energy_requirement: relevant only if intercepting reproductions. Determines whether the intercepted seed would have had enough energy to count as a successful reproduction. + n_sparse_max: either an int or None. If set to int, we will use a budget for + the amounts of agent operations allowed at each step. + return_metrics: if True, returns reproduction metrics. May be extended to + return more metrics in the future. Returns: an updated environment. If intercept_reproduction is True, returns also the number of successful reproductions intercepted. """ + assert (not do_reproduction or + (enable_asexual_reproduction or enable_sexual_reproduction)) etd = config.etd if excl_fs is None: excl_fs = ((etd.types.AIR, air_cell_op), (etd.types.EARTH, earth_cell_op)) - + if mutate_programs: agent_params, _ = vmap(mutator.split_params)(programs) else: @@ -135,17 +157,25 @@ def step_env( else: ku, key = jr.split(key) if mutate_programs: - env, programs = env_perform_reproduce_update( + # metrics are supported only here. + repr_result = env_perform_reproduce_update( ku, env, repr_programs, config, repr_f, mutate_programs, programs, - mutator.mutate) + mutator.mutate, enable_asexual_reproduction, + enable_sexual_reproduction, does_sex_matter, sexual_mutator.mutate, + mutator.split_params, agent_logic.get_sex, n_sparse_max, + return_metrics) + if return_metrics: + (env, programs), metrics = repr_result + else: + env, programs = repr_result else: env = env_perform_reproduce_update( - ku, env, repr_programs, config, repr_f) + ku, env, repr_programs, config, repr_f, n_sparse_max=n_sparse_max) # parallel updates k1, key = jr.split(key) env = env_perform_parallel_update( - k1, env, par_programs, config, agent_logic.par_f) + k1, env, par_programs, config, agent_logic.par_f, n_sparse_max) # energy absorbed and generated by materials. env = process_energy(env, config) @@ -153,11 +183,12 @@ def step_env( # exclusive updates k1, key = jr.split(key) env = env_perform_exclusive_update( - k1, env, excl_programs, config, excl_fs, agent_logic.excl_f) + k1, env, excl_programs, config, excl_fs, agent_logic.excl_f, n_sparse_max) # increase age. env = env_increase_age(env, etd) rval = (env, programs) if mutate_programs else env rval = (rval, n_successful_repr) if intercept_reproduction else rval + rval = (rval, metrics) if return_metrics else rval return rval