Using functions instead of tensors

[adamjuliangoldstein]

Hi there—

The most challenging part of using and debugging pymdp has, for me, been keeping track of the dimensions of the tensors and their meaning. I found the attached slide from @conorheins recent livestream very helpful for writing my own code; nonetheless, it remains difficult for me to read and comprehend code at a glance. I was emailing with him about an idea I had and he suggested I post it here.

In brief, the idea is to create an abstraction layer for A, B, C, and D tensors so they can be implemented via functions rather that directly as tensors. To reuse the example from the screenshot, I could implement an A tensor via something like the following function, which would take a dictionary of observations and a dictionary of hidden state values:

def A(self, observations, hidden_state_values):
  obs = observations['Loc']
  if hidden_state_values['y']*3 + hidden_state_values['x'] = obs:
    return True
  else
    return False

Or, to take another contrived example:

def A(self, observations, hidden_state_values):
  if hidden_state_value['rain'] or hidden_state_value['snow']:
     return 0.4
  elif observations['wet']:
     return 0.3
  else:
     return 0.2

This could be implemented by having A() be a method of Agent (so I'd subclass Agent to implement my own A) or passing in A as a function to other methods.

Initially, this would be translated behind-the-scenes into an A tensor, e.g. by alphabetizing the keys of each dictionary so that the order of the keys was always the same. (There would need to be functions for each of A, B, C, etc. that used the same dictionaries.) This would mean these kinds of functions could be layered on top of the existing architecture.

Potentially, over time, there could be cases where there wouldn't be a need for an explicit A tensor, and the A function could be lazily evaluated only as needed; of course this would require a significant re-architecture.

If there's interest in this approach I'd be happy to contribute some ideas or code to make this a reality.