Internal change

PiperOrigin-RevId: 603479877 Change-Id: I1782a8de6c8ca35e9b0182708ea5bfe8acae38ef
google · Feb 1, 2024 · dd37045 · dd37045
1 parent ec075da
commit dd37045
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 11 deletions.
diff --git a/brax/contact.py b/brax/contact.py
@@ -35,7 +35,8 @@ def get(sys: System, x: Transform) -> Optional[Contact]:
   Returns:
     Contact pytree
   """
-  ncon = mjx.ncon(sys)
+  # TODO: use mjx.ncon.
+  ncon = mjx._src.collision_driver.ncon(sys)
   if not ncon:
     return None
 

diff --git a/brax/envs/env_test.py b/brax/envs/env_test.py
@@ -57,16 +57,6 @@ def testSpeed(self, backend, env_name, expected_sps):
     )
     self.assertGreater(mean_sps, expected_sps * 0.99)
 
-  @parameterized.parameters(['mjx', 'generalized', 'spring', 'positional'])
-  def test_render(self, backend):
-    env = envs.create(
-        'ant',
-        backend=backend,
-    )
-    state = jax.jit(env.reset)(jax.random.PRNGKey(0))
-    images = env.render([state.pipeline_state])
-    self.assertLen(images, 1)
-    self.assertEqual(images[0].shape, (240, 320, 3))
 
 
 if __name__ == '__main__':