from functools import partial import jax import jax.numpy as jnp from flax import struct import chex from agents.population_interface import AgentPopulation def add_partners_to_buffer(population, buffer, params_batch): ''' Add multiple partner checkpoints to the buffer. ''' def add_single_partner(carry_buffer, params): return population.add_agent(carry_buffer, params), None new_buffer, _ = jax.lax.scan( add_single_partner, buffer, params_batch ) return new_buffer def get_final_buffer(buffer): ''' Get the final buffer from the buffer generated by the last iteration of OEL. buffer is a PopulationBuffer object with fields params: (max_pop_size, 1, 1, max_pop_size, ...) scores: (max_pop_size, 1, max_pop_size) ages: (max_pop_size, 1, 1, max_pop_size) filled: (max_pop_size, 1, max_pop_size) filled_count: (max_pop_size, 1, 1) write_ptr: (max_pop_size, 1, 1) # configuration parameters buffer_size: int staleness_coef: float temp: float ''' # adding extra 1 dimension purely for compatibility with previous versions # TODO: remove this hack for v1 last_params = jax.tree.map(lambda x: x[-1][jnp.newaxis], buffer.params) last_scores = jax.tree.map(lambda x: x[-1][jnp.newaxis], buffer.scores) last_ages = jax.tree.map(lambda x: x[-1][jnp.newaxis], buffer.ages) last_filled = jax.tree.map(lambda x: x[-1][jnp.newaxis], buffer.filled) last_filled_count = jax.tree.map(lambda x: x[-1][jnp.newaxis], buffer.filled_count) last_write_ptr = jax.tree.map(lambda x: x[-1][jnp.newaxis], buffer.write_ptr) return PopulationBuffer( params=last_params, scores=last_scores, ages=last_ages, filled=last_filled, filled_count=last_filled_count, write_ptr=last_write_ptr, buffer_size=buffer.buffer_size, staleness_coef=buffer.staleness_coef, temp=buffer.temp, ) class PopulationBuffer(struct.PyTreeNode): """PyTree structure to store population parameters with fixed buffer size.""" params: chex.ArrayTree # Parameters for all agents in the buffer scores: chex.Array # Score for each agent ages: chex.Array # Age (staleness) of each agent filled: chex.Array # Boolean mask for filled slots filled_count: chex.Array # Count of filled slots write_ptr: chex.Array # FIFO write pointer for uniform/softmax eviction # Configuration parameters buffer_size: int = struct.field(pytree_node=False, default=100) staleness_coef: float = struct.field(pytree_node=False, default=0.3) temp: float = struct.field(pytree_node=False, default=1.0) class BufferedPopulation(AgentPopulation): """Population that maintains a buffer of agent parameters that grows over time. This extends the AgentPopulation class and provides methods to add new agents to the buffer and sample from existing agents based on their scores. """ def __init__(self, max_pop_size, policy_cls, sampling_strategy="plr", staleness_coef=0.3, temp=1.0): """Initialize a buffered population. Args: max_pop_size: Maximum number of agents in the population policy_cls: Agent policy class sampling_strategy: Method for sampling partners ('plr', 'uniform', or 'softmax'). staleness_coef: Weight for staleness in 'plr' sampling. If 0, only the score is used for sampling. If scores are 1s (default), sampling becomes uniform over filled slots (equivalent to 'uniform' strategy). temp: Temperature for score-based sampling distribution. """ super().__init__(max_pop_size, policy_cls) self.max_pop_size = max_pop_size self.staleness_coef = staleness_coef self.temp = temp self.sampling_strategy = sampling_strategy assert sampling_strategy in ["plr", "uniform", "softmax"], "sampling_strategy must be 'plr', 'uniform' or 'softmax'" @partial(jax.jit, static_argnums=(0,)) def reset_buffer(self, example_params): """Initialize the buffer with zeros of the appropriate shape. Args: example_params: Example parameters to determine shape and dtype Returns: A new PopulationBuffer """ # Initialize buffer with zeros like example_params params = jax.tree_map( lambda x: jnp.zeros((self.max_pop_size,) + x.shape, x.dtype), example_params ) return PopulationBuffer( params=params, scores=jnp.ones(self.max_pop_size), # Default score of 1 ages=jnp.zeros(self.max_pop_size, dtype=jnp.int32), filled=jnp.zeros(self.max_pop_size, dtype=bool), filled_count=jnp.zeros(1, dtype=jnp.int32), write_ptr=jnp.zeros(1, dtype=jnp.int32), buffer_size=self.max_pop_size, staleness_coef=self.staleness_coef, temp=self.temp, ) @partial(jax.jit, static_argnums=(0,)) def _get_uniform_sampling_dist(self, buffer): """Get uniform distribution over filled slots.""" filled_count = jnp.maximum(buffer.filled.sum(), 1) # Avoid division by zero return buffer.filled / filled_count @partial(jax.jit, static_argnums=(0,)) def _get_plr_sampling_dist(self, buffer): """Get distribution for sampling agents based on scores and staleness (PLR strategy). Args: buffer: The population buffer Returns: Distribution over agents """ # Score distribution (use scores only for filled slots). # Raise scores to the power of 1/temp before normalising so that temp # actually controls the sharpness of the distribution (higher temp -> flatter). # Requires scores > 0, which is ensured by the default score of 1.0. score_dist = jnp.where(buffer.filled, jnp.power(buffer.scores, 1.0 / buffer.temp), 0.0) score_sum = score_dist.sum() # sum of scores over filled slots # if filled slots exist, use score distribution. otherwise, uniform over filled slots score_dist = jnp.where( score_sum > 0, score_dist / score_sum, buffer.filled / jnp.maximum(buffer.filled.sum(), 1) # Uniform over filled slots ) # Staleness distribution staleness_scores = buffer.ages * buffer.filled staleness_sum = staleness_scores.sum() staleness_dist = jnp.where( staleness_sum > 0, staleness_scores / staleness_sum, score_dist # If no staleness, use score distribution ) # Combined distribution return (1 - buffer.staleness_coef) * score_dist + buffer.staleness_coef * staleness_dist @partial(jax.jit, static_argnums=(0,)) def _get_softmax_sampling_dist(self, buffer): """Get softmax distribution over filled slots based on scores.""" # Mask unfilled slots with -inf so they have 0 probability after softmax masked_scores = jnp.where(buffer.filled, buffer.scores / buffer.temp, -1e10) return jax.nn.softmax(masked_scores) @partial(jax.jit, static_argnums=(0,)) def _get_sampling_dist(self, buffer): """Select distribution based on the instance's sampling strategy.""" if self.sampling_strategy == "uniform": return self._get_uniform_sampling_dist(buffer) elif self.sampling_strategy == "plr": return self._get_plr_sampling_dist(buffer) elif self.sampling_strategy == "softmax": return self._get_softmax_sampling_dist(buffer) else: raise ValueError(f"Unknown sampling strategy: {self.sampling_strategy}") @partial(jax.jit, static_argnums=(0,)) def _get_next_insert_idx(self, buffer): """Get index for next insertion. - 'plr': while filling, uses filled_count; once full, evicts the agent with the lowest PLR sampling probability (argmin). - 'uniform' / 'softmax': always uses a FIFO write pointer that wraps around modulo buffer_size, so every slot is replaced in round-robin order once the buffer is full. Args: buffer: The population buffer Returns: Index to insert the next agent """ if self.sampling_strategy == "plr": # PLR: argmin eviction once full sampling_dist = self._get_plr_sampling_dist(buffer) return jax.lax.cond( jnp.less(buffer.filled_count[0], buffer.buffer_size), lambda: buffer.filled_count[0], lambda: jnp.argmin(sampling_dist) ) else: # FIFO: write_ptr already wraps correctly for both the initial fill # phase (0, 1, 2, ...) and subsequent overwrites. return buffer.write_ptr[0] @partial(jax.jit, static_argnums=(0,)) def add_agent(self, buffer, new_params, score=None): """Add a new agent to the buffer. Args: buffer: The population buffer new_params: Parameters of the new agent to add score: Optional score for the new agent (default to 1.0) Returns: Updated population buffer """ # Find index to insert insert_idx = self._get_next_insert_idx(buffer) # Default score if score is None: score = 1.0 # Update parameters, filled, count, score, age new_buffer_params = jax.tree_map( lambda x, p: x.at[insert_idx].set(p), buffer.params, new_params ) new_filled = buffer.filled.at[insert_idx].set(True) new_filled_count = jnp.minimum(buffer.filled_count + 1, jnp.array([buffer.buffer_size])) new_scores = buffer.scores.at[insert_idx].set(score) # Reset age for new agent, increment for others ages_plus_one = (buffer.ages + 1) * buffer.filled new_ages = ages_plus_one.at[insert_idx].set(0) # Advance FIFO write pointer (used by uniform/softmax strategies) new_write_ptr = (buffer.write_ptr + 1) % buffer.buffer_size return buffer.replace( params=new_buffer_params, filled=new_filled, filled_count=new_filled_count, scores=new_scores, ages=new_ages, write_ptr=new_write_ptr, ) @partial(jax.jit, static_argnums=(0,)) def update_scores(self, buffer, indices, new_scores): """Supports updating scores for specific agents in the buffer. Args: buffer: The population buffer indices: Indices of agents to update (scalar or array) new_scores: New scores for these agents (scalar or array) Returns: Updated population buffer """ updated_scores = buffer.scores.at[indices].set(new_scores) return buffer.replace(scores=updated_scores) @partial(jax.jit, static_argnums=(0, 2)) def sample_agent_indices(self, buffer, n, rng, needs_resample_mask=None): """Sample n agent indices from the buffer based on scores and staleness. Optionally takes a mask to indicate which samples are actually used, for correct age updates. Args: buffer: The population buffer n: Number of indices to sample (must be static, typically num_envs) rng: Random key needs_resample_mask: Optional boolean mask of shape (n,) indicating which sampled indices are actually used. If None, assumes all n are used. sampling_strategy: Sampling strategy to use. If None, uses the instance's sampling strategy. Returns: indices: Indices of sampled agents (shape (n,)) new_buffer: Updated buffer with correctly incremented ages """ rng, sample_rng = jax.random.split(rng, 2) buffer_has_agents = jnp.greater(buffer.filled.sum(), 0) # Select distribution based on the instance's sampling strategy sampling_dist = self._get_sampling_dist(buffer) def handle_empty_buffer(): # If buffer is empty, return random indices and unchanged buffer rand_indices = jax.random.randint(sample_rng, (n,), 0, buffer.buffer_size) return rand_indices, buffer def handle_filled_buffer(): # Always sample n indices using the chosen distribution indices = jax.random.choice( sample_rng, jnp.arange(buffer.buffer_size), shape=(n,), p=sampling_dist, replace=True ) def update_ages_plr(current_buffer, sampled_indices, needs_mask): ''' Update the ages of the buffer based on the sampled indices. ''' # Initialize reset mask for buffer entries reset_mask = jnp.zeros(current_buffer.buffer_size, dtype=bool) # Default to assuming all samples are used if mask is not provided actual_needs_mask = needs_mask if needs_mask is not None else jnp.ones(n, dtype=bool) # Loop through samples to build the reset_mask without dynamic slicing def update_reset_mask(i, current_reset_mask): buffer_idx = sampled_indices[i] is_needed = actual_needs_mask[i] new_reset_mask = jax.lax.cond( is_needed, lambda mask: mask.at[buffer_idx].set(True), lambda mask: mask, current_reset_mask ) return new_reset_mask reset_mask = jax.lax.fori_loop(0, n, update_reset_mask, reset_mask) # Compute final ages increment_mask = current_buffer.filled & (~reset_mask) new_ages = current_buffer.ages new_ages = jnp.where(increment_mask, new_ages + 1, new_ages) new_ages = jnp.where(reset_mask, 0, new_ages) return current_buffer.replace(ages=new_ages) def skip_age_update(current_buffer, sampled_indices, needs_mask): '''Do nothing, return buffer unchanged ''' return current_buffer # Apply age updates only for 'plr' strategy to skip for-i loop if possible. if self.sampling_strategy == "plr": updated_buffer = update_ages_plr(buffer, indices, needs_resample_mask) else: updated_buffer = buffer return indices, updated_buffer indices, new_buffer = jax.lax.cond( buffer_has_agents, handle_filled_buffer, handle_empty_buffer ) return indices, new_buffer def gather_agent_params(self, buffer: PopulationBuffer, agent_indices): """Gather the parameters of agents specified by agent_indices. Args: buffer: The population buffer agent_indices: Indices with shape (num_envs,), each in [0, buffer_size) Returns: Gathered parameters with shape (num_envs, ...) """ def gather_leaf(leaf): # leaf shape: (buffer_size, ...) return jax.vmap(lambda idx: leaf[idx])(agent_indices) return jax.tree_map(gather_leaf, buffer.params) def get_actions(self, buffer, agent_indices, obs, done, avail_actions, hstate, rng, env_state=None, aux_obs=None, test_mode=False): """Get actions from agents in the buffer. Callers pass obs/done/avail_actions with shape (num_envs, ...) and hstate as returned by policy.init_hstate(num_envs). For recurrent policies (RNN, S5) hstate leaves have shape (..., num_envs, H) with the env axis at -2. For MLP pass hstate=None. Args: buffer: The population buffer agent_indices: Indices with shape (num_envs,), each in [0, buffer_size) obs: Observations with shape (num_envs, obs_dim) done: Done flags with shape (num_envs,) avail_actions: Available actions with shape (num_envs, num_actions) hstate: Hidden state in policy.init_hstate(num_envs) format, or None rng: Random key env_state: Environment state with shape (num_envs, ...) or None aux_obs: Optional auxiliary vector to append to observation test_mode: Whether to use test mode (deterministic actions) Returns: actions: Actions with shape (num_envs,) new_hstate: New hidden state in the same format as hstate """ gathered_params = self.gather_agent_params(buffer, agent_indices) num_envs = agent_indices.squeeze().shape[0] rngs_batched = jax.random.split(rng, num_envs) # We vmap over the per-env params axis only, keeping the (seq_len, batch) # dims that recurrent policies (S5, RNN) require inside each forward call. # Each per-env invocation receives (1, 1, ...) shaped obs/done/avail and # a hstate with a singleton env axis (..., 1, H). def _get_action_one_env(params_i, obs_i, done_i, avail_i, hstate_i, rng_i): # obs_i: (obs_dim,) done_i: () avail_i: (num_actions,) # hstate_i: pytree with leaves (..., H) — env axis stripped by outer vmap act_i, _, _, new_h_i = self.policy_cls.get_action_value_policy( params=params_i, obs=obs_i[jnp.newaxis, jnp.newaxis, :], # (1, 1, obs_dim) done=done_i[jnp.newaxis, jnp.newaxis], # (1, 1) avail_actions=avail_i[jnp.newaxis, jnp.newaxis, :], # (1, 1, num_actions) hstate=jax.tree.map(lambda h: h[..., jnp.newaxis, :], hstate_i), rng=rng_i, ) # Squeeze (1,1) action → scalar; drop the inserted singleton env axis return act_i.squeeze(), jax.tree.map(lambda h: h[..., 0, :], new_h_i) # Move env axis to position 0 so vmap strips it per-env, then restore. if hstate is not None: hstate_env_first = jax.tree.map(lambda h: jnp.moveaxis(h, -2, 0), hstate) else: hstate_env_first = None actions, new_hstate_env_first = jax.vmap(_get_action_one_env)( gathered_params, obs, done, avail_actions, hstate_env_first, rngs_batched ) if new_hstate_env_first is not None: new_hstate = jax.tree.map(lambda h: jnp.moveaxis(h, 0, -2), new_hstate_env_first) else: new_hstate = None return actions, new_hstate