google / google/brax

TracedConfig does not influence rendering

Open
#265 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
3.2k
Forks
349
PR merge metrics
No merged PRs in 30d

Description

I was trying to extend the domain randomization examples to include changing collider shapes using the code below. However, the renderer doesn't pick up the changes. Specifically this occurs because the mesh information comes from `json_format.MessageToDict` which does not pick up any of the changes in TracedConfig.

I don't know enough about the calls json_format.MessageToDict is making into the internal structures to know what to override, but if anyone knew and wanted to point me in that direction, I would be happy to make a PR.

```
def scale_bodies(config, body_scale_dict: dict):
"""Constructs tree and in_axes objects for a joint socket randomizer.
Adds an offset to any joints that match a key appearing in joint_key. If
no joint_key, then does nothing.
Args:
env: Environment to randomize over
body_scale_dict
Returns:
Tuple of a pytree containing the randomized values packed into a tree
structure parseable by the TracedConfig class
"""

custom_tree = {'bodies': []}

for b in config.bodies:

def scale_body(b, x):
colliders = []
for c in b.colliders:
collider = {
'position': {
'x': c.position.x * x,
'y': c.position.y * x,
'z': c.position.z * x
},
'capsule': {
'length': c.capsule.length * x,
'radius': c.capsule.radius * x
}
}

colliders.append(collider)
return {'colliders': colliders}

if any([key in b.name for key in body_scale_dict.keys()]):
custom_tree['bodies'].append(scale_body(b, body_scale_dict[b.name]))

else:
custom_tree['bodies'].append(scale_body(b, 1.0))

return custom_tree
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.