Repository: projectmesa/mesa
Pull request #1894 (2023-11-30 11:56:43)

Native support for multiple agent types
_On 2023-12-17 this PR was fully updated based on all feedback._

This PR adds a new `agents` dictionary to the Mesa `Model` class, enabling native support for handling multiple agent types within models. This way all modules can know which agents and agents type are in the model at any given time, by calling `model.agents`.

#### Motivation
NetLogo has had agent types, called [`breeds`](https://ccl.northwestern.edu/netlogo/docs/dict/breed.html), built-in from the start. It works perfectly in all NetLogo components, because it's a first class citizen and all components need to be designed to consider different breeds.

In Mesa, agent types are an afterthought at best. Almost nothing is currently designed with multiple agent types in mind. That has caused several issues and limitations over the years, including:

- https://github.com/projectmesa/mesa/pull/348
- https://github.com/projectmesa/mesa/pull/1142
- https://github.com/projectmesa/mesa/pull/1162

Especially in scheduling, space and datacollection, lack of a native, consistent construct for agent types severely limits the possibilities. With the discussion about patches and "empty" this discussion done again. You might want empty to refer to all agents or only a subset of types or single type. That's currently cumbersome to implement.

Basically, by always having dictionary available of which agents of which types are in the model, you can always trust on a consistent construct to iterate over agents and agent types.

#### Conceptualisation
@tpike3 has made a really nice conceptual overview how Mesa now looks with this construct:

![Screenshot_312](https://github.com/projectmesa/mesa/assets/15776622/a72417ff-c401-4cfd-9af3-b8e471e4fd51)

#### Implementation Details
- The `Model` class now uses a `defaultdict(list)` to manage agents. This eliminates the need for explicit checks when adding new agent types.
- Updated the `Agent` class to automatically register itself with the model upon creation, leveraging the `defaultdict` functionality.
- Modified the `remove` method in the `Agent` class, ensuring safe removal of agents.

#### Advantages
In the agent initialization, the agent gets added to the Model's `agents` dictionary. Thereby, you can always, from everywhere:
- Acces the available agent types (using the `model.agent_types` property)
- Iterate over the different agent types, by:
  ```Python
  for agent_type, agents in self.agents.items():
      yield agent_type, agents
  ```
- Get all the agents of a certain type by `model.agents_by_type[AgentType]`

This will make proper multiple agent support in all Mesa components (the scheduler, grid, datacollector and visualisation) easier, because there is now a consistent model variable they can access.

#### Backwards compatibility
This PR itself breaks nothing, it just adds a dictionary to the model.

If you have a single agent, looping over the agents will just mean that only one agent type is iterated over (since there is always at least one agent type). So this work well with both model with a single and multiple agent types.

#### Future implications
While this PR itself won't break current models, if we modify the scheduler, datacollector or grids to use this construct, it might start breaking models. By using weak references ([`weakref`](https://docs.python.org/3/library/weakref.html)), we might be able to limit backwards breaking changes.

#### Review
Since this is a very important change to the core of Mesa, I recommend really critically reviewing this PR, both conceptually (what are the implications for users) and implementation wise (is this the most efficient, scalable and flexible solution).

Any feedback is appreciated!

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/projectmesa/mesa/pull/1894?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) Report
Attention: `6 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`76229a5`)](https://app.codecov.io/gh/projectmesa/mesa/commit/76229a5424e473b1ca0d08124a84fe7a4a65f526?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 77.35% compared to head [(`074024d`)](https://app.codecov.io/gh/projectmesa/mesa/pull/1894?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 76.96%.

| [Files](https://app.codecov.io/gh/projectmesa/mesa/pull/1894?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) | Patch % | Lines |
|---|---|---|
| [mesa/model.py](https://app.codecov.io/gh/projectmesa/mesa/pull/1894?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa#diff-bWVzYS9tb2RlbC5weQ==) | 25.00% | [6 Missing :warning: ](https://app.codecov.io/gh/projectmesa/mesa/pull/1894?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    #1894      +/-   ##
==========================================
- Coverage   77.35%   76.96%   -0.39%     
==========================================
  Files          15       15              
  Lines        1007     1016       +9     
  Branches      220      221       +1     
==========================================
+ Hits          779      782       +3     
- Misses        197      203       +6     
  Partials       31       31              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/projectmesa/mesa/pull/1894?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).

rht: I'd say a less drastic change would be to modify the schedulers. Currently, there needs to be a permutation of [random activation, staged activation, simultaneous activation, discrete event, ...] x [1 agent, multiple agent]. This needs to be reduced to 1D of [random activation, staged activation, simultaneous activation, discrete event, ...] that is multiple type by default.

There are 2 problems:
- if you specify `agents_by_type` in the model, the implementation-specific collection container will be tied to the model class and can't be changed. There could be a fast backend typed dict implementation via Numba/Cython/etc, and if you solder in the `agents_by_type`, it is no longer modular.
- the agents container is replicated 3x in 3 objects: the model, scheduler, and space.
rht: > the agents container is replicated 3x in 3 objects: the model, scheduler, and space.

This is not the case given that you said the scheduler will access `model.agents_by_type` directly, so that it is only 2x.
rht: > There could be a fast backend typed dict implementation via Numba/Cython/etc, and if you solder in the agents_by_type, it is no longer modular.

Actually, it could still be modular even in the model class. There could be an argument `container_backend` that could be `default`, `numba_typeddict`, etc.

Note this all can be implemented without breaking compatibility by modifying the schedulers to by multiple type by default, so can happen without having to wait until Mesa 3.0. The remaining issue of whether to put the `agents_by_type` in the model or in the scheduler can be decided later.
rht: Another problem is that this is limited to grouping based on agent classes. Sometimes you want to "groupby" agents with various other metrics (inside/outside a bar in the El Farol model, the condition in Epstein civil violence model, or solvent/insolvent bank). This is the same issue I raised in the multi-agent data collection thread https://github.com/projectmesa/mesa/issues/348#issuecomment-1779050949. See also https://github.com/tpike3/multilevel_mesa.
EwoutH: Thanks for the extensive replies!

> if you specify agents_by_type in the model, the implementation-specific collection container will be tied to the model class and can't be changed. There could be a fast backend typed dict implementation via Numba/Cython/etc, and if you solder in the agents_by_type, it is no longer modular.

Sorry, I don’t understand the problem fully. Could you maybe give an example?

> This is not the case given that you said the scheduler will access model.agents_by_type directly, so that it is only 2x.

Is there a fundamental reason we couldn’t keep everything in the model and let both the scheduler and the grid access if from there.

> The remaining issue of whether to put the agents_by_type in the model or in the scheduler can be decided later.

For me, model feels like the right place. It’s the big overarching construct, from which other modules can access.

> Another problem is that this is limited to grouping based on agent classes. Sometimes you want to "groupby" agents with various other metrics (inside/outside a bar in the El Farol model, the condition in Epstein civil violence model, or solvent/insolvent bank).

I think you’re spot on here: this is another problem I’m suggesting native support for different kind (classes) of agents throughout the model, like NetLogo has breeds. You’re suggesting a universal way to filter agents on current state values. State values are parametric differences, different classes structural. So I think that are two different problems.
rht: I suppose I need to describe some concrete examples. I replaced `model.agents_by_type` with `model.agents` so that it is more concise.

```python
# staged activation by type and condition
# Both wolf and sheep
model.agents.do_each("move")
model.agents.get_by(class=Sheep).do_each("eat_or_die")
model.agents.get_by(class=Wolf).do_each("hunt_or_die")
# Both wolf and sheep
model.agents.get_by(condition=lambda a: a.check_is_ready()).do_each("reproduce")

# RiskAverseTrader and AdaptiveTrader
# This is a string because a group could be a class that both classes inherit from, or for some conditional.
# The group argument is for when it is final, i.e. unchanged after the model initialization. For dynamic grouping, use the condition argument.
model.agents.get_by(group="trader").do_each("trade")
```

- It's a bit more verbose than [NetLogo's](https://github.com/NetLogo/models/blob/master/Sample%20Models/Biology/Wolf%20Sheep%20Predation.nlogo).
- It's not declarative, because I suppose the scheduler is meant to describe chronology of actions.
EwoutH: This is a great example and I really like the cleanness. It reads well from left to right, like proper Python.

It's also nice to see the step not being one solid thing, but being modular and flexible. The Staged and ByType schedulers are also not that necessary anymore, and could maybe be replaced.

I have another idea, I will try to work it out tomorrow. But I already like this discussion and talking about the different options!
tpike3: > I started typing out this long detailed reply, only come to the conclusion that I basically ended up with the same solution you did.
> 
> * `model.agents` solves the problem to access all agents from everywhere
> * `get_by()` solves the problem to easily and quickly select a subset of agents
> 
> I have still a few questions however:
> 
> 1. What is `agents` exactly? Still a dict? Or a dataframe?
> 2. How does `get_by()` work on `agents`? Is it a method of it? (I imagined get_by as a model method, returning a set of agents)
> 3. What arguments does `do_each()` take? An agent or model function? As string? And probably something to randomize the order?
> 
> I would really like to see an implementation for your approach, since I think it's the right direction! We might still encounter some minor things but generally I really like it.

@EwoutH So this would be my take based on this conversation

Agents are python objects that are stored in a dictionary in [`mesa.time.py`](https://github.com/projectmesa/mesa/blob/main/mesa/time.py). Type only exists right now in [`RandomActivationByType`](https://github.com/projectmesa/mesa/blob/76229a5424e473b1ca0d08124a84fe7a4a65f526/mesa/time.py#L269) to do `get_by()` should be fairly trivial to reference the dictionary key which will then return a list of those agent objects.*

`do_each` would then reference the function in that agent type object, So you would [randomize the order of agents in the `list` which is already part of  `RandomActivationByType`](https://github.com/projectmesa/mesa/blob/76229a5424e473b1ca0d08124a84fe7a4a65f526/mesa/time.py#L291) returned when you called `get_by()` and then iterate through the list executing the function passed in via `do each()`

*This leads to the question of whether the agents by type should be more basic (e.g. part of the `time.py` base scheduler), but code a little test a little :)
 

EwoutH: @tpike3 looking forward to your new conceptualisation!

AgentPy also uses a construct like this in almost all their [example models](https://agentpy.readthedocs.io/en/latest/model_library.html):

```Python
class VirusModel(ap.Model):
    def setup(self):
        # Create agents
        self.agents = ap.AgentList(self, self.p.population, Person)
```
```Python
class WealthModel(ap.Model):
    def setup(self):
        # Create agents
        self.agents = ap.AgentList(self, self.p.agents, WealthAgent)
```
So the concept of having one single variable that contains all the agents - in our case stored by type - in not new in Python ABM modelling.
EwoutH: **Todo:**
- [ ] Write tests for `create_agent()` method
- [ ] Update Boltzmann Wealth Model / tutorial
EwoutH: @quaquel thanks for you feedback. I looked at it, and have the following thoughts:

1. Using class names (`.__name__`) instead of direct class reference: I think going with class references is more logical in Mesa. We might want to filter, select, check or do other things depending on the type. Having the class reference themselves is then more flexible, it allows direct method access, inspection and type checking.
    - As example, we might have a class `Predator`, which is a subclass of `Animal`. If we want to filter all animals, we can easily do that with the class reference (with `isinstance()`), but not with the class name.
2. Using the `__subclasses__` dunder method: I didn't know this existed and this is very smart! The only concern I have with using that `__subclasses__()` method in Python does not track instantiation; it only tracks class definitions. So if we have some definitions but don't use them in a model run, there might be classes where there are no instances in the model of, leading to unexpected behavior.
    ```Python
    class Model:
        @property
        def agent_types(self):
            return Agent.__subclasses__()
    ```
    Considering that, an variable that's updated by the `create_agent()` and `remove_agent()` methods might be more fitting.
3. Specifying `type` as Property in the Agent Class instead of assigning it on init: I see the advantages of having it be a property and it's probably best practice. My only concern is the readability for less experienced programmers. I don't think I've ever encountered a model where the Agent class is changed during the model run, because that's the only case I can think of where it would be beneficial to define type as a property.
    ```Python
    class Agent:
        @property
        def type(self):
            return self.__class__
    ```
4. Using Weak References is a really interesting idea. It would looks something like this right?
    ```Python
    class Model:
        def __init__(self, *args: Any, **kwargs: Any) -> None:
            self.agents_by_type = weakref.WeakValueDictionary()
    
        def create_agent(self, agent_type: type, *args, **kwargs):
            """Creates an agent of the given type with a unique ID."""
            agent = agent_type(self.next_id(), self, *args, **kwargs)
            self.agents_by_type.setdefault(agent.type, weakref.WeakSet()).add(agent)
            return agent
    ```
    It would provide benefits in automatic cleanup and provide a backup if agents are removed using something else than the `remove_agent()` method, preventing memory leaks. However, I think the model would be the only place that they are explicitly referenced, assuming we modify the schedulers to use the model's `agents_by_type` in the future. If that only reference is a weak reference, it might . Considering that, I'm slightly leaning towards a robust `remove_agent()` method and keeping it a regular dict. That would be in line with the concept that in the end, the Model knows what Agents are part of it.
      - It might be interesting to use weak references in the scheduler in the future however!

Thanks again for you suggestions, I'm learning a lot diving deeper in those Python constructs and thinking critically about how they would apply in our situation. So if you have the time, keep it coming!

<hr>

@rht, @Corvince and @wang-boy, I'm also really curious on your takes on this!
quaquel: 

Just some thoughts without diving too deeply into MESA itself but coming at it as a generic devs/abm modeler with ample Python experience who also happens to teach ABM with MESA to MSc. students (and has build a pPython connector to NetLogo without ever having built a single NetLogo model). 

I agree that having an agent_by_type overview at the model level is a good idea.

Similarly, keeping track of agents within the model should not be a concern of the user. This means using weakrefs everywhere (so in all schedulers and spaces) except in one place. You need to use hard refs in one place because, otherwise, agent instances are garbage collected once all references have been removed.`Model.agent_by_type` might be an ideal place for this one hard reference. Using weakrefs means no endless bookkeeping and removes possible locations for memory leaks.

What I don't like is burdening the user with additional api calls when creating and killing agent instances (i.e., `create_agent` and `remove_agent`). In particular, since all of this can, in principle, be handled within the agent class. As long as the user subclasses this and calls super in the `__init__`, the rest should just work. This means probably tying into some of Python's dunder methods. For example, `__new__` and `__subclasses__` can play a role in this. 

When creating a new Agent instance, this should handle the creation of the identifier (do you even really need this....?) as well as adding the instance to `Model.agent_by_type`. If you want to remove an agent, you probably need a kill method on the agent class. This kill method should then remove the agent from `Model.agent_by_type`. This means that you have to think carefully about the data structure for `Model.agent_by_type`. It seems like `defauldict(set)` might be the easiest solution, although some benchmarking is probably in order. The reason for a kill method on the agent is that there is no Python equivalent for `__new__` or `__init__`. The destructor dunder method is `__del__`, but this is only called once all references to the instance have been removed. 

As an aside, I think looking at NetLogo for how you want to handle anything is strange. NetLogo is based on a now 50-year-old language. Why confine yourself to choices made 50 years ago if you can access a much more modern language with very sophisticated OO features?
EwoutH: Right, I forgot that Agents have access to their model instance from very start, since they’re initiated with it. I’m my head is was a problem, because I thought some things needed to happen from the model scope (so in a model method like `create_agent`), but since the agent knows it’s model, it can just happen from the agent `__init__`.

Thanks, I’m going to try to move the create and kill logic into the agent.
EwoutH: I reimplemented this PR fully in c29ab08 based on the discussions and updated the PR description. In short:
- I removed the `create_agent()` model method, that's now handled in the Agent init
- I replaced and renamed the `agent_by_type` dict to a `agents` defaultdict
- I added a `remove()` method to the Agent
- I use `type(Agent)` instead of `Agent.__class__` (its more pythonic)
- I added a `agent_types` property that gives a list of all agent classes in the model

With these changes, there is now no backwards incompatible behavior and users don't need to modify existing models.
EwoutH: **Final review checklist**
 - [x] PR description is updated including conceptual model
 - [x] Docstring is updated
 - [x] Tests are added
 - [x] Commits are cleaned up:
   - b8b490c contains the feature changes
   - 6b849da fixes old test
   - 542a9b8 updates the model docstring
   - a74b120 adds tests for new functionality
 - [x] CI passes
 - [x] Code coverage is increased

@tpike3 and @jackiekazil please do final review. Let's get this merged!
Corvince: I am sorry for putting on brakes on this, but I think one detail needs to be changed. 

This stores agents in sets. While this appears to be the right datatype it has one big disadvantage: In Python the order of items in a set is non-deterministic. So in contrast to dictionaries they don't preserve insertion order. 

Since mesa is used for research we shouldn't introduce non-deterministic behavior to allow completely reproducible results. 

Since dictionaries work very similar, you can probably (ab)use dictionary keys instead of using a set.

Otherwise this really looks very clean and since it has no breaking changes very mergeable. 
EwoutH: > In Python the order of items in a set is non-deterministic. So in contrast to dictionaries they don't preserve insertion order.

Fair point. Changed from a defaultdict of sets to a defaultdict of lists (see a8bb15f).
quaquel: You are taking a massive performance hit by going to lists when you want to remove agents. 

I also fail to understand why, for research purposes, the order needs to be deterministic. The primary purpose of this is to keep track of agents per type. If you want deterministic behavior, you can always do a sort on ID and return a list. This would solve the performance loss while also ensuring deterministic behavior. Also, there are more sophisticated data structures that could be created that have the retrieval speed for removal of a set while providing deterministic behavior (e.g., ordered sets or some more involved linked list kind of thing).

So the question becomes: what is done more frequently: removing agents from the dict or retrieving the collection of agents for a specific agent type? 
EwoutH: I did a benchmark to see the impact. Since most models never remove agents, it tool the [Wolf-Sheep model](https://github.com/projectmesa/mesa-examples/tree/main/examples/wolf_sheep) from mesa-examples and ran 5 times with the default settings, 125 times. Here are the results:

![benchmark_results](https://github.com/projectmesa/mesa/assets/15776622/fa0b7a0e-195d-410b-89fd-10913f390be8)

The t-test is insignificant.

```Python
from scipy import stats
stats.ttest_ind(list_results, set_results)

TtestResult(statistic=-1.2073182676004544, pvalue=0.22845980311728672, df=248.0)
```
[mesa_with_wolfsheep_benchmark.zip](https://github.com/projectmesa/mesa/files/13701384/mesa_with_wolfsheep_benchmark.zip)

Given that with a model that frequently deletes agents, the performance impact is unnoticeable, I would say lets go with the lists.

If we find some weird edge-case later we can always patch it. Perfect is the enemy of done. 
EwoutH: @quaquel did some benchmarks (thanks!):

```Python
self.agents = defaultdict(list)
self.agents[type(agent)].append(agent)
```
Performance for adding and removing 1000 agents: 2.06 ms ± 3.11 µs

```Python
self.agents = defaultdict(set)
self.agents[type(agent)].add(agent)
```
Performance for adding and removing 1000 agents: 439 µs ± 1.3 µs

```Python
self.agents = defaultdict(dict)
self.agents[type(agent)][agent] = None
```
Performance for adding and removing 1000 agents: 460 µs ± 2.11 µs

All options are imperfect:
- Lists are slow (and allow duplicates)
- Sets don't keep insertion order (making runs non-deterministic)
- dicts are ugly in usage (due to their redundant `None` values)

What we need in an OrderedSet: https://discuss.python.org/t/add-orderedset-to-stdlib/12730/85

Unfortunately, that doesn't exist in the standard library (yet).

There are some packages, the most popular being ordered-set: https://pypi.org/project/ordered-set/

```Python
from ordered_set import OrderedSet

self.agents: defaultdict[type, OrderedSet] = defaultdict(OrderedSet)
```
But we have to add that to the requirements.

Personally, I'm happy with:
- Using lists
- Using the OrderedSet from ordered-set

The dict with all values None I find simply too ugly from and UX point of view. Teaching Mesa is complicated enough as it's already. You wouldn't have to specify `.keys()` everytime you want some set of agents.
EwoutH: Benchmarked ordered-set, it's extremely slow (22.9 ms ± 1.09 ms, written in pure Python).

There's another package called [Ordered-set-37](https://pypi.org/project/Ordered-set-37/) and it's a lot faster, almost as fast as the other set/dict options (615 µs ± 2.34 µs).

So, let's vote:
- 🚀: use `defaultdict(dict) = None`
- 🎉: use `defaultdict(OrderedSet)`
- ❤️: use `defaultdict(list)`
Corvince: Thanks for the performance benchmarks! I think for large lists, list should even perform much worse. 

I voted for dict. It's a bit ugly, but it's also what happens for Ordered Set 37 and I think explicitly handling it in mesa instead of a hidden in a dependency is more explicit
tpike3: > ```python
> self.agents = defaultdict(dict)
> self.agents[type(agent)][agent] = None
> ```

This is really great @EwoutH! Substantive discussion and forward progress. I am very much leaning to option 1 but instead of the `None` would be possible to do 

 self.agents[type(agent)][agent.unique_id] = agent`
 
 This will get rid of the `None` values and although you will have things like `.keys()` I think that can also be a plus as it will give users more flexibility.  In addition I would not be opposed to adding those simple functions (e.g. `.keys()`, `.values()` to `model` so users can get lists of agent ids or agent objects in an effort to mitigate your UX concerns. Let me know what you think
tpike3: > Thanks for the performance benchmarks! I think for large lists, list should even perform much worse.
> 
> I voted for dict. It's a bit ugly, but it's also what happens for Ordered Set 37 and I think explicitly handling it in mesa instead of a hidden in a dependency is more explicit

I agree and tried to think of a middle ground, let me know what you think
Corvince: Ah I didn't catch the last part of your comment. Dict should only be used internally and we should return a list via a getter. Turning the keys into a list should be negligible in terms of performance 
EwoutH: Thanks for both your responses. My biggest concern is that with the dict you can't simply do
```Python
model.agents[SomeType]
```
anlymore. Now you have to add `.keys()`, and that's just ugly. I already hear a hunderd students ask "What does the `.keys()`?" do and then I have to explain it's a stupid implementation detail.

You can indeed solve that by hiding everything behind utility function, and while I'm a big fan of utility function, in this case I think it just shouldn't be needed. It doesn't follow the spirit of [PEP 20](https://peps.python.org/pep-0020/).

> Beautiful is better than ugly.
> Explicit is better than implicit.
> Simple is better than complex.
> Flat is better than nested.
> Readability counts.
> Special cases aren't special enough to break the rules.

Having agents be a dict is just so much more explainable than having it be a nested dict, even if the values are just `None`.
tpike3: I registered my vote, but this is a great discussion and I am excited to see how it ends up. My two cents for what is worth are very much in line with @Corvince 

- I do not think we should import the Ordered Set 37 when we can just handle it natively
- I like the performance of the dictionaries particularly as I saw someone run a model with Mesa last November with 22,921,301 nodes (https://github.com/njiang8/Shareable-Reusable-Synthetic-Population-AND-Three-Agent-Based-Models )

However, I do appreciate @EwoutH  concerns having my own students struggle with dictionaries. Regardless, great job @EwoutH however this ends up it is going to be an awesome move forward.
EwoutH: Okay, it seems the majority is for a `dict` with `None` values. Let's do that.

I propose to handle all agent selection utility in the proposed `select_agents()` method from #1905.
EwoutH: Implemented in 2154cb1. Let's get this merged.

Patch:
diff --git a/mesa/agent.py b/mesa/agent.py
index 5a8966cf..e29b233c 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -20,19 +20,34 @@ if TYPE_CHECKING:
 
 
 class Agent:
-    """Base class for a model agent."""
+    """
+    Base class for a model agent in Mesa.
+
+    Attributes:
+        unique_id (int): A unique identifier for this agent.
+        model (Model): A reference to the model instance.
+        self.pos: Position | None = None
+    """
 
     def __init__(self, unique_id: int, model: Model) -> None:
-        """Create a new agent.
+        """
+        Create a new agent.
 
         Args:
-            unique_id (int): A unique numeric identified for the agent
-            model: (Model): Instance of the model that contains the agent
+            unique_id (int): A unique identifier for this agent.
+            model (Model): The model instance in which the agent exists.
         """
         self.unique_id = unique_id
         self.model = model
         self.pos: Position | None = None
 
+        # Register the agent with the model using defaultdict
+        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)
+
     def step(self) -> None:
         """A single step of the agent."""
 
diff --git a/mesa/model.py b/mesa/model.py
index 1ac0edd6..96398b44 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -8,6 +8,7 @@ Core Objects: Model
 from __future__ import annotations
 
 import random
+from collections import defaultdict
 
 # mypy
 from typing import Any
@@ -16,7 +17,19 @@ from mesa.datacollection import DataCollector
 
 
 class Model:
-    """Base class for models."""
+    """Base class for models in the Mesa ABM library.
+
+    This class serves as a foundational structure for creating agent-based models.
+    It includes the basic attributes and methods necessary for initializing and
+    running a simulation model.
+
+    Attributes:
+        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.
+    """
 
     def __new__(cls, *args: Any, **kwargs: Any) -> Any:
         """Create a new model object and instantiate its RNG automatically."""
@@ -31,16 +44,19 @@ class Model:
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Create a new model. Overload this method with the actual code to
-        start the model.
-
-        Attributes:
-            schedule: schedule object
-            running: a bool indicating if the model should continue running
+        start the model. Always start with super().__init__() to initialize the
+        model object properly.
         """
 
         self.running = True
         self.schedule = None
         self.current_id = 0
+        self.agents: defaultdict[type, dict] = defaultdict(dict)
+
+    @property
+    def agent_types(self) -> list:
+        """Return a list of different agent types."""
+        return list(self.agents.keys())
 
     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
new file mode 100644
index 00000000..561208f2
--- /dev/null
+++ b/tests/test_agent.py
@@ -0,0 +1,16 @@
+from mesa.agent import Agent
+from mesa.model import Model
+
+
+def test_agent_removal():
+    class TestAgent(Agent):
+        pass
+
+    model = Model()
+    agent = TestAgent(model.next_id(), model)
+    # Check if the agent is added
+    assert agent in model.agents[type(agent)]
+
+    agent.remove()
+    # Check if the agent is removed
+    assert agent not in model.agents[type(agent)]
diff --git a/tests/test_datacollector.py b/tests/test_datacollector.py
index a45d8891..a56ca477 100644
--- a/tests/test_datacollector.py
+++ b/tests/test_datacollector.py
@@ -55,6 +55,7 @@ class MockModel(Model):
     schedule = BaseScheduler(None)
 
     def __init__(self):
+        super().__init__()
         self.schedule = BaseScheduler(self)
         self.model_val = 100
 
diff --git a/tests/test_model.py b/tests/test_model.py
index 35ebb4a0..c54634d8 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1,3 +1,4 @@
+from mesa.agent import Agent
 from mesa.model import Model
 
 
@@ -40,3 +41,13 @@ def test_reset_randomizer(newseed=42):
     assert model._seed == oldseed
     model.reset_randomizer(seed=newseed)
     assert model._seed == newseed
+
+
+def test_agent_types():
+    class TestAgent(Agent):
+        pass
+
+    model = Model()
+    test_agent = TestAgent(model.next_id(), model)
+    assert test_agent in model.agents[type(test_agent)]
+    assert type(test_agent) in model.agent_types
diff --git a/tests/test_time.py b/tests/test_time.py
index dda68032..5d6d66f9 100644
--- a/tests/test_time.py
+++ b/tests/test_time.py
@@ -71,6 +71,7 @@ class MockModel(Model):
                               'staged' creates a StagedActivation scheduler.
                               The default scheduler is a BaseScheduler.
         """
+        super().__init__()
         self.log = []
         self.enable_kill_other_agent = enable_kill_other_agent
 
diff --git a/tests/test_visualization.py b/tests/test_visualization.py
index 55e53cdb..ca9f0949 100644
--- a/tests/test_visualization.py
+++ b/tests/test_visualization.py
@@ -33,6 +33,7 @@ class MockModel(Model):
     """Test model for testing"""
 
     def __init__(self, width, height, key1=103, key2=104):
+        super().__init__()
         self.width = width
         self.height = height
         self.key1 = (key1,)