Internal change

PiperOrigin-RevId: 603557707 Change-Id: Iaae8b6bb0d7c3a367f8e0b1862c438a9978a1902
google · Feb 2, 2024 · 5c4740e · 5c4740e
1 parent 74ad71d
commit 5c4740e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/brax/io/json.py b/brax/io/json.py
@@ -171,7 +171,8 @@ def dumps(sys: System, states: List[State]) -> Text:
     }
 
     if geom['name'] in ('Mesh', 'Box'):
-      vert, face = _get_mesh(sys.mj_model, sys.geom_dataid[id_])
+      # TODO: use sys.geom_dataid.
+      vert, face = _get_mesh(sys.mj_model, sys.mj_model.geom_dataid[id_])
       geom['vert'] = vert
       geom['face'] = face
 

diff --git a/brax/visualizer/visualizer.py b/brax/visualizer/visualizer.py
@@ -20,7 +20,7 @@
 
 from absl import app
 from absl import flags
-from brax.base import Actuator, Convex, Mesh, Plane, Transform
+from brax.base import Actuator
 from brax.generalized import pipeline as generalized
 from brax.io import html
 from brax.io import mjcf