gaussian_renderer 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (29) hide show
  1. gaussian_renderer-0.2.0/LICENSE +21 -0
  2. gaussian_renderer-0.2.0/PKG-INFO +127 -0
  3. gaussian_renderer-0.2.0/README.md +96 -0
  4. gaussian_renderer-0.2.0/pyproject.toml +76 -0
  5. gaussian_renderer-0.2.0/setup.cfg +4 -0
  6. gaussian_renderer-0.2.0/src/gaussian_renderer/__init__.py +35 -0
  7. gaussian_renderer-0.2.0/src/gaussian_renderer/batch_splat.py +221 -0
  8. gaussian_renderer-0.2.0/src/gaussian_renderer/core/__init__.py +28 -0
  9. gaussian_renderer-0.2.0/src/gaussian_renderer/core/batch_rasterization.py +350 -0
  10. gaussian_renderer-0.2.0/src/gaussian_renderer/core/gaussiandata.py +74 -0
  11. gaussian_renderer-0.2.0/src/gaussian_renderer/core/gs_renderer.py +207 -0
  12. gaussian_renderer-0.2.0/src/gaussian_renderer/core/super_splat_loader.py +678 -0
  13. gaussian_renderer-0.2.0/src/gaussian_renderer/core/util_gau.py +299 -0
  14. gaussian_renderer-0.2.0/src/gaussian_renderer/gs_renderer_motrixsim.py +134 -0
  15. gaussian_renderer-0.2.0/src/gaussian_renderer/gs_renderer_mujoco.py +162 -0
  16. gaussian_renderer-0.2.0/src/gaussian_renderer/simple_viewer.py +287 -0
  17. gaussian_renderer-0.2.0/src/gaussian_renderer/supersplat_compress.py +395 -0
  18. gaussian_renderer-0.2.0/src/gaussian_renderer/transform_gs_model.py +253 -0
  19. gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/PKG-INFO +127 -0
  20. gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/SOURCES.txt +27 -0
  21. gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/dependency_links.txt +1 -0
  22. gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/entry_points.txt +4 -0
  23. gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/requires.txt +27 -0
  24. gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/top_level.txt +1 -0
  25. gaussian_renderer-0.2.0/tests/test_gaussiandata.py +35 -0
  26. gaussian_renderer-0.2.0/tests/test_motrixsim_renderer.py +133 -0
  27. gaussian_renderer-0.2.0/tests/test_mujoco_renderer.py +113 -0
  28. gaussian_renderer-0.2.0/tests/test_ply_io.py +45 -0
  29. gaussian_renderer-0.2.0/tests/test_utils.py +36 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 Yufei Jia
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,127 @@
1
+ Metadata-Version: 2.4
2
+ Name: gaussian_renderer
3
+ Version: 0.2.0
4
+ Summary: A Gaussian Splatting Renderer and Tools
5
+ Requires-Python: >=3.10
6
+ Description-Content-Type: text/markdown
7
+ License-File: LICENSE
8
+ Requires-Dist: numpy
9
+ Requires-Dist: torch
10
+ Requires-Dist: scipy
11
+ Requires-Dist: plyfile
12
+ Requires-Dist: trimesh
13
+ Requires-Dist: gsplat
14
+ Provides-Extra: viewer
15
+ Requires-Dist: glfw; extra == "viewer"
16
+ Requires-Dist: PyOpenGL; extra == "viewer"
17
+ Provides-Extra: motrix
18
+ Requires-Dist: motrixsim>=0.5.0; extra == "motrix"
19
+ Provides-Extra: mujoco
20
+ Requires-Dist: mujoco; extra == "mujoco"
21
+ Provides-Extra: shs
22
+ Requires-Dist: einops; extra == "shs"
23
+ Requires-Dist: e3nn; extra == "shs"
24
+ Provides-Extra: dev
25
+ Requires-Dist: ruff; extra == "dev"
26
+ Requires-Dist: mypy; extra == "dev"
27
+ Requires-Dist: pytest; extra == "dev"
28
+ Requires-Dist: einops; extra == "dev"
29
+ Requires-Dist: e3nn; extra == "dev"
30
+ Dynamic: license-file
31
+
32
+ # Gaussian Renderer
33
+
34
+ A Gaussian Splatting Renderer and Tools package.
35
+
36
+ [中文文档](README_zh.md)
37
+
38
+ This repository primarily provides rendering capabilities for Gaussian Splatting models. It is developed as a component of the **DISCOVERSE** project.
39
+
40
+ For detailed usage within the simulation environment, see: [https://github.com/TATP-233/DISCOVERSE](https://github.com/TATP-233/DISCOVERSE)
41
+
42
+ ## Requirements
43
+
44
+ Python >= 3.10
45
+
46
+ ## Installation
47
+
48
+ ```bash
49
+ uv add gaussian-renderer
50
+ # or: pip install gaussian-renderer
51
+ ```
52
+
53
+ From source:
54
+ ```bash
55
+ git clone https://github.com/TATP-233/GaussainRenderer.git
56
+ cd GaussainRenderer
57
+ uv pip install .
58
+ # or: pip install .
59
+ ```
60
+
61
+ ### Optional extras
62
+
63
+ ```bash
64
+ uv add "gaussian-renderer[viewer]" # OpenGL viewer (glfw, PyOpenGL)
65
+ uv add "gaussian-renderer[mujoco]" # MuJoCo integration
66
+ uv add "gaussian-renderer[motrix]" # MotrixSim integration
67
+
68
+ # Combine as needed
69
+ uv add "gaussian-renderer[viewer,mujoco]"
70
+ # or: pip install ".[viewer,mujoco]"
71
+ ```
72
+
73
+ ## Usage
74
+
75
+ ### Command-line tools
76
+
77
+ **`gs-viewer`** — OpenGL viewer for `.ply` models
78
+ ```bash
79
+ gs-viewer path/to/model.ply
80
+ ```
81
+ Controls: Left mouse = rotate, Right/Middle = pan, Scroll = zoom, Up/Down = SH degree, Drag & drop = load file
82
+
83
+ **`gs-compress`** — Compress 3DGS PLY to SuperSplat format
84
+ ```bash
85
+ gs-compress input.ply
86
+ gs-compress input.ply -o output.ply
87
+ gs-compress models/ # batch
88
+ ```
89
+
90
+ **`gs-transform`** — Apply translation/rotation/scale to a model
91
+ ```bash
92
+ gs-transform input.ply -o output.ply -t 0 1 0 -s 2.0
93
+ gs-transform input.ply -r 0 0 0 1 # rotation quaternion xyzw
94
+ # --compress: save as compressed PLY
95
+ ```
96
+
97
+ ### Python API
98
+
99
+ ```bash
100
+ uv run python -m gaussian_renderer.simple_viewer path/to/model.ply
101
+ uv run python -m gaussian_renderer.supersplat_compress input.ply
102
+ uv run python -m gaussian_renderer.transform_gs_model input.ply
103
+ ```
104
+
105
+ ## Development
106
+
107
+ ```bash
108
+ uv pip install ".[dev]"
109
+ # or: pip install ".[dev]"
110
+ make lint # ruff check
111
+ make format # ruff format
112
+ make typecheck # mypy
113
+ make test # pytest
114
+ make ci # all of the above
115
+ ```
116
+
117
+ ## Citation
118
+
119
+ ```bibtex
120
+ @article{jia2025discoverse,
121
+ title={DISCOVERSE: Efficient Robot Simulation in Complex High-Fidelity Environments},
122
+ author={Yufei Jia and Guangyu Wang and Yuhang Dong and Junzhe Wu and Yupei Zeng and Haonan Lin and Zifan Wang and Haizhou Ge and Weibin Gu and Chuxuan Li and Ziming Wang and Yunjie Cheng and Wei Sui and Ruqi Huang and Guyue Zhou},
123
+ journal={arXiv preprint arXiv:2507.21981},
124
+ year={2025},
125
+ url={https://arxiv.org/abs/2507.21981}
126
+ }
127
+ ```
@@ -0,0 +1,96 @@
1
+ # Gaussian Renderer
2
+
3
+ A Gaussian Splatting Renderer and Tools package.
4
+
5
+ [中文文档](README_zh.md)
6
+
7
+ This repository primarily provides rendering capabilities for Gaussian Splatting models. It is developed as a component of the **DISCOVERSE** project.
8
+
9
+ For detailed usage within the simulation environment, see: [https://github.com/TATP-233/DISCOVERSE](https://github.com/TATP-233/DISCOVERSE)
10
+
11
+ ## Requirements
12
+
13
+ Python >= 3.10
14
+
15
+ ## Installation
16
+
17
+ ```bash
18
+ uv add gaussian-renderer
19
+ # or: pip install gaussian-renderer
20
+ ```
21
+
22
+ From source:
23
+ ```bash
24
+ git clone https://github.com/TATP-233/GaussainRenderer.git
25
+ cd GaussainRenderer
26
+ uv pip install .
27
+ # or: pip install .
28
+ ```
29
+
30
+ ### Optional extras
31
+
32
+ ```bash
33
+ uv add "gaussian-renderer[viewer]" # OpenGL viewer (glfw, PyOpenGL)
34
+ uv add "gaussian-renderer[mujoco]" # MuJoCo integration
35
+ uv add "gaussian-renderer[motrix]" # MotrixSim integration
36
+
37
+ # Combine as needed
38
+ uv add "gaussian-renderer[viewer,mujoco]"
39
+ # or: pip install ".[viewer,mujoco]"
40
+ ```
41
+
42
+ ## Usage
43
+
44
+ ### Command-line tools
45
+
46
+ **`gs-viewer`** — OpenGL viewer for `.ply` models
47
+ ```bash
48
+ gs-viewer path/to/model.ply
49
+ ```
50
+ Controls: Left mouse = rotate, Right/Middle = pan, Scroll = zoom, Up/Down = SH degree, Drag & drop = load file
51
+
52
+ **`gs-compress`** — Compress 3DGS PLY to SuperSplat format
53
+ ```bash
54
+ gs-compress input.ply
55
+ gs-compress input.ply -o output.ply
56
+ gs-compress models/ # batch
57
+ ```
58
+
59
+ **`gs-transform`** — Apply translation/rotation/scale to a model
60
+ ```bash
61
+ gs-transform input.ply -o output.ply -t 0 1 0 -s 2.0
62
+ gs-transform input.ply -r 0 0 0 1 # rotation quaternion xyzw
63
+ # --compress: save as compressed PLY
64
+ ```
65
+
66
+ ### Python API
67
+
68
+ ```bash
69
+ uv run python -m gaussian_renderer.simple_viewer path/to/model.ply
70
+ uv run python -m gaussian_renderer.supersplat_compress input.ply
71
+ uv run python -m gaussian_renderer.transform_gs_model input.ply
72
+ ```
73
+
74
+ ## Development
75
+
76
+ ```bash
77
+ uv pip install ".[dev]"
78
+ # or: pip install ".[dev]"
79
+ make lint # ruff check
80
+ make format # ruff format
81
+ make typecheck # mypy
82
+ make test # pytest
83
+ make ci # all of the above
84
+ ```
85
+
86
+ ## Citation
87
+
88
+ ```bibtex
89
+ @article{jia2025discoverse,
90
+ title={DISCOVERSE: Efficient Robot Simulation in Complex High-Fidelity Environments},
91
+ author={Yufei Jia and Guangyu Wang and Yuhang Dong and Junzhe Wu and Yupei Zeng and Haonan Lin and Zifan Wang and Haizhou Ge and Weibin Gu and Chuxuan Li and Ziming Wang and Yunjie Cheng and Wei Sui and Ruqi Huang and Guyue Zhou},
92
+ journal={arXiv preprint arXiv:2507.21981},
93
+ year={2025},
94
+ url={https://arxiv.org/abs/2507.21981}
95
+ }
96
+ ```
@@ -0,0 +1,76 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61.0"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "gaussian_renderer"
7
+ dynamic = ["version"]
8
+ description = "A Gaussian Splatting Renderer and Tools"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ dependencies = [
12
+ "numpy",
13
+ "torch",
14
+ "scipy",
15
+ "plyfile",
16
+ "trimesh",
17
+ "gsplat",
18
+ ]
19
+
20
+ [project.optional-dependencies]
21
+ viewer = [
22
+ "glfw",
23
+ "PyOpenGL",
24
+ ]
25
+
26
+ motrix = [
27
+ "motrixsim >= 0.5.0",
28
+ ]
29
+
30
+ mujoco = [
31
+ "mujoco",
32
+ ]
33
+
34
+ shs = [
35
+ "einops",
36
+ "e3nn",
37
+ ]
38
+
39
+ dev = [
40
+ "ruff",
41
+ "mypy",
42
+ "pytest",
43
+ "einops",
44
+ "e3nn",
45
+ ]
46
+
47
+ [project.scripts]
48
+ gs-viewer = "gaussian_renderer.simple_viewer:main"
49
+ gs-compress = "gaussian_renderer.supersplat_compress:main"
50
+ gs-transform = "gaussian_renderer.transform_gs_model:main"
51
+
52
+ [tool.setuptools.dynamic]
53
+ version = {attr = "gaussian_renderer.__version__"}
54
+
55
+ [tool.setuptools.packages.find]
56
+ where = ["src"]
57
+ include = ["gaussian_renderer*"]
58
+
59
+ [tool.ruff]
60
+ line-length = 120
61
+ target-version = "py310"
62
+
63
+ [tool.ruff.lint]
64
+ select = ["E", "F", "I"]
65
+ ignore = ["E501", "F401", "E402", "F841"]
66
+
67
+ [tool.mypy]
68
+ python_version = "3.10"
69
+ ignore_missing_imports = true
70
+ warn_unused_ignores = false
71
+ disable_error_code = ["assignment", "call-overload", "union-attr", "arg-type", "dict-item", "index", "return-value", "operator", "override", "type-arg", "attr-defined"]
72
+
73
+ [tool.pytest.ini_options]
74
+ testpaths = ["tests"]
75
+ python_files = ["test_*.py"]
76
+ pythonpath = ["src"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,35 @@
1
+ # SPDX-License-Identifier: MIT
2
+ #
3
+ # MIT License
4
+ #
5
+ # Copyright (c) 2025 Yufei Jia
6
+ #
7
+ # Permission is hereby granted, free of charge, to any person obtaining a copy
8
+ # of this software and associated documentation files (the "Software"), to deal
9
+ # in the Software without restriction, including without limitation the rights
10
+ # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11
+ # copies of the Software, and to permit persons to whom the Software is
12
+ # furnished to do so, subject to the following conditions:
13
+ #
14
+ # The above copyright notice and this permission notice shall be included in all
15
+ # copies or substantial portions of the Software.
16
+ #
17
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18
+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19
+ # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20
+ # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21
+ # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22
+ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23
+ # SOFTWARE.
24
+
25
+ from .batch_splat import BatchSplatConfig, BatchSplatRenderer, MjxBatchSplatRenderer, MtxBatchSplatRenderer
26
+ from .core.batch_rasterization import batch_env_render, batch_render, batch_update_gaussians
27
+ from .core.gaussiandata import GaussianBatchData, GaussianData
28
+ from .core.gs_renderer import GSRenderer
29
+ from .core.super_splat_loader import is_super_splat_format, load_super_splat_ply, save_super_splat_ply
30
+ from .core.util_gau import load_ply, save_ply, transform_shs
31
+ from .gs_renderer_motrixsim import GSRendererMotrixSim
32
+ from .gs_renderer_mujoco import GSRendererMuJoCo
33
+
34
+ __version__ = "0.2.0"
35
+ __author__ = "Yufei Jia"
@@ -0,0 +1,221 @@
1
+ # SPDX-License-Identifier: MIT
2
+ #
3
+ # MIT License
4
+ #
5
+ # Copyright (c) 2025 Yufei Jia
6
+ #
7
+ # Permission is hereby granted, free of charge, to any person obtaining a copy
8
+ # of this software and associated documentation files (the "Software"), to deal
9
+ # in the Software without restriction, including without limitation the rights
10
+ # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11
+ # copies of the Software, and to permit persons to whom the Software is
12
+ # furnished to do so, subject to the following conditions:
13
+ #
14
+ # The above copyright notice and this permission notice shall be included in all
15
+ # copies or substantial portions of the Software.
16
+ #
17
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18
+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19
+ # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20
+ # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21
+ # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22
+ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23
+ # SOFTWARE.
24
+
25
+ """
26
+ Minimal batch Gaussian splat renderer that works with mjx state tensors.
27
+
28
+ Usage pattern (3 steps):
29
+ 1) init: cfg = BatchSplatConfig(...); renderer = BatchSplatRenderer(cfg)
30
+ 2) loop: gsb = renderer.batch_update_gaussians(body_pos, body_quat)
31
+ 3) loop: rgb, depth = renderer.batch_env_render(gsb, cam_pos, cam_xmat, height, width, fovy)
32
+
33
+ Where body_pos/body_quat/cam_pos/cam_xmat can come from mjx (MuJoCo JAX) state via
34
+ `torch.utils.dlpack.from_dlpack(state.data.xxx)`.
35
+ """
36
+
37
+ from dataclasses import dataclass
38
+ from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple
39
+
40
+ import numpy as np
41
+ import torch
42
+ from torch import Tensor
43
+
44
+ from .core.batch_rasterization import (
45
+ batch_env_render as _batch_env_render,
46
+ )
47
+ from .core.batch_rasterization import (
48
+ batch_update_gaussians as _batch_update_gaussians,
49
+ )
50
+ from .core.gaussiandata import GaussianBatchData, GaussianData
51
+ from .core.util_gau import load_ply
52
+
53
+
54
+ @dataclass
55
+ class BatchSplatConfig:
56
+ # Mapping from MuJoCo body name -> PLY path (local coordinate frame).
57
+ body_gaussians: Dict[str, str]
58
+ # Optional static/background PLY that is not attached to any body.
59
+ background_ply: Optional[str] = None
60
+ # Device to place tensors on; if None, will pick CUDA if available else CPU.
61
+ device: Optional[torch.device] = None
62
+ # Minibatch size for rendering
63
+ minibatch: Optional[int] = None
64
+
65
+
66
+ class BatchSplatRenderer:
67
+ def __init__(self, cfg: BatchSplatConfig, body_name_to_id: Optional[Dict[str, int]]) -> None:
68
+ assert torch.cuda.is_available()
69
+ device = cfg.device or torch.device("cuda")
70
+
71
+ self.device = device
72
+ self.minibatch = cfg.minibatch
73
+
74
+ xyz_list: List[Tensor] = []
75
+ rot_list: List[Tensor] = []
76
+ scale_list: List[Tensor] = []
77
+ opacity_list: List[Tensor] = []
78
+ sh_list: List[Tensor] = []
79
+ max_sh_dim = 0
80
+
81
+ self.gs_idx_start: List[int] = []
82
+ self.gs_idx_end: List[int] = []
83
+ self.gs_body_ids: List[int] = []
84
+
85
+ # Load per-body gaussians
86
+ for body_name, ply_path in cfg.body_gaussians.items():
87
+ g = load_ply(ply_path)
88
+ max_sh_dim = max(max_sh_dim, g.sh.shape[1])
89
+ start = len(torch.cat(xyz_list)) if xyz_list else 0
90
+ end = start + len(g.xyz)
91
+ self.gs_idx_start.append(start)
92
+ self.gs_idx_end.append(end)
93
+ if body_name not in body_name_to_id:
94
+ raise ValueError(f"Body '{body_name}' not found in mj_model; available: {list(body_name_to_id.keys())}")
95
+ self.gs_body_ids.append(body_name_to_id[body_name])
96
+ xyz_list.append(torch.tensor(g.xyz, device=device, dtype=torch.float32))
97
+ rot_list.append(torch.tensor(g.rot, device=device, dtype=torch.float32))
98
+ scale_list.append(torch.tensor(g.scale, device=device, dtype=torch.float32))
99
+ opacity_list.append(torch.tensor(g.opacity, device=device, dtype=torch.float32))
100
+ sh_list.append(torch.tensor(g.sh, device=device, dtype=torch.float32))
101
+
102
+ # Optional background/static gaussian
103
+ if cfg.background_ply:
104
+ g = load_ply(cfg.background_ply)
105
+ max_sh_dim = max(max_sh_dim, g.sh.shape[1])
106
+ xyz_list.append(torch.tensor(g.xyz, device=device, dtype=torch.float32))
107
+ rot_list.append(torch.tensor(g.rot, device=device, dtype=torch.float32))
108
+ scale_list.append(torch.tensor(g.scale, device=device, dtype=torch.float32))
109
+ opacity_list.append(torch.tensor(g.opacity, device=device, dtype=torch.float32))
110
+ sh_list.append(torch.tensor(g.sh, device=device, dtype=torch.float32))
111
+
112
+ if sh_list:
113
+ for i, sh in enumerate(sh_list):
114
+ cur_dim = sh.shape[1]
115
+ if cur_dim < max_sh_dim:
116
+ pad = torch.zeros(sh.shape[0], max_sh_dim - cur_dim, device=sh.device, dtype=sh.dtype)
117
+ sh_list[i] = torch.cat([sh, pad], dim=1)
118
+
119
+ # Concatenate template
120
+ xyz_all = torch.cat(xyz_list, dim=0)
121
+ rot_all = torch.cat(rot_list, dim=0)
122
+ scale_all = torch.cat(scale_list, dim=0)
123
+ opacity_all = torch.cat(opacity_list, dim=0)
124
+ sh_all = torch.cat(sh_list, dim=0).reshape(xyz_all.shape[0], -1, 3).contiguous()
125
+
126
+ self.template = GaussianData(
127
+ xyz=xyz_all,
128
+ rot=rot_all,
129
+ scale=scale_all,
130
+ opacity=opacity_all,
131
+ sh=sh_all,
132
+ )
133
+
134
+ self.gs_idx_start = torch.tensor(self.gs_idx_start, dtype=torch.long, device=device)
135
+ self.gs_idx_end = torch.tensor(self.gs_idx_end, dtype=torch.long, device=device)
136
+ self.gs_body_ids = torch.tensor(self.gs_body_ids, dtype=torch.long, device=device)
137
+
138
+ # Precompute point-to-body mapping for vectorized update
139
+ num_points = len(self.template)
140
+ self.dynamic_mask = torch.zeros(num_points, dtype=torch.bool, device=device)
141
+ self.point_to_body_idx = torch.zeros(num_points, dtype=torch.long, device=device)
142
+
143
+ for i in range(len(self.gs_idx_start)):
144
+ start = self.gs_idx_start[i]
145
+ end = self.gs_idx_end[i]
146
+ body_id = self.gs_body_ids[i]
147
+ self.dynamic_mask[start:end] = True
148
+ self.point_to_body_idx[start:end] = body_id
149
+
150
+ def batch_update_gaussians(self, body_pos: Tensor, body_quat: Tensor, scalar_first=True) -> GaussianBatchData:
151
+ """Update gaussians using body poses.
152
+
153
+ Args:
154
+ body_pos: (Nenv, Nbody, 3) torch tensor
155
+ body_quat: (Nenv, Nbody, 4) torch tensor (wxyz)
156
+ Returns:
157
+ GaussianBatchData with per-env gaussians.
158
+ """
159
+ # Ensure device
160
+ if not isinstance(body_pos, torch.Tensor):
161
+ body_pos = torch.tensor(body_pos, device=self.device, dtype=torch.float32)
162
+ if not isinstance(body_quat, torch.Tensor):
163
+ body_quat = torch.tensor(body_quat, device=self.device, dtype=torch.float32)
164
+ return _batch_update_gaussians(
165
+ self.template,
166
+ body_pos,
167
+ body_quat,
168
+ point_to_body_idx=self.point_to_body_idx,
169
+ dynamic_mask=self.dynamic_mask,
170
+ scalar_first=scalar_first,
171
+ )
172
+
173
+ def batch_env_render(
174
+ self,
175
+ gsb: GaussianBatchData,
176
+ cam_pos: Tensor,
177
+ cam_xmat: Tensor,
178
+ height: int,
179
+ width: int,
180
+ fovy: np.ndarray,
181
+ bg_imgs: Optional[Tensor] = None,
182
+ ):
183
+ """Render RGBD for batch envs and cameras."""
184
+ if not isinstance(cam_pos, torch.Tensor):
185
+ cam_pos = torch.tensor(cam_pos, device=self.device, dtype=torch.float32)
186
+ if not isinstance(cam_xmat, torch.Tensor):
187
+ cam_xmat = torch.tensor(cam_xmat, device=self.device, dtype=torch.float32)
188
+ return _batch_env_render(gsb, cam_pos, cam_xmat, height, width, fovy, bg_imgs=bg_imgs, minibatch=self.minibatch)
189
+
190
+
191
+ if TYPE_CHECKING:
192
+ import mujoco
193
+
194
+
195
+ class MjxBatchSplatRenderer(BatchSplatRenderer):
196
+ def __init__(self, cfg: BatchSplatConfig, mj_model: "mujoco.MjModel") -> None:
197
+ body_name_to_id = {}
198
+ for i in range(mj_model.nbody):
199
+ body_name = mj_model.body(i).name
200
+ if body_name:
201
+ body_name_to_id[body_name] = i
202
+ super().__init__(cfg, body_name_to_id)
203
+
204
+
205
+ if TYPE_CHECKING:
206
+ import motrixsim
207
+
208
+
209
+ class MtxBatchSplatRenderer(BatchSplatRenderer):
210
+ def __init__(self, cfg, mx_model: "motrixsim.MotrixSimModel") -> None:
211
+ body_name_to_id = {}
212
+ for i, link_name in enumerate(mx_model.link_names):
213
+ if link_name:
214
+ body_name_to_id[link_name] = i
215
+ super().__init__(cfg, body_name_to_id)
216
+
217
+ def batch_update_gaussians(self, body_pos: Tensor, body_quat: Tensor) -> GaussianBatchData:
218
+ return super().batch_update_gaussians(body_pos, body_quat, scalar_first=False)
219
+
220
+
221
+ __all__ = ["BatchSplatConfig", "BatchSplatRenderer", "MjxBatchSplatRenderer", "MtxBatchSplatRenderer"]
@@ -0,0 +1,28 @@
1
+ # SPDX-License-Identifier: MIT
2
+ #
3
+ # MIT License
4
+ #
5
+ # Copyright (c) 2025 Yufei Jia
6
+ #
7
+ # Permission is hereby granted, free of charge, to any person obtaining a copy
8
+ # of this software and associated documentation files (the "Software"), to deal
9
+ # in the Software without restriction, including without limitation the rights
10
+ # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11
+ # copies of the Software, and to permit persons to whom the Software is
12
+ # furnished to do so, subject to the following conditions:
13
+ #
14
+ # The above copyright notice and this permission notice shall be included in all
15
+ # copies or substantial portions of the Software.
16
+ #
17
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18
+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19
+ # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20
+ # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21
+ # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22
+ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23
+ # SOFTWARE.
24
+
25
+ from .batch_rasterization import batch_env_render, batch_render, batch_update_gaussians
26
+ from .gaussiandata import GaussianBatchData, GaussianData
27
+ from .gs_renderer import GSRenderer
28
+ from .util_gau import load_ply, save_ply, transform_shs