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