How to render observations generated from a learned transition model

[jamesheald]

I am learning a transition model that predicts future observations given actions in a model-based RL setting. I am wondering if it's possible to render these predicted observations? NB I am only learning to predict the observations not the full pipineline_state.

In openai gym mujoco environments, it's possible to to set the environment state by calling 'env.set_state(qpos, qvel)' (and this new state can subsequently be rendered). Is there a similar way to set the state (using observations) in brax? From what I can tell brax uses state.pipeline_state.x.pos and state.pipeline_state.x.rot to render images, so there are two parts to my question:

1. Is there a (straightforward) way to determine state.pipeline_state.x.pos and state.pipeline_state.x.rot from state.obs?
2. Can I assign new values to state.pipeline_state.x.pos and state.pipeline_state.x.rot (when I attempt to do this I get error messages 'FrozenInstanceError: cannot assign to field 'pos'' and 'FrozenInstanceError: cannot assign to field 'rot'', respectively)?

[simon-bachhuber]

You should try x, _ = kinematics.forward(sys, qpos, qvel). qvel is not really required as you don't care for the returned Motion vector.

You must use state.pipeline_state._replace(x=x)

[jamesheald]

Perfect. Would you mind clarifying the meaning of the underscore in _replace. I get an error message 'AttributeError: 'State' object has no attribute '_replace'' when calling state.pipeline_state._replace(x=x) but state.pipeline_state.replace(x=x) works fine.

[simon-bachhuber]

Oh then .replace is what i meant. For some other immutable objects like tuples it's called ._replace.