Merge branch 'upaxis' into 'main'

Fix up axis handling

See merge request omniverse/warp!627

CHANGELOG.md

@@ -32,6 +32,7 @@
 - Allow passing external arrays/tensors to Warp kernels directly via `__cuda_array_interface__` and `__array_interface__`
 - Add faster Torch interop path using `return_ctype` argument to `wp.from_torch()`
 - Handle incompatible CUDA driver versions gracefully
+- Fix handling of `upaxis` variable in `ModelBuilder` and the rendering thereof in `OpenGLRenderer`
 
 ## [1.2.2] - 2024-07-04
 

warp/render/render_opengl.py

@@ -150,7 +150,7 @@
 layout (location = 2) in vec2 aTexCoord;
 
 uniform mat4 view;
-uniform mat4 model;
+uniform mat4 inv_model;
 uniform mat4 projection;
 uniform vec3 viewPos;
 
@@ -160,7 +160,8 @@
 void main()
 {
     vec4 worldPos = vec4(aPos + viewPos, 1.0);
-    gl_Position = projection * view * worldPos;
+    gl_Position = projection * view * inv_model * worldPos;
+
     FragPos = vec3(worldPos);
     TexCoord = aTexCoord;
 }
@@ -1078,6 +1079,7 @@ def __init__(
         self._scaling = scaling
 
         self._model_matrix = self.compute_model_matrix(self._camera_axis, scaling)
+        self._inv_model_matrix = np.linalg.inv(self._model_matrix.reshape(4, 4)).flatten()
         self.update_view_matrix(cam_pos=camera_pos, cam_front=camera_front, cam_up=camera_up)
         self.update_projection_matrix()
 
@@ -1216,7 +1218,7 @@ def __init__(
 
         with self._sky_shader:
             self._loc_sky_view = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("view"))
-            self._loc_sky_model = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("model"))
+            self._loc_sky_inv_model = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("inv_model"))
             self._loc_sky_projection = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("projection"))
 
             self._loc_sky_color1 = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("color1"))
@@ -1752,26 +1754,28 @@ def compute_model_matrix(camera_axis: int, scaling: float):
         if camera_axis == 0:
             return np.array((0, 0, scaling, 0, scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, 0, 1), dtype=np.float32)
         elif camera_axis == 2:
-            return np.array((-scaling, 0, 0, 0, 0, 0, scaling, 0, 0, scaling, 0, 0, 0, 0, 0, 1), dtype=np.float32)
+            return np.array((0, scaling, 0, 0, 0, 0, scaling, 0, scaling, 0, 0, 0, 0, 0, 0, 1), dtype=np.float32)
 
         return np.array((scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, 1), dtype=np.float32)
 
     def update_model_matrix(self, model_matrix: Optional[Mat44] = None):
         from pyglet import gl
 
-        # fmt: off
         if model_matrix is None:
             self._model_matrix = self.compute_model_matrix(self._camera_axis, self._scaling)
         else:
             self._model_matrix = np.array(model_matrix).flatten()
-        # fmt: on
+        self._inv_model_matrix = np.linalg.inv(self._model_matrix.reshape((4, 4))).flatten()
+        # update model view matrix in shaders
         ptr = arr_pointer(self._model_matrix)
         gl.glUseProgram(self._shape_shader.id)
         gl.glUniformMatrix4fv(self._loc_shape_model, 1, gl.GL_FALSE, ptr)
         gl.glUseProgram(self._grid_shader.id)
         gl.glUniformMatrix4fv(self._loc_grid_model, 1, gl.GL_FALSE, ptr)
+        # sky shader needs inverted model view matrix
         gl.glUseProgram(self._sky_shader.id)
-        gl.glUniformMatrix4fv(self._loc_sky_model, 1, gl.GL_FALSE, ptr)
+        inv_ptr = arr_pointer(self._inv_model_matrix)
+        gl.glUniformMatrix4fv(self._loc_sky_inv_model, 1, gl.GL_FALSE, inv_ptr)
 
     @property
     def num_tiles(self):

warp/sim/model.py

@@ -1245,7 +1245,7 @@ def __init__(self, up_vector=(0.0, 1.0, 0.0), gravity=-9.80665):
         self.joint_axis_total_count = 0
 
         self.up_vector = wp.vec3(up_vector)
-        self.up_axis = wp.vec3(np.argmax(np.abs(up_vector)))
+        self.up_axis = int(np.argmax(np.abs(up_vector)))
         self.gravity = gravity
         # indicates whether a ground plane has been created
         self._ground_created = False
@@ -4478,7 +4478,9 @@ def finalize(self, device=None, requires_grad=False) -> Model:
 
             # enable ground plane
             m.ground_plane = wp.array(self._ground_params["plane"], dtype=wp.float32, requires_grad=requires_grad)
-            m.gravity = np.array(self.up_vector) * self.gravity
+            m.gravity = np.array(self.up_vector, dtype=wp.float32) * self.gravity
+            m.up_axis = self.up_axis
+            m.up_vector = np.array(self.up_vector, dtype=wp.float32)
 
             m.enable_tri_collisions = False
 

warp/sim/render.py

@@ -67,13 +67,15 @@ def __init__(
             path,
             scaling=1.0,
             fps=60,
-            up_axis="Y",
+            up_axis=None,
             show_rigid_contact_points=False,
             contact_points_radius=1e-3,
             show_joints=False,
             **render_kwargs,
         ):
             # create USD stage
+            if up_axis is None:
+                up_axis = "XYZ"[model.up_axis]
             super().__init__(path, scaling=scaling, fps=fps, up_axis=up_axis, **render_kwargs)
             self.scaling = scaling
             self.cam_axis = "XYZ".index(up_axis.upper())