Spaces:
Sleeping
Sleeping
| 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'" | |
| 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, | |
| ) | |
| 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 | |
| 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 | |
| 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) | |
| 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}") | |
| 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] | |
| 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, | |
| ) | |
| 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) | |
| 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 |