Support for drag force to implement gym swimmer environment

[benelot]

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.

[benelot]

As soon as we defined this, I can make a PR to get this into the code.

[cdfreeman-google]

Hello!

The Force primitive will appear in the next github sync, and I think the most natural place to slot this is as a new kind of Force. We can collect the necessary metadata about colliders for different bodies there, and feed it into the Force calculation in a similar way to how we do this for collision calculations.

I think we will need some kind of get_cross_sectional_area fn call for each primitive (is it really surface area in the direction of motion, or is it cross sectional area?). I don't think such a function is too hard to write down for our existing primitives, but if it requires thinking about occlusion, then it may take some additional work. We could get a beta version up and running pretty quickly without it, though, I think.

I'll ping this thread again once we sync Forces in!

[benelot]

Perfect!

[benelot]

@erwincoumans

[benelot]

I think we will need some kind of get_cross_sectional_area fn call for each primitive

It seems pretty confusing on the net as I just find out, but most descriptions of the area is the area some body covers when you look at it from the direction of motion. So cross-section would be correct only for a few edge-cases, but for instance for a tilted capsule, you do not just want a slice like an ellipsoid from the cylinder part, but area of the whole object projected flat. So the correct wording seems to be projected area in the direction of motion or something.

[cdfreeman-google]

That makes sense. Thankfully swimmer is only in 2d, which makes this a bit simpler.

[cdfreeman-google]

Alright, check out https://github.com/google/brax/blob/main/brax/physics/forces.py for the implementation of our simplest force. It's used to power the cart in the new inverted_pendulum and inverted_double_pendulum environments

[erwincoumans]

MuJoCo implements a heavily simplified viscosity model, based on the inertial box and its bounding sphere as far as I can tell, not on the actual collision geometry.

[benelot]

Alright, check out https://github.com/google/brax/blob/main/brax/physics/forces.py for the implementation of our simplest force. It's used to power the cart in the new inverted_pendulum and inverted_double_pendulum environments

I see how you define the thruster to the system config and that you apply a step in the pendulum env. Is the truster then just added as an action in step? Because in the swimmer's case, the force is not an action, but a reaction to the limb velocity and is applied to the COMs of the limbs. I am trying to get how this is done, but any hint could be helpful.

[cdfreeman-google]

Alright, check out https://github.com/google/brax/blob/main/brax/physics/forces.py for the implementation of our simplest force. It's used to power the cart in the new inverted_pendulum and inverted_double_pendulum environments

I see how you define the thruster to the system config and that you apply a step in the pendulum env. Is the truster then just added as an action in step? Because in the swimmer's case, the force is not an action, but a reaction to the limb velocity and is applied to the COMs of the limbs. I am trying to get how this is done, but any hint could be helpful.

Two possible paths forward here: 1) Put the details of the force calculations in the env step and pipe the relevant action value directly into the step function of the brax system. 2) Create a "passive" force class that does the requisite calculation, but which doesn't depend on user input in any way. This would need to be added to the config too, just like thruster is, now.

Pros for way 1) is that it could be easily accomplished right now by attaching thrusters to each body, and putting a little extra logic in the env step function. Cons for way 1) is that that force would be replicated across physics substeps instead of calculated different at each substep, so there might be a small accuracy penalty.

Pros for way 2) is that it removes all of the complexity from the env step, and puts it directly in a configurable widget that we can just attach to different bodies. Cons for way 2) is that it introduces a lot of complexity for a single configurable widget used by only one environment ;) .

It's probably worth trying way 1) for now, just because it's easier. Basically, all you'd need to do is pass in the requisite force values into the right indices of action of the system step. The force and actuator objects can report which indices correspond to which controllable objects, but, generically, forces come after actuators in the list. Let me know if that makes sense--happy to hop on a video call if it would be helpful.

[benelot]

We implemented a first version that is going to be in the repo soon.