newton/examples/example_anymal_c_walk_physx_policy.py (3)

41-45: Fix docstring formatting.

The docstring has formatting issues with missing spaces and line breaks.

-def quat_rotate_inverse(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-    """Rotate a vector by the inverse of a quaternion along the last dimension of q and v.    Args:
-    q: The quaternion in (x, y, z, w). Shape is (..., 4).
-    v: The vector in (x, y, z). Shape is (..., 3).    Returns:
-    The rotated vector in (x, y, z). Shape is (..., 3).
-    """
+def quat_rotate_inverse(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """Rotate a vector by the inverse of a quaternion along the last dimension of q and v.
+    
+    Args:
+        q: The quaternion in (x, y, z, w). Shape is (..., 4).
+        v: The vector in (x, y, z). Shape is (..., 3).
+    
+    Returns:
+        The rotated vector in (x, y, z). Shape is (..., 3).
+    """

---

200-203: Consider implementing the suggested optimization.

The inline comment correctly identifies an optimization opportunity. Instead of launching apply_mjc_control_kernel at every simulation step, you could update positions and velocity targets only at the policy inference rate.

Would you like me to help implement this optimization by directly updating self.solver.mjw_data.ctrl[0] and avoiding unnecessary kernel launches?

---

236-236: Remove whitespace from blank line.

-        example = Example(stage_path=args.stage_path, headless=args.headless)
-        
+        example = Example(stage_path=args.stage_path, headless=args.headless)
+
         # Download the policy from the newton-assets repository

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

newton/examples/example_anymal_c_walk_physx_policy.py (1)

26-26: Remove unused import.

The os module is imported but never used in the code.

-import os
-
 import torch

newton/examples/example_anymal_c_walk_physx_policy.py (3)

39-45: Fix docstring formatting.

The docstring has formatting issues - the Args and Returns sections are not properly structured.

 @torch.jit.script
 def quat_rotate_inverse(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-    """Rotate a vector by the inverse of a quaternion along the last dimension of q and v.    Args:
-    q: The quaternion in (x, y, z, w). Shape is (..., 4).
-    v: The vector in (x, y, z). Shape is (..., 3).    Returns:
-    The rotated vector in (x, y, z). Shape is (..., 3).
+    """Rotate a vector by the inverse of a quaternion along the last dimension of q and v.
+    
+    Args:
+        q: The quaternion in (x, y, z, w). Shape is (..., 4).
+        v: The vector in (x, y, z). Shape is (..., 3).
+    
+    Returns:
+        The rotated vector in (x, y, z). Shape is (..., 3).
     """

---

200-202: Consider implementing the suggested performance optimization.

The inline comment mentions a valid performance optimization: avoiding repeated kernel launches by updating control targets only at the policy frequency rather than every simulation step.

Would you like me to help implement this optimization as mentioned in the PR objectives?

---

236-236: Remove whitespace from blank line.

The blank line contains unnecessary whitespace.

         example = Example(stage_path=args.stage_path, headless=args.headless)
-        
+
         # Download the policy from the newton-assets repository

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro