is it not possible to just edit the disablebit for collisions? DisableBit.CONTACT?

DisableBit.CONTACT still zeros ncon_out and therefore discards the contacts injected from Newton, see function below. I could move the flag check up but it might have side effects. Inputs from mujoco experts are welcome.

def collision(m: Model, d: Data):
  """Collision detection."""

  # zero collision-related arrays
  wp.launch(
    _zero_collision_arrays,
    dim=d.nconmax,
    inputs=[
      d.nworld,
      d.ncon_hfield.shape[1],
      d.ncon,
      d.ncon_hfield.reshape(-1),
      d.collision_hftri_index,
      d.ncollision,
    ],
  )

  if d.nconmax == 0:
    return

  dsbl_flgs = m.opt.disableflags
  if (dsbl_flgs & DisableBit.CONSTRAINT) | (dsbl_flgs & DisableBit.CONTACT):
    return

  if m.opt.broadphase == int(BroadphaseType.NXN):
    nxn_broadphase(m, d)
  else:
    sap_broadphase(m, d)

  narrowphase(m, d)

I see your point. What does your pipeline look like, you will generate contacts before the call to mjwarp.step()? How are you getting the correct geom positions in this case? FK does run as part of the step as well.

would it make sense to give you a callback or something along these lines?

A callback does not really help since the contacts get already passed into newton's step function. I simply convert them to mujoco. So far it works quite good but I need to test more examples. The following code shows the step function in solver_mujoco.py from newton. The signature of step is fixed, so I don't know if a callback to run contact detection again (=inefficient) makes sense.

    def step(self, state_in: State, state_out: State, control: Control, contacts: Contacts, dt: float):
        if self.use_mujoco:
            self.apply_mjc_control(self.model, state_in, control, self.mj_data)
            if self.update_data_interval > 0 and self._step % self.update_data_interval == 0:
                # XXX updating the mujoco state at every step may introduce numerical instability
                self.update_mjc_data(self.mj_data, self.model, state_in)
            self.mj_model.opt.timestep = dt
            self.mujoco.mj_step(self.mj_model, self.mj_data)
            self.update_newton_state(self.model, state_out, self.mj_data)
        else:
            self.apply_mjc_control(self.model, state_in, control, self.mjw_data)
            if self.update_data_interval > 0 and self._step % self.update_data_interval == 0:
                self.update_mjc_data(self.mjw_data, self.model, state_in)
            self.mjw_model.opt.timestep.fill_(dt)
            with wp.ScopedDevice(self.model.device):
                use_mujoco_contacts = False
                if not use_mujoco_contacts:
                    self.convert_contacts_to_mjwarp(self.model, state_in, contacts)
                    self.mujoco_warp.step(self.mjw_model, self.mjw_data, False)
                else:
                    self.mujoco_warp.step(self.mjw_model, self.mjw_data, True)
                print("step done")
            self.update_newton_state(self.model, state_out, self.mjw_data)
        self._step += 1
        return state_out

Does not make any sense to run contact gen twice indeed. So you are relying on the newton FK to get shape positions? I still don't really get how that would give you updated geom positions. Plus there is a lot of duplication that could be removed if you are relying on an external forward kinematics step anyway.

Back to the main topic, I guess in this setup your proposal makes sense. Is there a world where we would generate contacts for part of the geom in MjWarp, and other geoms have externally generated contacts? Maybe it would make sense to design the API around this.

The MuJoCo Way™ for public API functions to take only Model and Data, and control flow is dictated by Model properties. So I think you have two options here:

1. Don't call step. Instead call the sub-functions of step (skipping collision detection), all of which are in the public API.
2. Propose a property on Model that gives you the behavior you want.

I think either of those we could work with. Long term, if you think you'll be customizing how you use mjwarp a lot, I think 1) is the better bet.

@erikfrey I added a new flag to the options in the model. I tried option 1 but it involves too much code I would need to copy over to newton.