Internal change

PiperOrigin-RevId: 415374066 Change-Id: Id567f1d181c724105be1a983a052d054384731f8
google · Dec 10, 2021 · 4672000 · 4672000
1 parent 24773b5
commit 4672000
Show file tree

Hide file tree

Showing 37 changed files with 857 additions and 225 deletions.
diff --git a/brax/envs/ant.py b/brax/envs/ant.py
@@ -225,8 +225,9 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   child_offset { x: -0.1 y: -0.1 }
   parent: "$ Torso"
   child: "Aux 1"
-  stiffness: 5000.0
-  angular_damping: 35
+  stiffness: 18000.0
+  angular_damping: 20
+  spring_damping: 80
   angle_limit { min: -30.0 max: 30.0 }
   rotation { y: -90 }
 }
@@ -236,8 +237,9 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   child_offset { x: -0.2 y: -0.2 }
   parent: "Aux 1"
   child: "$ Body 4"
-  stiffness: 5000.0
-  angular_damping: 35
+  stiffness: 18000.0
+  angular_damping: 20
+  spring_damping: 80
   rotation: { z: 135 }
   angle_limit {
     min: 30.0
@@ -250,8 +252,9 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   child_offset { x: 0.1 y: -0.1 }
   parent: "$ Torso"
   child: "Aux 2"
-  stiffness: 5000.0
-  angular_damping: 35
+  stiffness: 18000.0
+  angular_damping: 20
+  spring_damping: 80
   rotation { y: -90 }
   angle_limit { min: -30.0 max: 30.0 }
 }
@@ -261,8 +264,9 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   child_offset { x: 0.2 y: -0.2 }
   parent: "Aux 2"
   child: "$ Body 7"
-  stiffness: 5000.0
-  angular_damping: 35
+  stiffness: 18000.0
+  angular_damping: 20
+  spring_damping: 80
   rotation { z: 45 }
   angle_limit { min: -70.0 max: -30.0 }
 }
@@ -272,8 +276,9 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   child_offset { x: 0.1 y: 0.1 }
   parent: "$ Torso"
   child: "Aux 3"
-  stiffness: 5000.0
-  angular_damping: 35
+  stiffness: 18000.0
+  angular_damping: 20
+  spring_damping: 80
   rotation { y: -90 }
   angle_limit { min: -30.0 max: 30.0 }
 }
@@ -286,8 +291,9 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   }
   parent: "Aux 3"
   child: "$ Body 10"
-  stiffness: 5000.0
-  angular_damping: 35
+  stiffness: 18000.0
+  angular_damping: 20
+  spring_damping: 80
   rotation { z: 135 }
   angle_limit { min: -70.0 max: -30.0 }
 }
@@ -297,8 +303,9 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   child_offset { x: -0.1 y: 0.1 }
   parent: "$ Torso"
   child: "Aux 4"
-  stiffness: 5000.0
-  angular_damping: 35
+  stiffness: 18000.0
+  angular_damping: 20
+  spring_damping: 80
   rotation { y: -90 }
   angle_limit { min: -30.0 max: 30.0 }
 }
@@ -308,8 +315,9 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   child_offset { x: -0.2 y: 0.2 }
   parent: "Aux 4"
   child: "$ Body 13"
-  stiffness: 5000.0
-  angular_damping: 35
+  stiffness: 18000.0
+  angular_damping: 20
+  spring_damping: 80
   rotation { z: 45 }
   angle_limit { min: 30.0 max: 70.0 }
 }

diff --git a/brax/envs/halfcheetah.py b/brax/envs/halfcheetah.py
@@ -261,7 +261,7 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
 }
 joints {
   name: "bthigh"
-  stiffness: 8000.0
+  stiffness: 25000
   parent: "torso"
   child: "bthigh"
   parent_offset {
@@ -272,16 +272,16 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   rotation {
     z: 90.0
   }
-  spring_damping: 100.0
+  spring_damping: 60
   angle_limit {
     min: -29.793806076049805
     max: 60.16056823730469
   }
-  limit_strength: 2000.0
-}
+  limit_strength: 1000
+  }
 joints {
   name: "bshin"
-  stiffness: 5000.0
+  stiffness: 25000
   parent: "bthigh"
   child: "bshin"
   parent_offset {
@@ -293,16 +293,16 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   rotation {
     z: 90.0
   }
-  spring_damping: 100.0
+  spring_damping: 60
   angle_limit {
     min: -44.97718811035156
     max: 44.97718811035156
   }
-  limit_strength: 1200.0
-}
+  limit_strength: 1000
+  }
 joints {
   name: "bfoot"
-  stiffness: 5000.0
+  stiffness: 25000
   parent: "bshin"
   child: "bfoot"
   parent_offset {
@@ -314,16 +314,16 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   rotation {
     z: 90.0
   }
-  spring_damping: 100.0
+  spring_damping: 60
   angle_limit {
     min: -22.918312072753906
     max: 44.97718811035156
   }
-  limit_strength: 400.0
-}
+  limit_strength: 1000
+  }
 joints {
   name: "fthigh"
-  stiffness: 8000.0
+  stiffness: 25000
   parent: "torso"
   child: "fthigh"
   parent_offset {
@@ -334,16 +334,16 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
   rotation {
     z: 90.0
   }
-  spring_damping: 100.0
+  spring_damping: 60
   angle_limit {
     min: -57.295780181884766
     max: 40.1070442199707
   }
-  limit_strength: 2000.0
-}
+  limit_strength: 1000
+  }
 joints {
   name: "fshin"
-  stiffness: 5000.0
+  stiffness: 25000
   parent: "fthigh"
   child: "fshin"
   parent_offset {
@@ -360,11 +360,11 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
     min: -68.75493621826172
     max: 49.847328186035156
   }
-  limit_strength: 400.0
-}
+  limit_strength: 1000
+  }
 joints {
   name: "ffoot"
-  stiffness: 3500.0
+  stiffness: 25000
   parent: "fshin"
   child: "ffoot"
   parent_offset {
@@ -381,8 +381,8 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
     min: -28.647890090942383
     max: 28.647890090942383
   }
-  limit_strength: 200.0
-}
+  limit_strength: 1000
+  }
 actuators {
   name: "bthigh"
   joint: "bthigh"
@@ -400,7 +400,7 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
 actuators {
   name: "bfoot"
   joint: "bfoot"
-  strength: 60.0
+  strength: 60
   torque {
   }
 }
@@ -414,7 +414,7 @@ def _get_obs(self, qp: brax.QP, info: brax.Info) -> jp.ndarray:
 actuators {
   name: "fshin"
   joint: "fshin"
-  strength: 60.0
+  strength: 60
   torque {
   }
 }