rollout function is incompatible with onehot_obs Agent?

[Arun-Niranjan]

With an initial setup like so:

from jax import random as jr
from pymdp.agent import Agent
from pymdp.envs import TMaze, rollout

key = jr.PRNGKey(0)
batch_size=1
T = 5
env = TMaze(
    batch_size=batch_size,
)
A, A_dependencies = env.generate_A()
B, B_dependencies = env.generate_B()
agent = Agent(
    A=A,
    B=B,
    A_dependencies=A_dependencies,
    B_dependencies=B_dependencies,
    batch_size=batch_size,
    policy_len=T,
    onehot_obs=True,
)
_, info, _ = rollout(agent, env, num_timesteps=T, rng_key=key)

rollout raises an exception when trying to call agent.infer_states(...).

This does not occur when setting onehot_obs=False in the Agent constructor. We should either:

1. fix rollout to handle one hot encoded observations (preferable)
2. remove the flag from the agent constructor

if we keep support for one hot encoded observations, we should add unit tests to cover this option

[conorheins]

Good catch, and more reason for us to get unit tests for rollout function integrated sooner rather than later (#219)

[Arun-Niranjan]

Do we actually want to enforce one-hot encoded observations? Or is the real value in being able to supply a vector that represents the belief about an observation?

This could either be a belief about future observations that haven't occurred yet, or past/present observations that are sampled with uncertainty (e.g. a sensor reports a value with a standard deviation). In the case of certainty the distribution tends towards a delta function i.e. one-hot encoded.

Entirely possible I'm confusing multiple concepts here, and that uncertainty in observations should only be handled explicitly in the agent's hidden states.

[conorheins]

Yeah there are use cases when you want K-vector ("one hot") observations -- the uncertainty case is the most common one, but actually we shouldn't call this one hot observation, tbh, maybe should call it categorical observations instead since we use the other entries of the K vector to encode uncertainty.

A subtlely different but related case is hierarchical models where one agents observations are the hidden state posteriors of another agent (a lower level POMDP).

Either way, we should support both onehot(really: categorical) observations as well as integer index ones

[Arun-Niranjan]

perhaps this is over the top, but would it also be possible to support changing observation vector dimensions? I.e. during Agent operation, it learns that an observation can take a new value - or if after a certain time if an observation value is never observed, it removes that option entirely.

My instinct is that this wouldn't play nicely with JAX, or that there's a smarter way to implement this with a heirarchical model and being clever with priors - structure learning is tricky 😅

[tverbele]

The jax-compatible way of doing the above is to allocate for some maximum, i.e. have a state or obs dimension of K, initialize your agent only using k < K. Then over time you can start using extra dimensions and parameter slices as data comes in. The other option is to completely re-create your Agent object with expanded parameters, but this implies also having to re-jit all the functions for the new sizes.

[Arun-Niranjan]

I think for now the good middle ground to strike before a v1 release is either a K-vector or an individual value for observations, and support structure learning in the observation space as a follow on iteration

A subtlely different but related case is hierarchical models where one agents observations are the hidden state posteriors of another agent (a lower level POMDP).

This is the 'killer app' in my mind, anything we can do to make the developer experience easier to achieve this will pay dividends