search

google-deepmind / mujoco

Multi-Joint dynamics with Contact. A general purpose physics simulator.
https://mujoco.org
Apache License 2.0
8.19k stars 819 forks source link

Implementing Contact Force Penalization in MJX: Transitioning from Brax's `pipeline_state.qf_constraint` #1209

Closed bheijden closed 11 months ago

bheijden commented 11 months ago

Hi MuJoCo Team,

I am a PhD student and working on a project involving reinforcement learning. I'm in the process of transitioning a sim2real pusher task from Brax to MJX.

In Brax, I leveraged the pipeline_state.qf_constraint feature to integrate penalties for substantial contact forces into my reward function. I am seeking a parallel approach in MJX for achieving the same functionality. More precisely, I require the contact forces exerted on specific bodies, as identified by their names in the XML configuration (for example, <body name="arm/gripper_link" ... >). This includes both the magnitude and direction of these forces in Cartesian coordinates.

Any guidance or suggestions on how to replicate or approximate this feature in MJX would be helpful!

Thank you for your assistance!

EDIT: Data.qfrc_constraint seems to provide applied generalized forces. Is there an easy way to convert them to forces on the bodies in Cartesian space?

btaba commented 11 months ago

Hi @bheijden

Indeed qfrc_constraint is similar to the one in brax. For cartesian contact forces, it looks like you'll want something like mj_contactForce, see https://github.com/google-deepmind/mujoco/issues/997#issuecomment-1647992310 which isn't currently implemented in MJX but is available in MuJoCo after a mjx.get_data

bheijden commented 11 months ago

Oke, thanks :)!