3-dim tensors can't be float32 + prohibited 4-dim tensors

[MRiabov]

Hello Danijar,
Thanks for the amazing repo! Great work.

So I have a sizeable state space (more than in minecraft, for sure), and I have at least a dozen 3-dim float32 tensors.

They seem to be prohibited by this assertion:

class Encoder(nj.Module):
    # ...encoder params...
    def __init__(self, obs_space, **kw):
        assert all(len(s.shape) <= 3 for s in obs_space.values()), obs_space
        self.obs_space = obs_space
        self.veckeys = [k for k, s in obs_space.items() if len(s.shape) <= 2]
        self.imgkeys = [k for k, s in obs_space.items() if len(s.shape) == 3] #< collected here
        self.depths = tuple(self.depth * mult for mult in self.mults)
        self.kw = kw


#which is later asserted as uint:
          assert all(x.dtype == jnp.uint8 for x in imgs)

Is there a reason for prohibition of other dtypes? For me, it's a fact of environment that I need to use float32 and not int32 because I'm working with precision data. It's simply unavoidable to use the float32.

I also have a few 4-dimensional tensors which I should use. What edits should be made to support them? (there is an assertion that prohibits shape to be >3dim state.

Cheers,
Maksym

[MRiabov]

Hello @danijar,
Any ideas on this? I'm blocked by this.
Should I just expand the the imgkeys to support 4-dim tensors?

Also, I would add that I have a lot of data which is high-dimensional but not really an "image". Think points in a pointcloud.
Should it go to veckeys and not imgkeys? Since the data isn't really related and is more "SQL-like".
Cheers.