Monolithic Model/State vs Solver-specific Model/State

[etaoxing]

WIth the current monolithic Model/State design in Newton, it is difficult to add custom attributes in order to read Solver-specific data. See #871 for an example of the modifications needed to (ModelBuilder, Model, State, SolverMujoco) in order to read MjData.qfrc_actuator.

This challenge is not limited to SolverMujoco. For example, SolverFeatherstone also defines custom Model and State attributes inside Solver:

newton/newton/_src/solvers/featherstone/solver_featherstone.py

Line 205 in 709422d

def allocate_model_aux_vars(self, model):

newton/newton/_src/solvers/featherstone/solver_featherstone.py

Line 239 in 709422d

def allocate_state_aux_vars(self, model, target, requires_grad):

---

Since not all aspects of MjWarp are exposed to users in Newton currently, users may be discouraged from using newton.SolverMujoco, and instead prefer to use mujoco_warp directly, or use lightweight wrappers on top of it such as mjlab, in order to access the full capabilities of MjWarp.

newton/newton/_src/solvers/mujoco/solver_mujoco.py

Lines 2514 to 2521 in 709422d

	def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
	if nworld == 1:
	return
	model_fields_to_expand = [
	# "qpos0",
	# "qpos_spring",
	"body_pos",

---

If instead there are custom Model/State for each solver, then Newton may be able to directly use MjWarp.types.Model and MjWarp.types.Data (State). I think this may also make multi-physics solver coupling more explicit. One potential API for backwards compatibility:

model = ModelNewton(model_featherstone, model_vbd)
model._default == "featherstone"
model.body_X_com == model.featherstone.body_X_com
model.particle_hessians  # raise AttributeError
model.vbd.particle_hessians = ...

state = StateNewton(state_featherstone, state_vbd)
state._default == "featherstone"
state.joint_tau == state.featherstone.joint_tau

At the very least, I think there should be a solvers/SOLVER/types.py in order to clarify what attributes each solver supports. Currently there is no consistency for how solvers add custom attributes. Style3D, ImplicitMPM, and Featherstone, Mujoco all do it differently. There are also stateful warp arrays as attributes on specific solvers that should belong in Model or State.

---

I am happy to make an attempt towards this change, but given the complexity and scale of this, thought it would be helpful to hear thoughts before trying.

[etaoxing]

@eric-heiden @mmacklin

[eric-heiden]

Thanks for raising this important issue @etaoxing! We are also looking into related problems of parsing/declaring custom, solver-specific attributes in #710.
The approach of reconciling different solver-specific models may also be a useful solution. Let's discuss this more.

@nvlukasz, @adenzler-nvidia, @Milad-Rakhsha-NV for viz

[etaoxing]

Thanks for the pointer Eric! Milad's PR on USD custom parsing looks great, just read through it. It seems like this is a good opportunity to make a wider change.

I think a major pain point for me right now is that newton.SolverMujoco is ~3000 LOC, while the sole functionality (to the best of my understanding) is to convert between Mujoco/Newton conventions and basically has no actual physics sim logic. For comparison, newton.SolverVBD is also ~3000 LOC and has actual physics sim logic. I would hope that the final solution yields a lightweight implementation for SolverMujoco that just does:

class SolverMujoco(SolverBase):
    def step(self, model: Model, state_in: State, state_out: State, control: Control, contacts: Contacts, dt: float):
        data = data_fn(state_in, state_out, control, contacts)  # remap to a Data obj, need to handle when state_in and state_out are swapped
        assert type(model) == mjwarp.types.Model
        assert type(data) == mjwarp.types.Data
        mjwarp.step(model, data)
        # state_out is still updated

Sidenote: I think Solver should be state-less and should go back to passing model into Solver.step(). For example, gravity randomization currently requires newton.SolverMujoco.update_model_properties().

With a lightweight implementation, Newton would be more resilient to upstream changes in mujoco_warp without having to update the solver code in Newton.

In the case of Solver-specific Model/State, we would expect each Solver to define its own custom attributes and conventions (e.g. quaternion scalar order, transformation frame). Then we can either lazy-convert attributes when the user requests it in Newton (inside utils.selection.*View) or prepare views for the user beforehand inside ModelBuilderNewton.build(), ModelNewton, StateNewton. The highest level Newton APIs can use Newton conventions, and then convert to Solver-specific conventions (with no conversion if they already follow Newton).

[adenzler-nvidia]

Thanks for the comments - I agree resolving this very soon is important to keep making progress on having a complete integration of MuJoCo Warp. I also share your pain of the size and complexity of the SolverMuJoCo file, we need to reorganize this.

I'm not sure if I understand your last paragraph correctly - don't we already have this? You can just access the MjWarp model and state fields directly instead of using the newton model/state if you want direct access to some MjWarp data that isn't converted, or don't want to pay the price for the conversion? I think all you need in that case is a way to skip all the conversion overhead if you're not taking advantage of that.

[etaoxing]

That is true, we can always directly access MjWarp and do conversions separately. In that case, why should users bother with using Newton's Model/State? There are two conversion types:

• mapping different attribute names (soon to be addressed by #710?)
• conventions

I am happy to pay the price of conversion in order to use MJWarp in Newton's unified API across multiple solvers. But if I have to separately manage newton.Model / State, newton.SolverMujoco.mj_model / mj_state, newton.SolverMujoco.mjw_model / mjw_state myself, then I'm not sure if using Newton makes sense rather than just using mujoco_warp directly.

I am primarily advocating for some more consistency in when/where conversions happen (in order to declutter SolverMujoco). One potential solution if Newton's Model/State is hierarchical:

model = ModelNewton(model_mujoco, model_featherstone)  # resolver order from PR #710
model.mujoco.qpos == model_mujoco.qpos  # no conversion
type(model.mujoco) == mujoco_warp.types.Model
model.featherstone.joint_q == model_featherstone.joint_q
model.joint_q == conversion(model_mujoco.qpos)  # ModelNewton and StateNewton handle conversion from solver-specific to Newton conventions

state = model.state()
state.mujoco == model_mujoco.state()
type(state.mujoco) == mujoco_warp.types.State
state.featherstone == model_featherstone.state()
state.joint_q == conversion(state_mujoco.qpos)

# then SolverMujoco.mjw_model and SolverMujoco.mjw_data can be removed
# and no conversion occurs inside Solver*

[adenzler-nvidia]

I guess we didn't really communicate our plans very well in this regard. The plan was always that a newton user will never have to deal with the MjWarp model/state structs unless they are doing something very exotic. By version 1.0 of newton, all of MjWarp will be drivable through the newton model/state, with appropriate conversions or solver specific model fields where needed. MuJoCo will always be the odd one out in newton as we are bringing in a stable and battle proven solver that already has a stable API.

Unfortunately progess on this has been slow but should accelerate as we near beta2, your feedback here is similar to what many others are telling us.

We will also clean up solverMuJoCo during this process, but there will likely always be some complexity involved there because of the conversions. Maybe it's an option to move the conversions to a separate file and I'm sure we can find ways to make the high-level flow easier to follow.

All of that being said, we still need a way to handle solver-specific data in the model/state structures so your input is very valuable here. We need to fast-track the design here for sure otherwise we'll never have a world where a users does not need to know about MuJoCo model/state.

[etaoxing]

thanks @adenzler-nvidia!! The transparency on the roadmap is very helpful and reassuring.

I suppose it is also worth mentioning the elephant in the room: Should Newton just match all the conventions of Mujoco? (minus the quaternion ordering 😅) And we modify the other solvers to do conversion or directly use Mujoco conventions instead?

[mmacklin]

Hey @etaoxing, thanks for the comments - wanted to address a couple of things:

Sidenote: I think Solver should be state-less and should go back to passing model into Solver.step(). For example, gravity randomization currently requires newton.SolverMujoco.update_model_properties().

I think the issue here is how the solver knows what to update - we don't want to do this remapping / recreation every step() (since that would be inefficient) and so how would the Solver know which internal state needs to be updated? Although it's an desirable goal to be stateless I don't think it's practical to say that Solver's should have no internal state in general, since almost every solver will have some internal allocations for buffers / data structures for it to be efficient. However they should appear stateless in the sense that if you give them the same inputs (state_in, contacts, etc) then they should give the same outputs.

I agree that the conversion code is pretty dense though, so I think this could be refactored to at least be easier to navigate.

I suppose it is also worth mentioning the elephant in the room: Should Newton just match all the conventions of Mujoco? (minus the quaternion ordering 😅) And we modify the other solvers to do conversion or directly use Mujoco conventions instead?

We definitely want to align with MuJoCo where it makes sense, but there are a number of other solvers that don't follow or fit with MuJoCo's conventions (particular maximal coordinate models). Also we have a number of users coming from IsaacLab, which has its own conventions and behaviors that are derived from PhysX. We would also like them to have a good and familiar experience with the API. That said, we should absolutely share as much as we can, especially around things like collision representations.

On a general note, inheritance is one way to go about specializing Model and State classes, but I think in general we should prefer composition, and take advantage of Python's dynamic typing where we can extend these objects with data from multiple sources. e.g.: so we can combine an MPM simulation with MuJoCo simply by virtue of the way we set arrays on the shared Model.

Cheers,
Miles