Repository: projectmesa/mesa
Pull request #1916 (2023-12-19 10:10:43)

Introduce AgentSet class
_This PR is based on the discussion in https://github.com/projectmesa/mesa/discussions/1912#discussioncomment-7890448 for the initial discussion._

This PR introduces the `AgentSet` class in the Mesa agent-based modeling framework, along with significant changes to the agent management process. The `AgentSet` class is designed to encapsulate and manage a collection of agents, providing methods for efficient selection, sorting, shuffling, and applying actions to groups of agents. This addition aims to enhance the framework's scalability and flexibility in handling agent operations.

### Key Changes

1. **Agent Class Modifications**: The `Agent` class has been updated to manage agent registration and removal within the model's `_agents` attribute directly. This approach eliminates the need for separate registration and removal methods.
2. **Model Class Enhancements**: The `Model` class now utilizes the `AgentSet` class. The `agents` property returns an `AgentSet` instance, representing all agents in the model, thus streamlining agent access and manipulation.
3. **AgentSet Functionality**: The `AgentSet` class includes methods like `select`, `shuffle`, `sort`, and `do` for more intuitive operations on agent collections. These methods facilitate tasks like filtering agents, randomizing their order, sorting, and applying actions to each agent.
4. **Efficient Agent Management**: `AgentSet` maintains weak references to agents, enabling efficient management of agent lifecycles and garbage collection. It uses a `WeakKeyDictionary` to store agents.

### Usage examples

1. **Selecting and Filtering Agents**:
    ```python
    # Select agents with a specific attribute
    selected_agents = model.agents.select(lambda agent: agent.attribute > threshold)
    ```

2. **Shuffling Agents**:
    ```python
    # Randomize agent order
    shuffled_agents = model.agents.shuffle()
    ```

3. **Sorting Agents**:
    ```python
    # Sort agents by wealth, high to low
    sorted_agents = model.agents.sort("wealth")
    ```

4. **Applying Actions to Agents**:
    ```python
    # Apply a method to each agent
    model.agents.do('method_name')
    ```

5. **Retrieving Agent Attributes**:
    ```python
    # Get a specific attribute from all agents
    attributes = model.agents.get('attr_name')
    ```

### Advanced Chaining examples

1. **Filter, Shuffle, and Apply an Action**:
    Chaining `select`, `shuffle`, and `do` allows for filtering, randomizing, and then performing an action on the selected agents.
    ```python
    # Select agents meeting a condition, shuffle them, and then apply a method
    model.agents.select(lambda a: a.condition).shuffle().do('action_method')
    ```

2. **Sort, Select Top N, and Get Attribute**:
    This chain sorts agents, selects the top N agents based on the sort, and retrieves a specific attribute from these agents.
    ```python
    # Sort agents, select the top N, and get an attribute
    top_agents_attributes = model.agents.sort(key=lambda a: a.rank).select(n=5).get('attribute')
    ```

3. **Complex Filtering and Action Application**:
    Combining multiple filters and then applying an action. This chain can be particularly useful for complex conditional operations.
    ```python
    # Apply a method to agents that meet multiple conditions
    model.agents.select(lambda a: a.condition1).select(lambda a: a.condition2).do('method')
    ```

4. **Combined Set Operations and Action Application**:
    If set-like operations (union, intersection, difference) are implemented, they can be combined with action application.
    ```python
    # Combine two AgentSets and apply a method to the combined set
    (agentset1 + agentset2).do('collaborative_method')
    ```

5. **Shuffle, Sort, and Sequential Action Execution**:
    This chain first randomizes the order of agents, then sorts them based on an attribute, and finally applies a method sequentially.
    ```python
    # Shuffle, sort, and then apply a method in sequence
    model.agents.shuffle().sort(key=lambda a: a.priority).do('sequential_method')
    ```

6. **Nested Chain Operations**:
    Combining chaining with nested operations for more sophisticated scenarios.
    ```python
    # Nested operation: Select agents, then within that selection, shuffle and apply an action
    model.agents.select(lambda a: a.group == 'X').shuffle().do('group_specific_method')
    ```

### Replacing schedulers 
You can use the new AgentSet functionality to replacer the current schedulers and create entirely new ones.

#### 1. BaseScheduler

The `BaseScheduler` activates each agent once per step, in the order they were added.

```python
class MyModel(Model):
    def step(self):
        self.agents.do('step')
```

#### 2. RandomActivation

`RandomActivation` activates each agent once per step, in random order.

```python
class MyModel(Model):
    def step(self):
        self.agents.shuffle().do('step')
```

#### 3. SimultaneousActivation

In `SimultaneousActivation`, all agents are activated simultaneously. This can be simulated by first collecting all agents' decisions and then applying them.

```python
class MyModel(Model):
    def step(self):
        self.agents.do("step")
        self.agents.do("advance")
```

#### 4. StagedActivation

`StagedActivation` allows for multiple stages per step. For instance, agents first "move", then "eat", then "sleep".

```python
class MyModel(Model):
    def step(self):
        self.agents.do('move')
        self.agents.do('eat')
        self.agents.do('sleep')
```

#### 5. RandomActivationByType

`RandomActivationByType` activates each type of agent in random order.

```python
class MyModel(Model):
    def step(self):
        # Assuming AgentType1 and AgentType2 are subclasses of Agent
        self.agents.select(agent_type=AgentType1).shuffle().do('step')
        self.agents.select(agent_type=AgentType2).shuffle().do('step')
```

These examples demonstrate how the `AgentSet` class can be used to replicate the behavior of different schedulers in Mesa. The `AgentSet` class provides a more flexible and concise way to define the activation and interaction patterns among agents within the model.


#### Future work
- Adding more Sequence and MutableSet methods (e.g., `__reversed__`, `index`, `count`, `clear`, `pop`, `__ior__`) for enhanced performance and functionality.
- Rewriting current schedulers using AgentSets.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/projectmesa/mesa/pull/1916?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) Report
Attention: `14 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`9495a5a`)](https://app.codecov.io/gh/projectmesa/mesa/commit/9495a5acb6256104e95b02c71f6a8eecffb412f1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 77.53% compared to head [(`3bb8dde`)](https://app.codecov.io/gh/projectmesa/mesa/pull/1916?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 76.97%.

| [Files](https://app.codecov.io/gh/projectmesa/mesa/pull/1916?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) | Patch % | Lines |
|---|---|---|
| [mesa/agent.py](https://app.codecov.io/gh/projectmesa/mesa/pull/1916?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa#diff-bWVzYS9hZ2VudC5weQ==) | 56.66% | [12 Missing and 1 partial :warning: ](https://app.codecov.io/gh/projectmesa/mesa/pull/1916?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) |
| [mesa/model.py](https://app.codecov.io/gh/projectmesa/mesa/pull/1916?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa#diff-bWVzYS9tb2RlbC5weQ==) | 90.90% | [1 Missing :warning: ](https://app.codecov.io/gh/projectmesa/mesa/pull/1916?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1916      +/-   ##
==========================================
- Coverage   77.53%   76.97%   -0.57%     
==========================================
  Files          15       15              
  Lines        1015     1051      +36     
  Branches      221      228       +7     
==========================================
+ Hits          787      809      +22     
- Misses        197      210      +13     
- Partials       31       32       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/projectmesa/mesa/pull/1916?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).

EwoutH: I won’t be working on it between now and out meeting tomorrow, so take it as far as you like. Maybe there is also still some stuff from #1911 that can be ported over.
quaquel: I'll see what I can do tomorrow morning. I just addressed my comments on your first submission and added the inplace keyword argument. 
Corvince: Awesome PR
EwoutH: Thanks for working on It. I will try to do an initial review tomorrow between 9 and 10 CET, otherwise around noon (feel free to continue regardless).
EwoutH: I’m leaning towards just documenting properly that you can use `agentset.copy()` or `copy(agentset)’ instead of the whole `inplace` thing everywhere. It just seems to add unnecessary complexity to every method while `copy` is a general solution.
quaquel: I disagree with the copy thing. First, it adds another method call into any chain. Second, the API now follows Pandas, so users of Mesa are likely to be familiar with it. Third, in particular, in the case of a large number of agents, there are substantial memory advantages to using `inplace=True.`. 
EwoutH: Thanks for your perspective. Sound arguments, let me further think about it. Also curious what others think.

I would propose all the "getting data back" things (like `get_each`) to be done in a future PR. That could be done together with the redesign of the data collector (see #348).
Corvince: > Thanks for your perspective. Sound arguments, let me further think about it. Also curious what others think.
> 
> I would propose all the "getting data back" things (like `get_each`) to be done in a future PR. That could be done together with the redesign of the data collector (see #348).

I concur with keeping inplace in the API, since I think copying should be the default. We will see how that turns out performance wise, but we shouldnt worry about that just yet.

I also agree we should keep this PR small at first and add further functions later. This is a very solid basis and can/should be used for lots of things so not overcomplicate things now. 
quaquel: > I also agree we should keep this PR small at first and add further functions later. This is a very solid basis and can/should be used for lots of things so not overcomplicate things now.


Keeping the PR focused is the right way to go. However, it should also be feature-complete, and the basic structure should be solid. Otherwise, the code base keeps changing repeatedly to reflect new ideas. My suggestion would be to focus on making sure AgentSet has the following things:

- Complete behavior of an ordered set 
- inclusion of `shuffle`, `select`, and `sort`
- ensure that you can call methods and retrieve attributes on the agents in the set. So, for me, `do_each` and `get_each` are within scope, although I am unsure about their naming. It also makes sense to ensure `do_each` has a return. This makes `do_each` much more useful across the board.

What can be discussed or kept for later, in my view, are:

- support for listlike indexing. Since `AgentSet` is ordered, so, conceptually, it makes sense to be able to retrieve by index. It also would mean there is no need for a separate `AgentList` class (contra AgentPy). It is possible to add this by implementing `__getitem__`. There are, however, some tricky questions because of the use of weak references. Should the index just refer to whatever is the fifth item in the ordered set, or should it return None or raise an exception if the agent has been removed by the garbage collector?
- the sophistication of `select`. The present implementation is bare bones because it only accepts a boolean callable. So selecting the first _n_ items  is not supported through `select` at the moment. It can however be done through slicing `agents[0:5]`, if indexing is supported
- Adding methods like `to_dataframe` or `as_dataframe`. 

EwoutH: Inheriting from [`MutableSet`](https://docs.python.org/3/library/collections.abc.html#collections.abc.MutableSet) is an interesting approach. Since we're using it as a mixin, do we comply with these notes?

![Screenshot_315](https://github.com/projectmesa/mesa/assets/15776622/f84e2bf8-f41e-4742-abad-41e2eb6cbbf4)

Also remarkably, there is a [reference](https://code.activestate.com/recipes/576694/) to an `OrderedSet` recipe in the official Python docs.
EwoutH: If either of you has anything to add to the Python discussion on `OrderedSets`, please do so. I'm trying to make the case but your probably have more specialized knowledge about this:  https://discuss.python.org/t/add-orderedset-to-stdlib/12730/92
quaquel: We follow 1, but I am not sure we need to worry about 2 (it involves subset and superset comparisons as far as I can tell); 3 is likewise unnecessary.

While googling this, I came across some discussions on pickling the set. Is this something we want to support?

Corvince: > > I also agree we should keep this PR small at first and add further functions later. This is a very solid basis and can/should be used for lots of things so not overcomplicate things now.
> 
> Keeping the PR focused is the right way to go. However, it should also be feature-complete, and the basic structure should be solid. Otherwise, the code base keeps changing repeatedly to reflect new ideas. My suggestion would be to focus on making sure AgentSet has the following things:
> 
> * Complete behavior of an ordered set
> * inclusion of `shuffle`, `select`, and `sort`
> * ensure that you can call methods and retrieve attributes on the agents in the set. So, for me, `do_each` and `get_each` are within scope, although I am unsure about their naming. It also makes sense to ensure `do_each` has a return. This makes `do_each` much more useful across the board.

The question is should do_each return the return values the called function or return the AgentSet itself. The latter would allow to keep on chaining, for example:

```python
agents.do_each(step).select(n=10).do_each(move)
```

But getting the return values could also be valueable, so maybe we can also add `agents.apply(func)`, which returns whatever the function returns. Just thinking out loud here.


quaquel: Unit tests have been added, and it seems pretty complete.
 
However, there are several outstanding issues:
1. `select` is very barebones because it only takes a callable -> bool. Do we want to support other arguments, such as `n=5` for the first five items?
2. `sort` likewise depends on a callable, so you cannot do a simple sort on attributes. In my view, this is such a common use case that we should support it. It is straightforward to do. if `key` is a string, use `getattr` with the key. 
3. The current class is not pickle-able. This can be an issue if you want to run in parallel and have created AgentSets somewhere during the initialization of the model. 
4.  Naming for `do_each` invokes the given method on each agent in the set, and `get_each` retrieves the given attribute from each agent in the set. 
quaquel: > The question is should do_each return the return values the called function or return the AgentSet itself. The latter would allow to keep on chaining, for example:
> 
> ```python
> agents.do_each(step).select(n=10).do_each(move)
> ```
> 
> But getting the return values could also be valueable, so maybe we can also add `agents.apply(func)`, which returns whatever the function returns. Just thinking out loud here.

Good point that do_each will stop chaining which might be undesirable. So having one method that calls a method but returns the agents, and another which calls the method and returns the results might be a nice way to enable both use cases. We could use 'do_each`, 'apply', and 'retrieve', where the latter 2 return the value from the method call or attribute call, while do_each returns the agents itself to allow continued chaining.
Corvince: > Unit tests have been added, and it seems pretty complete.
> 
> However, there are several outstanding issues:

Thanks for putting these together, great discussion points.

> 
> 1. `select` is very barebones because it only takes a callable -> bool. Do we want to support other arguments, such as `n=5` for the first five items?

I think n would be nice and also selecting by type would be a great addition, since this is what got all this started I think

> 2. `sort` likewise depends on a callable, so you cannot do a simple sort on attributes. In my view, this is such a common use case that we should support it. It is straightforward to do. if `key` is a string, use `getattr` with the key.

I think its a bit of a personal taste, but for sorting by attributes I could imagine a dedicated `sortBy` method.

> 3. The current class is not pickle-able. This can be an issue if you want to run in parallel and have created AgentSets somewhere during the initialization of the model.

This is a downer for me, because I used to love pickling objects for varying use cases. What would be required to make it pickleable? 

> 4. Naming for `do_each` invokes the given method on each agent in the set, and `get_each` retrieves the given attribute from each agent in the set.

I think do_each is pretty good in this context, because it does exactly what the name suggests. get_each could also be called collect_attribute, but I am not sure that is better. But having suggested that a generic collect function could also be nice in a future iteration. 


quaquel: 
> > 3. The current class is not pickle-able. This can be an issue if you want to run in parallel and have created AgentSets somewhere during the initialization of the model.
> 
> This is a downer for me, because I used to love pickling objects for varying use cases. What would be required to make it pickleable?
> 

It seems that we need to add `__setstate__` and `__getstate__`. See [the docs](https://docs.python.org/3/library/pickle.html#pickling-class-instances)
EwoutH: Great progress! 
> 1. `select` is very barebones because it only takes a callable -> bool. Do we want to support other arguments, such as `n=5` for the first five items?

It’s the end, the `AgentSet` should include all functionally as proposed in #1905 and implemented in #1911. Doesn’t all have to be with `select`, but should be possible though chaining or slicing.

> 2. `sort` likewise depends on a callable, so you cannot do a simple sort on attributes. In my view, this is such a common use case that we should support it. It is straightforward to do. if `key` is a string, use `getattr` with the key.

Agreed, also aligns with my remark on 1. Should take multiple sorting criteria and directions.

> 3. The current class is not pickle-able. This can be an issue if you want to run in parallel and have created AgentSets somewhere during the initialization of the model.

Already fixed, really nice!

> 4. Naming for `do_each` invokes the given method on each agent in the set, and `get_each` retrieves the given attribute from each agent in the set.

> We could use 'do_each`, 'apply', and 'retrieve', where the latter 2 return the value from the method call or attribute call, while do_each returns the agents itself to allow continued chaining.

I would argue for only one method that asks agents to do something. Otherwise it will just confuse the hell out of everyone. That method can have an optional keyword argument with what it returns. By default I think that should be an `AgentSet`, but we might want to support other types of data (list, maybe dict).

Since we might want to “get” other things or data formats later, I don’t think the method to get an attribute from all agents should be called `get_each`. `get_attribute` or `get_attribute_list` sounds more logical.

For me this will get more clear as we’re going to work on #348.
EwoutH: Maybe we can just call `do_each` simply `do`. In the `AgentSet`, it might be implied that the methods apply on all agents. Then we can also make multiple variants of `get`.

```Python
agents.do(step).select(n=10).shuffle().do(move)
```
EwoutH: [`IndexedSet`](https://boltons.readthedocs.io/en/latest/setutils.html#boltons.setutils.IndexedSet) was suggested on discuss.python.org. Functionality wise it looks great, but I don’t know how fast it is since it looks like pure Python.
quaquel: 1. Ok, I'll have a look and make sure it all is covered
2. Ok. I do like `sort_by` as well, but given your later points, I agree that having many methods that do slightly different things is not desirable. My hunch is to just have several kwargs, each with a different clear use case.
3. Was a bit tricky because of weakrefs, but yes it works.
4. I'll probably go with `do` and `get` or `get_attribute`. For do, I'll add a kwarg to control the return.  
quaquel: I believe I have completely covered the various comments and suggestions given so far. In short, `AgentSet` can now be pickled. I have changed the method names to `do` for calling methods on agents and `get` for retrieving attributes. `do` by default returns the agentset itself to enable continued changing: `agents.shuffle().do("stage").do("execute")`. However, you can use a keyword argument to return the results from the method call instead. `get` evidently always returns the results.

`select` and `sort` have some additional optional keyword arguments. `select` now accepts `n=` to select the __n__ first items. This can be used in conjunction with a callable. `sort` now accepts a callable or str. In the case of a str, it is assumed that you want to sort on an attribute, so it gets wrapped into `operator.attrgetter`. 

Selecting by agent type can already be handled by the callable. Moreover, there is also `model.get_agents_of_type()`, which returns an AgentSet. Since `model._agents` stores by agent type, this will likely be faster. However, if there is a need or desire to make `select` work with agent types as well, it would be relatively easy to add this analogous to how sort handles the situation where `key` is a string.  
EwoutH: Okay, I think this is mature enough to get out in the wild. In my opinion we try to move fast and break things, by releasing this with an experimental warning to 2.2.0. Then we can gather feedback while we develop more advanced features on top of the `AgentSet` in the mean time, and aim to release it stable in 2.3.0 (or 3.0.0 if that's next).

What we still need to do in this PR, that I can think of:
- [ ] Make sure tests are covering all current features and use cases
- [ ] Update docstring
- [ ] Update PR description and examples
quaquel: I have added tests that cover most of the stuff. I haven't tested with-coverage yet to be completely sure, however. 
I added docstrings to the key user-facing methods. 
go ahead on the PR text. 
Corvince: > I believe I have completely covered the various comments and suggestions given so far. In short, `AgentSet` can now be pickled. I have changed the method names to `do` for calling methods on agents and `get` for retrieving attributes. `do` by default returns the agentset itself to enable continued chaining: `agents.shuffle().do("stage").do("execute")`. However, you can use a keyword argument to return the results from the method call instead. `get` evidently always returns the results.
> 
> `select` and `sort` have some additional optional keyword arguments. `select` now accepts `n=` to select the **n** first items. This can be used in conjunction with a callable. `sort` now accepts a callable or str. In the case of a str, it is assumed that you want to sort on an attribute, so it gets wrapped into `operator.attrgetter`.

That sounds pretty good to me so far. 
> 
> Selecting by agent type can already be handled by the callable. Moreover, there is also `model.get_agents_of_type()`, which returns an AgentSet. Since `model._agents` stores by agent type, this will likely be faster. However, if there is a need or desire to make `select` work with agent types as well, it would be relatively easy to add this analogous to how sort handles the situation where `key` is a string.

While it is true that it can already be done by the callable, I think a a dedicated type keyword would make sense anyway. The advantage of having this inside the agentset and not only in get_agents_of_type is that we can subselect the agentset by type and not only select all the agents of a specified type. For example if a grid returns an Agentset for the agents in a grid cell and we only want to select a specific type there.


EwoutH: Thanks for all the work @quaquel! I'm going to clean the last things up and make it ready for review.
EwoutH: Todo:
- Easy selection of agents by type in the AgentSet
- Rename `reverse` to something more descriptive like `descending`.

I think by default, sorting should happen from high to low. So `descending=True` as default.
quaquel: 1. done
2. There are two Python defaults. `reverse` is used in default Python sort functions. `ascending` is used in pandas. I deliberately went with the default Python label.
EwoutH: Sorry, was also already working on it, overwritten your changes.
EwoutH: @jackiekazil @tpike3 @rht @wang-boyu and others,

Thanks to the hard work of @quaquel and critical eye of @Corvince, I'm quite proud to announce this PR is ready for review! Please look though the code and examples and share your thoughts.

I added the experimental warning in 7888007, and it works beautifully.

![Screenshot_317](https://github.com/projectmesa/mesa/assets/15776622/a48e9234-02fd-4bc6-b0db-1e54791101d5)
Corvince: Just two more comments from my side, otherwise I think we are good to go!
quaquel: > Just two more comments from my side, otherwise I think we are good to go!

both have been done.
Corvince: I approved, but would like to have 2 more eyes on this before I feel comfortable merging. But I think @EwoutH you can start to adapt this in your other prs
EwoutH: Thanks, I agree on at least another reviewer. No other PRs I have currently open have to be modified based on this PR (we can modify the schedulers, datacollector, and maybe grid in the future).
tpike3: @EwoutH This is great!!! 

Thanks @Corvince and @quaquel this is truly an amazing effort and incredible discussion. 

Patch:
diff --git a/mesa/agent.py b/mesa/agent.py
index e29b233c..0bb2f218 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -7,10 +7,15 @@ Core Objects: Agent
 # Remove this __future__ import once the oldest supported Python is 3.10
 from __future__ import annotations
 
+import contextlib
+import operator
+import warnings
+import weakref
+from collections.abc import MutableSet, Sequence
 from random import Random
 
 # mypy
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator
 
 if TYPE_CHECKING:
     # We ensure that these are not imported during runtime to prevent cyclic
@@ -41,12 +46,13 @@ class Agent:
         self.model = model
         self.pos: Position | None = None
 
-        # Register the agent with the model using defaultdict
-        self.model.agents[type(self)][self] = None
+        # register agent
+        self.model._agents[type(self)][self] = None
 
     def remove(self) -> None:
         """Remove and delete the agent from the model."""
-        self.model.agents[type(self)].pop(self)
+        with contextlib.suppress(KeyError):
+            self.model._agents[type(self)].pop(self)
 
     def step(self) -> None:
         """A single step of the agent."""
@@ -57,3 +63,279 @@ class Agent:
     @property
     def random(self) -> Random:
         return self.model.random
+
+
+class AgentSet(MutableSet, Sequence):
+    """
+    .. warning::
+        The AgentSet is experimental. It may be changed or removed in any and all future releases, including
+        patch releases.
+        We would love to hear what you think about this new feature. If you have any thoughts, share them with
+        us here: https://github.com/projectmesa/mesa/discussions/1919
+
+    A collection class that represents an ordered set of agents within an agent-based model (ABM). This class
+    extends both MutableSet and Sequence, providing set-like functionality with order preservation and
+    sequence operations.
+
+    Attributes:
+        model (Model): The ABM model instance to which this AgentSet belongs.
+
+    Methods:
+        __len__, __iter__, __contains__, select, shuffle, sort, _update, do, get, __getitem__,
+        add, discard, remove, __getstate__, __setstate__, random
+
+    Note:
+        The AgentSet maintains weak references to agents, allowing for efficient management of agent lifecycles
+        without preventing garbage collection. It is associated with a specific model instance, enabling
+        interactions with the model's environment and other agents.The implementation uses a WeakKeyDictionary to store agents,
+        which means that agents not referenced elsewhere in the program may be automatically removed from the AgentSet.
+    """
+
+    def __init__(self, agents: Iterable[Agent], model: Model):
+        """
+        Initializes the AgentSet with a collection of agents and a reference to the model.
+
+        Args:
+            agents (Iterable[Agent]): An iterable of Agent objects to be included in the set.
+            model (Model): The ABM model instance to which this AgentSet belongs.
+        """
+        self.model = model
+
+        if not self.model.agentset_experimental_warning_given:
+            self.model.agentset_experimental_warning_given = True
+            warnings.warn(
+                "The AgentSet is experimental. It may be changed or removed in any and all future releases, including patch releases.\n"
+                "We would love to hear what you think about this new feature. If you have any thoughts, share them with us here: https://github.com/projectmesa/mesa/discussions/1919",
+                FutureWarning,
+                stacklevel=2,
+            )
+
+        self._agents = weakref.WeakKeyDictionary()
+        for agent in agents:
+            self._agents[agent] = None
+
+    def __len__(self) -> int:
+        """Return the number of agents in the AgentSet."""
+        return len(self._agents)
+
+    def __iter__(self) -> Iterator[Agent]:
+        """Provide an iterator over the agents in the AgentSet."""
+        return self._agents.keys()
+
+    def __contains__(self, agent: Agent) -> bool:
+        """Check if an agent is in the AgentSet. Can be used like `agent in agentset`."""
+        return agent in self._agents
+
+    def select(
+        self,
+        filter_func: Callable[[Agent], bool] | None = None,
+        n: int = 0,
+        inplace: bool = False,
+        agent_type: type[Agent] | None = None,
+    ) -> AgentSet:
+        """
+        Select a subset of agents from the AgentSet based on a filter function and/or quantity limit.
+
+        Args:
+            filter_func (Callable[[Agent], bool], optional): A function that takes an Agent and returns True if the
+                agent should be included in the result. Defaults to None, meaning no filtering is applied.
+            n (int, optional): The number of agents to select. If 0, all matching agents are selected. Defaults to 0.
+            inplace (bool, optional): If True, modifies the current AgentSet; otherwise, returns a new AgentSet. Defaults to False.
+            agent_type (type[Agent], optional): The class type of the agents to select. Defaults to None, meaning no type filtering is applied.
+
+        Returns:
+            AgentSet: A new AgentSet containing the selected agents, unless inplace is True, in which case the current AgentSet is updated.
+        """
+
+        def agent_generator():
+            count = 0
+            for agent in self:
+                if filter_func and not filter_func(agent):
+                    continue
+                if agent_type and not isinstance(agent, agent_type):
+                    continue
+                yield agent
+                count += 1
+                # default of n is zero, zo evaluates to False
+                if n and count >= n:
+                    break
+
+        agents = agent_generator()
+
+        return AgentSet(agents, self.model) if not inplace else self._update(agents)
+
+    def shuffle(self, inplace: bool = False) -> AgentSet:
+        """
+        Randomly shuffle the order of agents in the AgentSet.
+
+        Args:
+            inplace (bool, optional): If True, shuffles the agents in the current AgentSet; otherwise, returns a new shuffled AgentSet. Defaults to False.
+
+        Returns:
+            AgentSet: A shuffled AgentSet. Returns the current AgentSet if inplace is True.
+        """
+        shuffled_agents = list(self)
+        self.random.shuffle(shuffled_agents)
+
+        return (
+            AgentSet(shuffled_agents, self.model)
+            if not inplace
+            else self._update(shuffled_agents)
+        )
+
+    def sort(
+        self,
+        key: Callable[[Agent], Any] | str,
+        ascending: bool = False,
+        inplace: bool = False,
+    ) -> AgentSet:
+        """
+        Sort the agents in the AgentSet based on a specified attribute or custom function.
+
+        Args:
+            key (Callable[[Agent], Any] | str): A function or attribute name based on which the agents are sorted.
+            ascending (bool, optional): If True, the agents are sorted in ascending order. Defaults to False.
+            inplace (bool, optional): If True, sorts the agents in the current AgentSet; otherwise, returns a new sorted AgentSet. Defaults to False.
+
+        Returns:
+            AgentSet: A sorted AgentSet. Returns the current AgentSet if inplace is True.
+        """
+        if isinstance(key, str):
+            key = operator.attrgetter(key)
+
+        sorted_agents = sorted(self._agents.keys(), key=key, reverse=not ascending)
+
+        return (
+            AgentSet(sorted_agents, self.model)
+            if not inplace
+            else self._update(sorted_agents)
+        )
+
+    def _update(self, agents: Iterable[Agent]):
+        """Update the AgentSet with a new set of agents.
+        This is a private method primarily used internally by other methods like select, shuffle, and sort.
+        """
+        _agents = weakref.WeakKeyDictionary()
+        for agent in agents:
+            _agents[agent] = None
+
+        self._agents = _agents
+        return self
+
+    def do(
+        self, method_name: str, *args, return_results: bool = False, **kwargs
+    ) -> AgentSet | list[Any]:
+        """
+        Invoke a method on each agent in the AgentSet.
+
+        Args:
+            method_name (str): The name of the method to call on each agent.
+            return_results (bool, optional): If True, returns the results of the method calls; otherwise, returns the AgentSet itself. Defaults to False, so you can chain method calls.
+            *args: Variable length argument list passed to the method being called.
+            **kwargs: Arbitrary keyword arguments passed to the method being called.
+
+        Returns:
+            AgentSet | list[Any]: The results of the method calls if return_results is True, otherwise the AgentSet itself.
+        """
+        res = [getattr(agent, method_name)(*args, **kwargs) for agent in self._agents]
+
+        return res if return_results else self
+
+    def get(self, attr_name: str) -> list[Any]:
+        """
+        Retrieve a specified attribute from each agent in the AgentSet.
+
+        Args:
+            attr_name (str): The name of the attribute to retrieve from each agent.
+
+        Returns:
+            list[Any]: A list of attribute values from each agent in the set.
+        """
+        return [getattr(agent, attr_name) for agent in self._agents]
+
+    def __getitem__(self, item: int | slice) -> Agent:
+        """
+        Retrieve an agent or a slice of agents from the AgentSet.
+
+        Args:
+            item (int | slice): The index or slice for selecting agents.
+
+        Returns:
+            Agent | list[Agent]: The selected agent or list of agents based on the index or slice provided.
+        """
+        return list(self._agents.keys())[item]
+
+    def add(self, agent: Agent):
+        """
+        Add an agent to the AgentSet.
+
+        Args:
+            agent (Agent): The agent to add to the set.
+
+        Note:
+            This method is an implementation of the abstract method from MutableSet.
+        """
+        self._agents[agent] = None
+
+    def discard(self, agent: Agent):
+        """
+        Remove an agent from the AgentSet if it exists.
+
+        This method does not raise an error if the agent is not present.
+
+        Args:
+            agent (Agent): The agent to remove from the set.
+
+        Note:
+            This method is an implementation of the abstract method from MutableSet.
+        """
+        with contextlib.suppress(KeyError):
+            del self._agents[agent]
+
+    def remove(self, agent: Agent):
+        """
+        Remove an agent from the AgentSet.
+
+        This method raises an error if the agent is not present.
+
+        Args:
+            agent (Agent): The agent to remove from the set.
+
+        Note:
+            This method is an implementation of the abstract method from MutableSet.
+        """
+        del self._agents[agent]
+
+    def __getstate__(self):
+        """
+        Retrieve the state of the AgentSet for serialization.
+
+        Returns:
+            dict: A dictionary representing the state of the AgentSet.
+        """
+        return {"agents": list(self._agents.keys()), "model": self.model}
+
+    def __setstate__(self, state):
+        """
+        Set the state of the AgentSet during deserialization.
+
+        Args:
+            state (dict): A dictionary representing the state to restore.
+        """
+        self.model = state["model"]
+        self._update(state["agents"])
+
+    @property
+    def random(self) -> Random:
+        """
+        Provide access to the model's random number generator.
+
+        Returns:
+            Random: The random number generator associated with the model.
+        """
+        return self.model.random
+
+
+# consider adding for performance reasons
+# for Sequence: __reversed__, index, and count
+# for MutableSet clear, pop, remove, __ior__, __iand__, __ixor__, and __isub__
diff --git a/mesa/model.py b/mesa/model.py
index 96398b44..22e65448 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -7,12 +7,14 @@ Core Objects: Model
 # Remove this __future__ import once the oldest supported Python is 3.10
 from __future__ import annotations
 
+import itertools
 import random
 from collections import defaultdict
 
 # mypy
 from typing import Any
 
+from mesa.agent import AgentSet
 from mesa.datacollection import DataCollector
 
 
@@ -27,8 +29,20 @@ class Model:
         running: A boolean indicating if the model should continue running.
         schedule: An object to manage the order and execution of agent steps.
         current_id: A counter for assigning unique IDs to agents.
-        agents: A defaultdict mapping each agent type to a dict of its instances.
-                Agent instances are saved in the nested dict keys, with the values being None.
+        _agents: A defaultdict mapping each agent type to a dict of its instances.
+                 This private attribute is used internally to manage agents.
+
+    Properties:
+        agents: An AgentSet containing all agents in the model, generated from the _agents attribute.
+        agent_types: A list of different agent types present in the model.
+
+    Methods:
+        get_agents_of_type: Returns an AgentSet of agents of the specified type.
+        run_model: Runs the model's simulation until a defined end condition is reached.
+        step: Executes a single step of the model's simulation process.
+        next_id: Generates and returns the next unique identifier for an agent.
+        reset_randomizer: Resets the model's random number generator with a new or existing seed.
+        initialize_data_collector: Sets up the data collector for the model, requiring an initialized scheduler and agents.
     """
 
     def __new__(cls, *args: Any, **kwargs: Any) -> Any:
@@ -51,12 +65,27 @@ class Model:
         self.running = True
         self.schedule = None
         self.current_id = 0
-        self.agents: defaultdict[type, dict] = defaultdict(dict)
+        self._agents: defaultdict[type, dict] = defaultdict(dict)
+
+        # Warning flags for current experimental features. These make sure a warning is only printed once per model.
+        self.agentset_experimental_warning_given = False
 
     @property
-    def agent_types(self) -> list:
+    def agents(self) -> AgentSet:
+        """Provides an AgentSet of all agents in the model, combining agents from all types."""
+        all_agents = itertools.chain(
+            *(agents_by_type.keys() for agents_by_type in self._agents.values())
+        )
+        return AgentSet(all_agents, self)
+
+    @property
+    def agent_types(self) -> list[type]:
         """Return a list of different agent types."""
-        return list(self.agents.keys())
+        return list(self._agents.keys())
+
+    def get_agents_of_type(self, agenttype: type) -> AgentSet:
+        """Retrieves an AgentSet containing all agents of the specified type."""
+        return AgentSet(self._agents[agenttype].values(), self)
 
     def run_model(self) -> None:
         """Run the model until the end condition is reached. Overload as


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 561208f2..3cc8f9ad 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -1,16 +1,208 @@
-from mesa.agent import Agent
+import pickle
+
+import pytest
+
+from mesa.agent import Agent, AgentSet
 from mesa.model import Model
 
 
-def test_agent_removal():
-    class TestAgent(Agent):
-        pass
+class TestAgent(Agent):
+    def get_unique_identifier(self):
+        return self.unique_id
 
+
+def test_agent_removal():
     model = Model()
     agent = TestAgent(model.next_id(), model)
     # Check if the agent is added
-    assert agent in model.agents[type(agent)]
+    assert agent in model.agents
 
     agent.remove()
     # Check if the agent is removed
-    assert agent not in model.agents[type(agent)]
+    assert agent not in model.agents
+
+
+def test_agentset():
+    # create agentset
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+
+    agentset = AgentSet(agents, model)
+
+    assert agents[0] in agentset
+    assert len(agentset) == len(agents)
+    assert all(a1 == a2 for a1, a2 in zip(agentset[0:5], agents[0:5]))
+
+    for a1, a2 in zip(agentset, agents):
+        assert a1 == a2
+
+    def test_function(agent):
+        return agent.unique_id > 5
+
+    assert len(agentset.select(test_function)) == 5
+    assert len(agentset.select(test_function, n=2)) == 2
+    assert len(agentset.select(test_function, inplace=True)) == 5
+    assert agentset.select(inplace=True) == agentset
+    assert all(a1 == a2 for a1, a2 in zip(agentset.select(), agentset))
+    assert all(a1 == a2 for a1, a2 in zip(agentset.select(n=5), agentset[:5]))
+
+    assert len(agentset.shuffle().select(n=5)) == 5
+
+    def test_function(agent):
+        return agent.unique_id
+
+    assert all(
+        a1 == a2
+        for a1, a2 in zip(agentset.sort(test_function, ascending=False), agentset[::-1])
+    )
+    assert all(
+        a1 == a2
+        for a1, a2 in zip(agentset.sort("unique_id", ascending=False), agentset[::-1])
+    )
+
+    assert all(
+        a1 == a2.unique_id for a1, a2 in zip(agentset.get("unique_id"), agentset)
+    )
+    assert all(
+        a1 == a2.unique_id
+        for a1, a2 in zip(
+            agentset.do("get_unique_identifier", return_results=True), agentset
+        )
+    )
+    assert agentset == agentset.do("get_unique_identifier")
+
+    agentset.discard(agents[0])
+    assert agents[0] not in agentset
+    agentset.discard(agents[0])  # check if no error is raised on discard
+
+    with pytest.raises(KeyError):
+        agentset.remove(agents[0])
+
+    agentset.add(agents[0])
+    assert agents[0] in agentset
+
+    # because AgentSet uses weakrefs, we need hard refs as well....
+    other_agents, another_set = pickle.loads(  # noqa: S301
+        pickle.dumps([agents, AgentSet(agents, model)])
+    )
+    assert all(
+        a1.unique_id == a2.unique_id for a1, a2 in zip(another_set, other_agents)
+    )
+    assert len(another_set) == len(other_agents)
+
+
+def test_agentset_initialization():
+    model = Model()
+    empty_agentset = AgentSet([], model)
+    assert len(empty_agentset) == 0
+
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+    assert len(agentset) == 10
+
+
+def test_agentset_serialization():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(5)]
+    agentset = AgentSet(agents, model)
+
+    serialized = pickle.dumps(agentset)
+    deserialized = pickle.loads(serialized)  # noqa: S301
+
+    original_ids = [agent.unique_id for agent in agents]
+    deserialized_ids = [agent.unique_id for agent in deserialized]
+
+    assert deserialized_ids == original_ids
+
+
+def test_agent_membership():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(5)]
+    agentset = AgentSet(agents, model)
+
+    assert agents[0] in agentset
+    assert TestAgent(model.next_id(), model) not in agentset
+
+
+def test_agent_add_remove_discard():
+    model = Model()
+    agent = TestAgent(model.next_id(), model)
+    agentset = AgentSet([], model)
+
+    agentset.add(agent)
+    assert agent in agentset
+
+    agentset.remove(agent)
+    assert agent not in agentset
+
+    agentset.add(agent)
+    agentset.discard(agent)
+    assert agent not in agentset
+
+    with pytest.raises(KeyError):
+        agentset.remove(agent)
+
+
+def test_agentset_get_item():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+
+    assert agentset[0] == agents[0]
+    assert agentset[-1] == agents[-1]
+    assert agentset[1:3] == agents[1:3]
+
+    with pytest.raises(IndexError):
+        _ = agentset[20]
+
+
+def test_agentset_do_method():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+
+    with pytest.raises(AttributeError):
+        agentset.do("non_existing_method")
+
+
+def test_agentset_get_attribute():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+
+    unique_ids = agentset.get("unique_id")
+    assert unique_ids == [agent.unique_id for agent in agents]
+
+    with pytest.raises(AttributeError):
+        agentset.get("non_existing_attribute")
+
+
+class OtherAgentType(Agent):
+    def get_unique_identifier(self):
+        return self.unique_id
+
+
+def test_agentset_select_by_type():
+    model = Model()
+    # Create a mix of agents of two different types
+    test_agents = [TestAgent(model.next_id(), model) for _ in range(4)]
+    other_agents = [OtherAgentType(model.next_id(), model) for _ in range(6)]
+
+    # Combine the two types of agents
+    mixed_agents = test_agents + other_agents
+    agentset = AgentSet(mixed_agents, model)
+
+    # Test selection by type
+    selected_test_agents = agentset.select(agent_type=TestAgent)
+    assert len(selected_test_agents) == len(test_agents)
+    assert all(isinstance(agent, TestAgent) for agent in selected_test_agents)
+    assert len(selected_test_agents) == 4
+
+    selected_other_agents = agentset.select(agent_type=OtherAgentType)
+    assert len(selected_other_agents) == len(other_agents)
+    assert all(isinstance(agent, OtherAgentType) for agent in selected_other_agents)
+    assert len(selected_other_agents) == 6
+
+    # Test with no type specified (should select all agents)
+    all_agents = agentset.select()
+    assert len(all_agents) == len(mixed_agents)
diff --git a/tests/test_model.py b/tests/test_model.py
index c54634d8..874d45f9 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -49,5 +49,5 @@ def test_agent_types():
 
     model = Model()
     test_agent = TestAgent(model.next_id(), model)
-    assert test_agent in model.agents[type(test_agent)]
+    assert test_agent in model.agents
     assert type(test_agent) in model.agent_types