Altering body measurements on fly (i.e. in state)

[peabody124]

Love brax! Thanks for all your work. The ability to run billions of sim steps in 30 minutes is amazing.

I am really hoping to find a way to change some of the body parameters on the fly. For example, to alter limb segment lengths in the reset function and store that information as an augmented part of the state. From my reading, all that information if bound up in System, which is expected to be static during training and inference.

Is this something that is feasible with brax, or is it complete antithetical to the design decisions? I can imagine ways to implement this for very specific a specific env that uses a custom superclass of System that pulls in the appropriate information from the state for each step, but am worried it might ultimately require overriding all of the forces/actuators/etc to keep the implementation "pure".

(p.s. if this is something that could be cleanly implemented in a general way, I'd be happy to work on a PR)

[peabody124]

@erikfrey do you have any thoughts on this?

[btaba]

Hi @peabody124 ! Thanks for using brax :D

This is similar to the discussion on #229 (comment) where modifying params in the System (without having to re-jit) can be done via an experimental wrapper, as is done in https://github.com/google/brax/blob/main/brax/experimental/tracing/wrappers.py

In the next major version of brax, the System will be traceable by design, so tracing over environment settings should be a bit easier to handle. Hope this helps somewhat for now!

[peabody124]

Thank you! @btaba I'll try it out. It appears much more elegant than my current approach having to overload the system and implementation of joints to take in additional parameters.

[peabody124]

@btaba thanks for that pointer. I managed to get it running ... somewhat. However, in my case I'm trying to have the environment modified over a continuous set of variables that are also being learned (yay differentiable simulator!). I adapted the DynamicWrapper to handle this but find that the memory usage blows up massively. For example things that were running on an A100 now request 350GB of GPU ram.

Is this to be expected? A much uglier approach that passes these scaling parameters to the physics engine runs without such an additional overhead.

def create_joint_scale_custom_tree(hip_width=1.0, thigh_length=1.0, leg_length=1.0, torso_length=1.0,
                                   shoulder_width=1.0, arm_length=1.0, forearm_length=1.0):

  scales = [torso_length, torso_length,
            hip_width, thigh_length, leg_length, hip_width, thigh_length, leg_length,
            shoulder_width, arm_length, forearm_length, shoulder_width, arm_length, forearm_length]

  def make_joint_scale(x):
    return {'parent_offset': {'x': x, 'y': x, 'z': x}}

  custom_tree = {'joints': [make_joint_scale(x) for x in scales]}
  return custom_tree


@pytree.register
class TracedConfig:
  """A wrapper around the config that faciltates jax tracing of static data."""

  __pytree_ignore__ = ('msg',)

  def __init__(self, msg: Any, custom_tree: Dict):
    self.msg = msg
    self.custom_tree = custom_tree

  def __getattr__(self, name: str):

    base = self.msg.__getattribute__(name)  # pytype:disable=attribute-error

    if name in self.custom_tree:
      if isinstance(base, Iterable):
        list_msg = []
        for o, b in zip(self.custom_tree[name], base):
          list_msg.append(TracedConfig(b, custom_tree=o))
        return list_msg
      elif not isinstance(self.custom_tree[name], dict):
        return self.custom_tree[name] * base
      else:
        return TracedConfig(base, self.custom_tree[name])

    return base


def post_process_fn(config, scales):
  custom_tree = create_joint_scale_custom_tree(scales[0])
  config = TracedConfig(config, custom_tree)
  return config


class AnthromomorphicWrapper(brax_env.Wrapper):
  """Wraps environment methods to allow vectorization over config state."""

  def __init__(self, env_fn: Callable[..., brax_env.Env]):
    self.env_fn = env_fn #pytree.register(env_fn)

    def init_env_fn():
      env = self.env_fn(post_process_fn=lambda x: post_process_fn(x, jp.ones((7,))))
      return env

    env = init_env_fn()
    super().__init__(env)

  def reset(self, rng, anthropomorphic_params=None):
    if anthropomorphic_params is not None:
      env = self.env_fn(post_process_fn=lambda x: post_process_fn(x, anthropomorphic_params))
    else:
      env = self.env_fn()
    return env.reset(rng)

  def step(self, state: brax_env.State, action: jp.ndarray, anthropomorphic_params=None) -> brax_env.State:
    if anthropomorphic_params is not None:
      env = self.env_fn(post_process_fn=lambda x: post_process_fn(x, anthropomorphic_params))
    else:
      env = self.env_fn()
    return env.step(state, action)

[peabody124]

I also opened an issue about the TracedConfig not influencing the renderer #265