Ant with a trap of boxes (similar to issue #136)

Hi!

I am trying to extend the [Ant](https://github.com/google/brax/blob/main/brax/envs/ant.py) basic brax locomotion environment by simply adding a Trap in front of the Ant.

To do so, I took the _SYSTEM_CONFIG of the ant and added a new body called Trap with three box colliders. This alone works. Now when I add the collide_include statement to have collisions between the ant body and the trap, I am having a warning `/usr/local/lib/python3.7/dist-packages/brax/physics/colliders.py:638: RuntimeWarning: divide by zero encountered in true_divide
t = jp.dot(normal, (p0 - a) / jp.abs(jp.dot(normal, capsule_normal)))`, even when the ant and the trap are far from each other.

To visualise the interactions, I made sure the Ant would have its starting position above the Trap so it would fall on the Trap. The result is that the Ant is ejected very far (into the sky...).

As this issue is very similar to #136 , I tried to remove the frozen statement and play with mass and inertia (just like described in issue #136) but none of my several tries ended up in acceptable physical behavior.

Below is the config file I used. Does this config file makes sense or did I made something wrong? (The differences with the original one are: the new body named "Trap", the include_collision statement between the body parts of the Ant and the "Trap", the include_collision statement between the "Trap" and the "Ground" and finally the defaults position of the Ant above the Trap to force collision between them.)

In case the file is good, I guess this use case definitely join issue #136; in which case it would be great to solve this issue so we can get to write plenty more environment with Brax ;)

Thanks!

PS: When replacing the boxes with capsules to construct my Trap, I was able to get acceptable physical behaviors.

```
_ANTTRAP_SYSTEM_CONFIG = """
bodies {
name: "$ Torso"
colliders {
capsule {
radius: 0.25
length: 0.5
end: 1
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 10
}
bodies {
name: "Aux 1"
colliders {
rotation { x: 90 y: -45 }
capsule {
radius: 0.08
length: 0.4428427219390869
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
}
bodies {
name: "$ Body 4"
colliders {
rotation { x: 90 y: -45 }
capsule {
radius: 0.08
length: 0.7256854176521301
end: -1
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
}
bodies {
name: "Aux 2"
colliders {
rotation { x: 90 y: 45 }
capsule {
radius: 0.08
length: 0.4428427219390869
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
}
bodies {
name: "$ Body 7"
colliders {
rotation { x: 90 y: 45 }
capsule {
radius: 0.08
length: 0.7256854176521301
end: -1
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
}
bodies {
name: "Aux 3"
colliders {
rotation { x: -90 y: 45 }
capsule {
radius: 0.08
length: 0.4428427219390869
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
}
bodies {
name: "$ Body 10"
colliders {
rotation { x: -90 y: 45 }
capsule {
radius: 0.08
length: 0.7256854176521301
end: -1
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
}
bodies {
name: "Aux 4"
colliders {
rotation { x: -90 y: -45 }
capsule {
radius: 0.08
length: 0.4428427219390869
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
}
bodies {
name: "$ Body 13"
colliders {
rotation { x: -90 y: -45 }
capsule {
radius: 0.08
length: 0.7256854176521301
end: -1
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
}
bodies {
name: "Ground"
colliders {
plane {}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
frozen { all: true }
}
bodies {
name: "Trap"
colliders {
position { x: 10.0 y: 0.0 z: 1.0 }
box {
halfsize { x: 0.5 y: 2.5 z: 0.5 }
}
}
colliders {
position { x: 8.5 y: 3.0 z: 1.0 }
box {
halfsize { x: 2.0 y: 0.5 z: 0.5 }
}
}
colliders {
position { x: 8.5 y: -3.0 z: 1.0 }
box {
halfsize { x: 2.0 y: 0.5 z: 0.5 }
}
}
inertia { x: 1.0 y: 1.0 z: 1.0 }
mass: 1
frozen { all: true }
}
joints {
name: "$ Torso_Aux 1"
parent_offset { x: 0.2 y: 0.2 }
child_offset { x: -0.1 y: -0.1 }
parent: "$ Torso"
child: "Aux 1"
stiffness: 18000.0
angular_damping: 20
spring_damping: 80
angle_limit { min: -30.0 max: 30.0 }
rotation { y: -90 }
}
joints {
name: "Aux 1_$ Body 4"
parent_offset { x: 0.1 y: 0.1 }
child_offset { x: -0.2 y: -0.2 }
parent: "Aux 1"
child: "$ Body 4"
stiffness: 18000.0
angular_damping: 20
spring_damping: 80
rotation: { z: 135 }
angle_limit {
min: 30.0
max: 70.0
}
}
joints {
name: "$ Torso_Aux 2"
parent_offset { x: -0.2 y: 0.2 }
child_offset { x: 0.1 y: -0.1 }
parent: "$ Torso"
child: "Aux 2"
stiffness: 18000.0
angular_damping: 20
spring_damping: 80
rotation { y: -90 }
angle_limit { min: -30.0 max: 30.0 }
}
joints {
name: "Aux 2_$ Body 7"
parent_offset { x: -0.1 y: 0.1 }
child_offset { x: 0.2 y: -0.2 }
parent: "Aux 2"
child: "$ Body 7"
stiffness: 18000.0
angular_damping: 20
spring_damping: 80
rotation { z: 45 }
angle_limit { min: -70.0 max: -30.0 }
}
joints {
name: "$ Torso_Aux 3"
parent_offset { x: -0.2 y: -0.2 }
child_offset { x: 0.1 y: 0.1 }
parent: "$ Torso"
child: "Aux 3"
stiffness: 18000.0
angular_damping: 20
spring_damping: 80
rotation { y: -90 }
angle_limit { min: -30.0 max: 30.0 }
}
joints {
name: "Aux 3_$ Body 10"
parent_offset { x: -0.1 y: -0.1 }
child_offset {
x: 0.2
y: 0.2
}
parent: "Aux 3"
child: "$ Body 10"
stiffness: 18000.0
angular_damping: 20
spring_damping: 80
rotation { z: 135 }
angle_limit { min: -70.0 max: -30.0 }
}
joints {
name: "$ Torso_Aux 4"
parent_offset { x: 0.2 y: -0.2 }
child_offset { x: -0.1 y: 0.1 }
parent: "$ Torso"
child: "Aux 4"
stiffness: 18000.0
angular_damping: 20
spring_damping: 80
rotation { y: -90 }
angle_limit { min: -30.0 max: 30.0 }
}
joints {
name: "Aux 4_$ Body 13"
parent_offset { x: 0.1 y: -0.1 }
child_offset { x: -0.2 y: 0.2 }
parent: "Aux 4"
child: "$ Body 13"
stiffness: 18000.0
angular_damping: 20
spring_damping: 80
rotation { z: 45 }
angle_limit { min: 30.0 max: 70.0 }
}
actuators {
name: "$ Torso_Aux 1"
joint: "$ Torso_Aux 1"
strength: 350.0
torque {}
}
actuators {
name: "Aux 1_$ Body 4"
joint: "Aux 1_$ Body 4"
strength: 350.0
torque {}
}
actuators {
name: "$ Torso_Aux 2"
joint: "$ Torso_Aux 2"
strength: 350.0
torque {}
}
actuators {
name: "Aux 2_$ Body 7"
joint: "Aux 2_$ Body 7"
strength: 350.0
torque {}
}
actuators {
name: "$ Torso_Aux 3"
joint: "$ Torso_Aux 3"
strength: 350.0
torque {}
}
actuators {
name: "Aux 3_$ Body 10"
joint: "Aux 3_$ Body 10"
strength: 350.0
torque {}
}
actuators {
name: "$ Torso_Aux 4"
joint: "$ Torso_Aux 4"
strength: 350.0
torque {}
}
actuators {
name: "Aux 4_$ Body 13"
joint: "Aux 4_$ Body 13"
strength: 350.0
torque {}
}
friction: 1.0
gravity { z: -9.8 }
angular_damping: -0.05
baumgarte_erp: 0.1
collide_include {
first: "$ Torso"
second: "Ground"
}
collide_include {
first: "$ Body 4"
second: "Ground"
}
collide_include {
first: "$ Body 7"
second: "Ground"
}
collide_include {
first: "$ Body 10"
second: "Ground"
}
collide_include {
first: "$ Body 13"
second: "Ground"
}
collide_include {
first: "$ Torso"
second: "Trap"
}
collide_include {
first: "$ Body 4"
second: "Trap"
}
collide_include {
first: "$ Body 7"
second: "Trap"
}
collide_include {
first: "$ Body 10"
second: "Trap"
}
collide_include {
first: "$ Body 13"
second: "Trap"
}
collide_include {
first: "Trap"
second: "Ground"
}
dt: 0.05
substeps: 10
defaults {
qps {
name: "$ Torso"
pos { x: 10.0 y: 0.0 z: 3.5 }
}
}
"""
```