Apply external force to the agent

[aliang8]

Hi,

I was wondering if anyone has pointers on how to apply an external force to an agent, specifically the HalfCheetah environment.

I read in the MuJoCo documents there is a xfrc_applied field that implements this. Is there a similar feature for BRAX environments? Thanks!

[aliang8]

Bumping this!

[btaba]

Hey @aliang8 , there isn't a similar feature in brax right now, but have you checked out MJX? We aim to have feature parity with MuJoCo using that physics engine (which is now also ported into brax via f0cf109)

[NicoJiao]

Hi, I encounter the same issue. Do you have any solution yet?
Best regards

[Lhy6900]

Is there any latest progress on this issue? I have also encountered difficulties.

[jongyaoY]

I would do something like this:

def process_action(action, pipeline_state):
    nv = pipeline_state.ctrl.shape[0]
    # Assuming the first nv elements are joint actions, the rest are external forces
    joint_action = action[:nv]
    xfrc_applied = action[nv:]
    xfrc_applied = jp.vstack([jp.zeros((1, 6)), action]) # zeros for the worldbody
    xfrc_applied = jp.reshape(xfrc_applied, (-1, 6)) # make sure the shape is (nbody, 6)
    
    pipeline_state = pipeline_state.replace(xfrc_applied=xfrc_applied)
    return joint_action, pipeline_state

in PipelineEnv.step:

pipeline_state0 = state.pipeline_state
action, pipeline_state0 = process_action(action, pipeline_state0)
pipeline_state = self.pipeline_step(pipeline_state0, action)

# ...other stuff

This should work with the mjx backend

[Lhy6900]

Yes, replace function is supported in mjx. Thank you.