Braindump of various considerations and unresolved questions for the initial design. I recommend first looking at the README.md, which gives a high-level overview of the design, and examples/bevy_snake, which goes into much more details on how might use entity-gym-rs in a Bevy app.

1 Bevy

1a NonSend Agent

The TrainAgent contains a Receiver/Sender which makes it non-Send. So the agent has to be a NonSend resource. Maybe that’s fine but I don’t really know full the implications. The alternative would be to wrap the TrainAgent in an Arc<Mutex<..>>. (EDIT: looks like eliminating the Mutex was a major performance win: https://github.com/entity-neural-network/entity-gym-rs/issues/2)

✅ 1b Dry app definition

Resolved by https://github.com/entity-neural-network/entity-gym-rs/commit/31950005a5384fd43a570145387580b1e9490f00 Currently the App construction is highly duplicated between the run and run_headless methods. Seems easy enough to factor out the common parts into a different functions. Not sure if that’s the idiomatic way of doing this is in Bevy.

❓❗ 1c Closing training environment

When creating an environment from Python, there needs to be some way to shut down the threads running the apps. Currently, this is done by having the Agent::act method return an Option. When the counterparts to the channels owned by the TrainAgent are dropped, TrainAgent will return a None value which can be handled by sending an AppExit event. This feels inelegant, ideally app shutdown should just be handled automatically somehow.

❓❗ 1d Training executor

The current way of spawning a thread for each app instance during training seems suboptimal.

✅ 1e Asset packaging and loading

https://github.com/entity-neural-network/entity-gym-rs/commit/8519bec3a5c792638a06dc58d94c0417b9c35565 implements a Bevy asset loader for RogueNet checkpoints. Currently can load an agent by passing a path to a checkpoint directory to a function. Haven’t looked into this yet, but presumably Bevy has an asset loading system that could give a nicer API. I can easily see games shipping with hundreds of snapshots of different agents that total many GB. There should also be a way of bundling checkpoints with Wasm apps.

✅ 1f Graphics overhead

At least for bevy-snake, this doesn't seem to be the bottleneck: https://github.com/entity-neural-network/entity-gym-rs/issues/2 Even when running with minimal plugins and disabling rendering, creating SpriteBundles and whatnot probably adds unnecessary overhead. Is there a simple way to avoid this by just adding a plugin to the App or something?

2 API

2a Featurizable

The Featurizable trait and way of constructing observations by creating a new struct for each entity rather than just passing the Position directly may seem a bit verbose, but this is actually done for a good reason. The input to the neural network is ultimately an untyped vector of floating point numbers. By materializing the field name of each feature, we can later add additional fields while still remaining compatible with networks trained on an previous version of the environment. Explicitly declaring a struct for each entity given to the network adds an additional layer of indirection that prevents changes to the internal representation from exposing us to backwards compatibility hazards. In theory, it would be possible to completely eliminate the intermediary structs by automatically deriving Featurizable for tuples etc., but I feel the downsides aren’t worth the small amount of boilerplate that this would eliminate. Not completely sure though, maybe both should be an option.

2b Multiple action types

The API currently supports only a single type, but it should also allow not just for multiple action types, but simultaneously performing multiple actions. Not sure how to make the type signatures work out in the current design without variadic generics, macros, or defining act2, act3, act4, …, etc methods.

✅ 2c Agent object safety

Mostly solved by https://github.com/entity-neural-network/entity-gym-rs/commit/0a1dcbd1bae14c4911a4c55ed26643aec791e7a3 The generic act methods prevents Agent from being trait safe, which means you can’t just a have a Box<dyn Agent>. The AnyAgent enum sort of gets around this, but it feels inelegant and won’t work if someone wanted to e.g. implement their own scripted Agent.

✅ 2d Typed Action/Entity API

Mostly solved by https://github.com/entity-neural-network/entity-gym-rs/commit/0a1dcbd1bae14c4911a4c55ed26643aec791e7a3 Having the entity/action methods take a trait as an argument allows for good performance and feels like a nice API (also seems consistent Bevy), but looses some flexibility and has a few annoyances. Maybe we need a more dynamic version of the API as well.

✅ 2e Python API boilerplate

Most of the boilerplate eliminated in https://github.com/entity-neural-network/entity-gym-rs/commit/79dfd6470e4f5ec30629a8cf32b73f38d9201059, I don't see how to do much better than this. Some of the boilerplate in python.rs could probably be eliminated somehow.

3 Misc

3a Entity name collisions

I ended up with a bunch of structs that have a Featurizable and a Component version which is a bit annoying. Seems like this could be a common occurrence, should probably come up with a good convention to prevent name collisions.

✅ 3b Package name

Looks like this might not be possible to solve when using PyO3, module name must be the same as the lib/crate name Ideally the Python package should be entity-gym-rs and the Rust package entity-gym, but for some reason maturin keeps insisting on using the Rust library name for both 🤷

When creating an environment from Python, there needs to be some way to shut down the threads running the apps. Currently, this is done by having the TrainAgent::act method return an Option<Action> rather than an Action directly. When the counterparts to the channels owned by the TrainAgent are dropped, TrainAgent will return a None value which can be handled by sending an AppExit event. Ideally, Agent::act could just return an action directory and app shutdown would be handled automatically without writing additional code. I don't currently know enough about Bevy to know how to make this work. Maybe there could be an entity-gym Bevy plugin that connects to the train agent and handles app shutdown?

The agent API is not 🚀 blazingly 🚀 fast 🚀 yet. Results from benchmarking both a low_level and agent/Bevy implementation of snake with benchmark.py:

Benchmarking multisnake...
Wall time: 2.97s
Throughput: 860.65K steps/s
Benchmarking bevy_snake_enn...
Wall time: 3.38s
Throughput: 37.86K steps/s
Total send: 110 ms
Total wait: 855 ms
Total collect: 24 ms
Total send: 12 ms
Total wait: 3053 ms
Total collect: 2 ms

Already not bad, but still a 20x gap. 800K steps/s is kind of overkill, I think if we can get to 200K/s or something, this will be sufficient for most practical applications.

End-to-end training throughput with examples/bevy_snake/train.py is now ~6000. It was previously closer to 4000, I think what made the difference was making the Player a NonSend resource to eliminate the Mutex on the TrainAgent. This seems to have resulted in a ~4x speedup.

With the Mutex:

NonSend resource without the Mutext:

Will get some profiles next to see what parts we can optimize.