This is probably a really silly question, but I have no experience whatsoever with JS and anything web related.

Im trying to add a torus primitive to the set of colliders. I made a simple mostly empty env to try it out in, and im using the HTML render functionality showcased in other notebook examples to try and debug my progress. Ive added colliders, changed the protobuf definition and compiled it, etc. Not much of a clue what im doing yet, but all parts of the brax repo that mention spheres now have a torus equivalent. The plane I added in my env shows up nicely; but my torus does not show up. Could be many reasons of course, but ive found at least one. Deleting createPlane in system.js in my repo, also does not stop my plane render code from working.

Further digging reveals that that html.py has a line import {Viewer} from 'https://cdn.jsdelivr.net/gh/google/[email protected]/js/viewer.js';

Unless im missing something, thats not referring to viewer.js inside my repo; so it makes sense im not seeing changes made there reflected in my notebook.

Now im not sure if this is just some WIP development code that got merged, or if im missing something fundamental about the zen-of-webdev here; but if you were me and looking to make a change that would allow me to do local development on this JS, and also would stand a chance of getting merged, what would you do? Do I need to locally host my viewer.js and link the HTML to that? But then how would I merge a non-broken PR if its supposed to refer to this CDN of a past brax release?

In general, is there something im missing about doing local development? Is debugging using the 3JS viz in a notebook the way to go?

As a number of machine learning projects use conda and it is better for conda is all dependencies are also on conda

Could Brax to added to conda-forge?

I'm trying to simulate a pool billiards tabletop that I've modeled with a plane as floor and 4 planes surrounding it, facing inward. If I now spawn a ball and bounce it off the walls, sometimes the collisions work well:

(left is top-down view, right is side view)

...but sometimes it doesn't work and the ball gets stuck in the wall:

I've made a Colab to reproduce the issue: https://colab.research.google.com/drive/1flnseQcjarIYM4G_rECTEAaps-kXPoey?usp=sharing

If anybody has any pointers, that'd be greatly appreciated.

Best, Flo

• Add optional float field "friction" to colliders; If unspecified, the global friction coefficient is applied
• Re-compile protobuf file (config_pb2.py)
• Introduce Collider dataclass to pass into _collide() and _collide_pair()
• Replace references to config.friction with per-body friction (Collider.friction)

Fixes #55

Please note:

1. I did some primitive tests within physics_test.py but observed that the simulation is currently unstable with friction due to lateral friction, no matter how large I set the number of substeps. Hence I ended up rolling back, in case someone else can advise. I'd appreciate any help writing tests for these changes in near future.

2. I am preserving the plane-body (and mesh-body?) collision behaviour in _collide(), i.e. only the body's coefficient matters and not the plane's. However, a plane's collision coefficient is pretty important for rolling and sliding motions. I'm not a simulation expert yet, but I'd love to see how e.g. Bullet or MuJoCo implements lateral friction.

I was trying to get drag force into my simulation in order to get the swimmer environment running (and possibly simple underwater simulations!) because ther rest of the swimmer env is just simple to port over to brax.

The drag force can be described as follows:

$F_d = - 1/2 \rho ||v||^2 A C_d unit_vector(v)$

where:

• \rho is the density of the fluid
• ||v||^2 is the magnitude of velocity squared
• A is the surface area in the direction of the velocity
• C_d is the Coeficient of drag, (friction of interaction)

\rho and C_d as well as the 1/2 could easily be merged into one constant as they initially play a minor role of determining the properties of the interaction.

This can in principle added easily to the location in code where the forces are applied to the different bodies. I got stuck with two problems:

• Where is the best location in the code to add this formula or how do you want to add it? I think you guys know a lot better than I do @cdfreeman-google
• Whereas the velocity of each body is certainly available, the surface area in the direction of velocity is certainly not (right?). I am pretty sure as I could not think of a part of brax that would require it. I think I will have to implemet for each collider shape a function that projects the collider onto a plane whose normal looks in the direction of the velocity. I am happy if there is an easier way to do this in brax.

https://github.com/google/brax/blob/8e58feb923ce86b7b8c7036a05429793bbc3fa65/brax/physics/math.py#L278

Little nitpick but

S = jnp.array([1., -1., -1., -1.])
def inv_quat(q):
	return q * S

Benchmarks as 10% faster on my laptop cpu at least; and I suspect the same would be more true of architectures more aggressively tuned for vectorization. Dont have any experience with TPUs and their compilers, but this formulation would also make it easier for a GPU compiler to get to the GPU-optimal compiled code I imagine.

For discussion: Beginning to implement external torque inputs. API might make sense to change, though, and is currently incomplete.

The idea is that the frozen field on the Thruster could indicate whether to freeze either force or torque axes, allowing anywhere from 0/1 to 6 DOF for a Thruster.

Currently this ignores the frozen field and provides all possible DOF. In addition, it might be worth adding a different strength field for torque versus force. However, I'm not going crazy in case a preferred approach is an entirely separate class for external torques than forces.

This builds on #94 and addresses #61

Hello,

Are you planning to create any multi-agent environment such as crowd simulation?

Is there also possibility to have a non-uniform terrain, walls etc in each environment?

so that each agent can be initialized in a random location for varying its experience.

(without that, I don't see a major advantage of parallel simulation capability of engine)

Sincerely, Kamer

Hello maintainers of Brax,

I really like your repo and would like to add support for height maps in order to train locomotion policies in uneven terrain.

This PR implements both :

• The visualization of height maps with the THREE interface.
• The collision handling between height maps and box corners.

I hope this might be useful to you and wish the best for your project.

Thanks for your great work. I just finish the installation and the verison of the libs are:

brax              0.0.12 
jax                0.3.7
jaxlib             0.3.7+cuda11.cudnn805

However, when I run the "learn" from the README, the log shows:

Traceback (most recent call last):
  File "/home/yangwang/brax-0.0.12/env/bin/learn", line 7, in <module>
    exec(compile(f.read(), __file__, 'exec'))
  File "/home/yangwang/brax-0.0.12/bin/learn", line 4, in <module>
    from brax.training import learner
  File "/home/yangwang/brax-0.0.12/brax/training/learner.py", line 26, in <module>
    from brax.training import apg
  File "/home/yangwang/brax-0.0.12/brax/training/apg.py", line 32, in <module>
    import optax
  File "/home/yangwang/brax-0.0.12/env/lib/python3.9/site-packages/optax/__init__.py", line 17, in <module>
    from optax import experimental
  File "/home/yangwang/brax-0.0.12/env/lib/python3.9/site-packages/optax/experimental/__init__.py", line 20, in <module>
    from optax._src.experimental.complex_valued import split_real_and_imaginary
  File "/home/yangwang/brax-0.0.12/env/lib/python3.9/site-packages/optax/_src/experimental/complex_valued.py", line 32, in <module>
    import chex
  File "/home/yangwang/brax-0.0.12/env/lib/python3.9/site-packages/chex/__init__.py", line 17, in <module>
    from chex._src.asserts import assert_axis_dimension
  File "/home/yangwang/brax-0.0.12/env/lib/python3.9/site-packages/chex/_src/asserts.py", line 26, in <module>
    from chex._src import asserts_internal as _ai
  File "/home/yangwang/brax-0.0.12/env/lib/python3.9/site-packages/chex/_src/asserts_internal.py", line 32, in <module>
    from chex._src import pytypes
  File "/home/yangwang/brax-0.0.12/env/lib/python3.9/site-packages/chex/_src/pytypes.py", line 40, in <module>
    CpuDevice = jax.lib.xla_extension.CpuDevice
AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'

I do not have a clue about this problem and I already use the latest version of both brax and jax (with GPU option).

Hope you can give me some suggestion.

Hello,

I've been experimenting with some swing-up pendulum environments, I thought Brax might find some of them useful / interesting.

I am proposing adding an Acrobot environment to Brax. The one I am submitting here is a bit different and more difficult than the one in I.E. gym, since the action state in continuous, and to "solve" the environment an agent must both swing up and balance the system. I like this environment because it is deceptively difficult for most model free RL with it (see for example this paper I wrote a couple of years ago).

I've got my own twist on APG that works well for this acrobot, but haven't yet been able to get good performance on it from any of the brax algorithms (including APG). I would be very interested to hear if anyone gets the brax RL working well with the environment!

In this PR I also remove some redundant / useless observations from the double inverted pendulum.

Let me know what you all think!

Edit: Also, here is a basic smoke test training set up in collab.

I have added mass property in URDF and upper/lower revolute joint to Brax model importer.

Test updated accordingly.

I have not added the properties for inertia since the internal model config differs from the official one, being the official one a 3x3 matrix.

Inertia property in Body is as a Vector3, I guess changing it to a matrix 3x3 can break other parts of the engine, if this change should be made I would gladly take it.

https://github.com/google/brax/blob/7eaa16b4bf446b117b538dbe9c9401f97cf4afa2/brax/physics/config.proto#L24-L35

How would the other limit properties fit in the engine computations? I tried to look for effort and velocity.

I saw limit_strength but without being sure of the units it used I didn't want to break anything.

For the case of velocity I did not find any similar parameter in the back end, are joints speed limited?
Regarding joint types "universal" and "spherical" it was not clear to me how are they written in URDF, they are not official. I saw similar ones in Gazebo, SDF and MuJoCo. Due to this I left them as they were.

Hello, I am trying to run some experiments using pytorch with the JaxToTorchWrapper. I'm running the default Training in Brax with PyTorch on GPUs on a local jupyter instance, but there are errors.

The error occurs when I use jax==0.4.1 and goes away when I use a lower version. It seems like it has to do with the new jax.Array type introduced in 0.4.1.

Environment:

• Python 3.10.7
• Cuda 11.8
• jax[cuda]
• brax==0.0.16

Thanks!

I tried running mujoco_convert.py on a sample MJCF file provided on the dm_control repository (this one, but others have similar issues). I get errors such as "unsupported geom type : None", because the type attribute of the elements is not always defined in the MJCF, since it relies on the defaults defined by MuJoCo in their XML reference. For instance, the documentation seems to indicate that the default value for the attribute "type" is "sphere".

I asked over at dm_control to know whether the parse_xml function worked as intended or if they planned on adding the default values to it, and they said it was the expected behaviour. So I think brax should take those into account after the parse_xml function. What do you think ?

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

This is an implementation of https://arxiv.org/pdf/2204.07137.pdf

Not sure if there is interest merging this into the main branch. This might be an algorithm worth supporting as it leverages the differentiable simulator to outperform PPO according to the paper.

Note that many of the environments don't actually have rewards that are differentiable w.r.t. the actions, in which case this algorithm performs poorly. For example, the fast environment used for testing APG and SHAC isn't. I added a fast_differentiable env and also made APG use this by default, after which the performance is much better.

Still could do with tuning for environments and replicating the performance benefits seen in the original manuscript.

Addressed #247

Hey 👋

I am working in my fork on supporting agent specific observations in the composer sub-package. Additionally I need this feature to support observation masking that is dependent on the current state of the environment.

Best approach I can see:

1. Each agent component already accepts observers a) Similar to how reward_fns is done b) This could support dynamic masking
2. Introduce two methods to flatten and unflatten the observation array so that each agent can be given it's specific observation.
3. Pass this sort of information inside the edge parameter to Composer: a) Could be part of its responsibility"automatically create necessary edge information among 2+ components" b) But maybe edge information is only meant to care about pairs of components

Let me know if this sounds cool with you guys, but otherwise I am going to make a start on this feature 😀