[BUG FIX] Fix mixed precision issue.

[abhijitmajumdar]

Description

When running camera renders across multiple environments with camera attached to a link (say the robot EE), the following error shows up

Traceback (most recent call last):
  File "/home/abhijit/.../main.py", line 991, in <module>
    main(args)
  File "/home/abhijit/.../main.py", line 923, in main
    pnp.step()
  File "/home/abhijit/.../main.py", line 637, in step
    self._scene.step()
  File "/home/abhijit/.../Genesis/genesis/utils/misc.py", line 144, in wrapper
    return method(self, *args, **kwargs)
  File "/home/abhijit/.../Genesis/genesis/engine/scene.py", line 767, in step
    self._visualizer.update(force=False, auto=refresh_visualizer)
  File "/home/abhijit/.../Genesis/genesis/vis/visualizer.py", line 166, in update
    self._viewer.update(auto_refresh=auto)
  File "/home/abhijit/.../Genesis/genesis/vis/viewer.py", line 143, in update
    self._pyrender_viewer.pending_buffer_updates |= self.context.update()
  File "/home/abhijit/.../Genesis/genesis/vis/rasterizer_context.py", line 818, in update
    self.visualizer.update_visual_states()
  File "/home/abhijit/.../Genesis/genesis/vis/visualizer.py", line 179, in update_visual_states
    camera.move_to_attach()
  File "/home/abhijit/.../Genesis/genesis/utils/misc.py", line 144, in wrapper
    return method(self, *args, **kwargs)
  File "/home/abhijit/.../Genesis/genesis/vis/camera.py", line 192, in move_to_attach
    link_T = gu.trans_quat_to_T(link_pos, link_quat)
  File "/home/abhijit/.../Genesis/genesis/utils/geom.py", line 872, in trans_quat_to_T
    quat_to_R(quat, out=T[..., :3, :3])
  File "/home/abhijit/.../Genesis/genesis/utils/geom.py", line 610, in quat_to_R
    return _np_quat_to_R(quat, out=out)
  File "/home/abhijit/.../.venv/lib/python3.10/site-packages/numba/core/dispatcher.py", line 424, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/abhijit/.../.venv/lib/python3.10/site-packages/numba/core/dispatcher.py", line 365, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Cannot unify array(float32, 2d, C) and array(float64, 2d, A) for 'out_.2', defined at /home/abhijit/.../Genesis/genesis/utils/geom.py (557)

File "Genesis/genesis/utils/geom.py", line 557:
def _np_quat_to_R(quat: np.ndarray, out: np.ndarray | None = None) -> np.ndarray:
    <source elided>

    s = 2.0 / np.sum(np.square(quat), -1)
    ^

During: typing of assignment at /home/abhijit/.../Genesis/genesis/utils/geom.py (557)

File "Genesis/genesis/utils/geom.py", line 557:
def _np_quat_to_R(quat: np.ndarray, out: np.ndarray | None = None) -> np.ndarray:
    <source elided>

    s = 2.0 / np.sum(np.square(quat), -1)
    ^

During: Pass nopython_type_inference

[Genesis] [14:00:52] [ERROR] TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Cannot unify array(float32, 2d, C) and array(float64, 2d, A) for 'out_.2', defined at /home/abhijit/.../Genesis/genesis/utils/geom.py (557)

File "Genesis/genesis/utils/geom.py", line 557:
def _np_quat_to_R(quat: np.ndarray, out: np.ndarray | None = None) -> np.ndarray:
    <source elided>

    s = 2.0 / np.sum(np.square(quat), -1)
    ^

During: typing of assignment at /home/abhijit/.../Genesis/genesis/utils/geom.py (557)

File "Genesis/genesis/utils/geom.py", line 557:
def _np_quat_to_R(quat: np.ndarray, out: np.ndarray | None = None) -> np.ndarray:
    <source elided>

    s = 2.0 / np.sum(np.square(quat), -1)
    ^

During: Pass nopython_type_inference
[Genesis] [14:00:52] [INFO] 💤 Exiting Genesis and caching compiled kernels...

This is a result of a feature I previously added and subsequent changes added to Genesis, and only triggers in the aforementioned scenario.

TLDR:

• move_to_attach is called when camera attached to a link is being rendered
• In case of env_idx, the pos and quat are of type np.float32 (default, as specified by the simulation precision) here, however envs_offset are returned as np.float64, which causes the final link_pos and link_quat passed to gu.trans_quat_to_T to have different types
• Internally trans_quat_to_T creates a T matrix passed down as out which is generated using pos.dtype, but later when quat derivates are added to T the above error occurs due to type mismatch (see traceback above to geom.py::trans_quat_to_T -> quat_to_R -> _np_quat_to_R)

Related Issue

Resolves Genesis-Embodied-AI/Genesis#

How Has This Been / Can This Be Tested?

Same steps described in #1323

Checklist:

• I read the CONTRIBUTING document.
• I followed the Submitting Code Changes section of CONTRIBUTING document.
• I tagged the title correctly (including BUG FIX/FEATURE/MISC/BREAKING)
• I updated the documentation accordingly or no change is needed.
• I tested my changes and added instructions on how to test it for reviewers.

• I have added tests to cover my changes.
• All new and existing tests passed.

[duburcqa]

Could you rather modify quat_to_R to force casting of input if the output is already defined for numpy ? This would be more generic.

[duburcqa]

I will try to clean up first the codebase to avoid mixing types in computations. Then we can merge this one. Because casting would slow down performance.

[duburcqa]

Superseded by #1395

[duburcqa]

After thinking twice, I would rather make sure that we do no mix different precision instead of hiding this mistake by cast the input arguments.

[abhijitmajumdar]

Superseded by #1395

@duburcqa The reference links back to this PR (I assume a typo). Could you link the correct one (I would like to keep track of when that gets merged in). Thanks

[duburcqa]

Superseded by #1398

[duburcqa]

It should be merged in less than 1 hour. Waiting for the CI to pass!

[duburcqa]

Thank you for your efforts!