[gym] Use of optimised utility fonction build_flatten and jiminy_py.tree.flatten don't give me exactly the same outputs

[mwulfman]

I have a script where I initialize an unwrapped environment and where I build the flatten_obs utility function that way:

obs, _ = env.reset(seed=42)
obs_flat = np.concatenate(tree.flatten(obs))
flatten_obs = build_flatten(env.observation, obs_flat)

Then when I do:

obs, _, terminated, truncated, info = env.step(env.action_space.sample())
obs_flat = np.concatenate(tree.flatten(obs))

or

obs, _, terminated, truncated, info = env.step(env.action_space.sample())
flatten_obs()

The obs_flat in both case are close but not equal. Is that the expected behaviour? I understand there's a clip operation between obs and env.observation, is there anything else? From my understanding, build_flatten was the optimised, memory-efficient, version of the np.concatenate(tree.flatten(.))

[duburcqa]

In this case, tree is an alias for jiminy_py.tree or deep mind library dm-tree ?

Maybe it is due to one being implemented in numpy, while the other is jitted.

[duburcqa]

You could try removing the decorator @nb.jit(nopython=True, cache=True) from _flatten and see if the difference vanishes.

[mwulfman]

It is jiminy.tree. I think my question is more how do I use build_flatten + any other need operation if they exist to do what I am trying to achieve. I'll try that. (I edited the post)

[duburcqa]

I understand there's a clip operation between obs and env.observation, is there anything else?

There is no other difference. (Besides being a copy rather than a reference, to make sure that the user can store the returned observation without asking itself whether or not it must be copied first).

I think my question is more how do I use build_flatten + any other need operation if they exist to do what I am trying to achieve.

I don't understand what you mean. You can use more generic utilities such as build_reduce or build_map. All the others are just derived from them.

[mwulfman]

Thank you for your help !

[duburcqa]

No problem ! Feel free to re-open this discussion if you have a specific need. build_reduce and build_map are quite tricky to use directly.