instadeepai / Mava

🦁 A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX

Home Page:https://id-mava.readthedocs.io

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

[FEATURE] Generalise Evaluator Fn

EdanToledo opened this issue · comments

Please describe the purpose of the feature. Is it related to a problem?

It would be very nice if the evaluator function was generic enough for future algorithms. It can be frustrating that a highly specific form is required. i.e. currently the network apply is force vmapped and it uses the default call function. Complex algorithms have complex networks and multiple different call functions and might not be suitable for this paradigm.

Describe the solution you'd like

An easy way to do this would be that a user passes in an "actor_fn" which has a specific input and output api. A user would create this and perhaps make a partial fn to fit the correct api. This could be suitable for feedforward and recurrent networks as long as the API is adhered to. I feel that although this is adding structure that could make a user confused, its a simple enough request that it wouldn't require the user to dig deep at all and would benefit creating new algorithms.