infer-actively / pymdp

A Python implementation of active inference for Markov Decision Processes
MIT License
419 stars 83 forks source link

Remove redundant operation for normalization in case of 3-tensors #121

Closed conorheins closed 10 months ago

conorheins commented 1 year ago

@OzanCatalVerses pointed out that norm_dist within the pymdp.utils module strangely treats the case of 3-tensors differently than other numbers of dimensions, see here and code below.

def norm_dist(dist):
    """ Normalizes a Categorical probability distribution (or set of them) assuming sufficient statistics are stored in leading dimension"""
    if dist.ndim == 3:
        new_dist = np.zeros_like(dist)
        for c in range(dist.shape[2]):
            new_dist[:, :, c] = np.divide(dist[:, :, c], dist[:, :, c].sum(axis=0))
        return new_dist
    else:
        return np.divide(dist, dist.sum(axis=0))

Will correct this to:

def norm_dist(dist):
    """ Normalizes a Categorical probability distribution (or set of them) assuming sufficient statistics are stored in leading dimension"""
      return np.divide(dist, dist.sum(axis=0))