In the same way we pass a Callable policy_search to the rollout function in pymdp/envs/rollout.py, one idea to make inference more configurable, is to similarly refactor the infer_states API so that a Config object (maybe a FrozenDict or something) within the Agent class can be used to determine which callables (e.g. smoothing_ovf or my_infer_states_function) are passed into the infer_states pipeline. The only constraints would be that they have consistent input/output shapes and are jittable
In the same way we pass a Callable
policy_searchto therolloutfunction inpymdp/envs/rollout.py, one idea to make inference more configurable, is to similarly refactor theinfer_statesAPI so that aConfigobject (maybe aFrozenDictor something) within theAgentclass can be used to determine which callables (e.g.smoothing_ovformy_infer_states_function) are passed into theinfer_statespipeline. The only constraints would be that they have consistent input/output shapes and are jittable