search

google / brax

Massively parallel rigidbody physics simulation on accelerator hardware.
Apache License 2.0
2.14k stars 234 forks source link

Fixing Bodies Dynamically #366

Open tomalner18 opened 1 year ago

tomalner18 commented 1 year ago

Hi, I'm looking for a way to fix/lock a body in its current position. Is it possible, for example, to set the joint ranges - in this case of a body with a 'free' joint - to the object's current q during a rollout. If you have any kind of example that would be amazing, thanks!

btaba commented 1 year ago

Hi @tomalner18 , thanks for using brax! You'd want to edit the dof limits

https://github.com/google/brax/blob/3b72c952f7c990e8e9e355f736e77f10f408bc14/brax/base.py#L294

here's an example

@jax.jit
def do_stuff(sys, state):
  state = pipeline.step(sys, state, jp.zeros(sys.act_size()))
  sys = sys.replace(dof=sys.dof.replace(limit=(sys.dof.limit[0].at[0].set(0), sys.dof.limit[1].at[0].set(0))))
  state = pipeline.step(sys, state, jp.zeros(sys.act_size()))
  return state

let us know if that works!