Separate observation processing from state inference

[conorheins]

Following the landing of #315 to v1.0.0_alpha and following up on @nikolamilovic-ft's comment there, in order to separate the optimized inference execution itself from the overhead introducing by padding observations, it makes sense to separate the processing of observations from the inference algorithm run upon them.

The idea would be to add a separate process_obs() method to the Agent class which can be overridden in a custom way (or be flexibly configured based on some kind of InferenceConfig if and when we get around to adding that, see e.g. #275) when using advanced inference algorithms where padding/preprocessing overhead can be segregated from the execution of the inference algorithm on those preprocessed observations. This is based on conversations with @praesc @nikolamilovic-ft where it can be useful to preprocess observations separately (even in a different runtime or at least jitted context) before running inference upon them.

The default behavior of process_obs() would be like what currently happens at the beginning of infer_states(), namely

def process_obs(self, observations, categorical_obs=False):
       use_categorical = categorical_obs or self.categorical_obs

        if not use_categorical:
            o_vec = [nn.one_hot(o, self.num_obs[m]) for m, o in enumerate(observations)]
        else:
            o_vec = observations
       
        return o_vec

When using optimized inference algorithms that require advanced padding techniques, this method could be overridden in an algo-specific way.

Then instead of calling only infer_states() in context of e.g., rollout.py::infer_and_plan(), we do instead

processed_obs = agent.process_obs(observations)
qs = agent.infer_states(processed_obs, ...)

This way, the processing could in theory (in custom designs) be separated out from the inference. If we want to keep the lines of code in rollout streamlined however, we could still create some default version of infer_states that includes both steps and basically is exactly what is currently happening in infer_states(), namely

def infer_states_fused(observations, categorical_obs=False)
     processed_obs = self.process_obs(observations, categorical_obs=categorical_obs) # with the default behavior, e.g. just categorical vs discrete index obs
     return agent.infer_states(processed_obs, ...)

[tverbele]

I'd say this is fine as long as this happens within infer_and_plan (and hence rollout) , i.e. a typical pymdp user should not be dealing with this (and focus the documentation on using rollout and infer_and_plan methods)

[conorheins]

Another idea

def infer_states(...., preprocess_fn: Callable = default_process_obs)
       
     observations = default_preprocess_fn(observations)
      # rest of inference

where default_process_obs is the current default of

def default_process_obs(observations, num_obs, categorical_obs=False):

        if not categorical_obs:
            o_vec = [nn.one_hot(o, num_obs[m]) for m, o in enumerate(observations)]
        else:
            o_vec = observations
       
        return o_vec

[praesc]

How about adding a simple extra boolean flag obs_preprocessed here

That would do something like:

        if obs_preprocessed or use_categorical: 
            o_vec = observations
        else:
            o_vec = [nn.one_hot(o, self.num_obs[m]) for m, o in enumerate(observations)]

That allows us to pass pre-processed observations.

[Arun-Niranjan]

I agree with the overall sentiment of this issue, but I think @tverbele makes a good point about trying to reduce complexity for the typical user. Optional overriding of a default method is likely the way to go.

The critical benefit I'd like to see is a reduction in the number of axis/dimension errors, and once observations are processed they can be consistently used across other publically exposed functions in pymdp e.g.