Skip to content

sim.sensors

crazyflow.sim.sensors

Classes

Functions

build_render_depth_fn(mjx_model, camera=0, resolution=(100, 100), geomgroup=(1, 1, 0, 0, 1, 1, 1, 1))

Build a depth renderer function for given camera and resolution.

Compiles the mjx model and rays directly into the rendering function for higher performance. The returned function takes a Sim object as input and returns depth images.

Source code in crazyflow/sim/sensors.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def build_render_depth_fn(
    mjx_model: mjx.Model,
    camera: int = 0,
    resolution: tuple[int, int] = (100, 100),
    geomgroup: tuple[int, ...] = (1, 1, 0, 0, 1, 1, 1, 1),
) -> Callable[[Sim], Array]:
    """Build a depth renderer function for given camera and resolution.

    Compiles the mjx model and rays directly into the rendering function for higher performance. The
    returned function takes a Sim object as input and returns depth images.
    """
    rays = _camera_rays(resolution=resolution, fov_y=jnp.pi / 4)[None, ...]
    ray_fn = jax.jit(
        partial(_render_rays, mjx_model=mjx_model, camera=camera, geomgroup=geomgroup, rays=rays),
        static_argnames=("mjx_model", "camera", "geomgroup", "rays"),
    )

    @requires_mujoco_sync
    def render_depth_fn(sim: Sim) -> Array:
        return ray_fn(mjx_data=sim.mjx_data)

    return render_depth_fn

render_depth(sim, camera=0, resolution=(100, 100), include_drone=False)

Render depth images using raycasting.

Note

Code has been adoped from https://github.com/Andrew-Luo1/jax_shac/blob/main/vision/2dof_ball.ipynb

Source code in crazyflow/sim/sensors.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
@requires_mujoco_sync
def render_depth(
    sim: Sim, camera: int = 0, resolution: tuple[int, int] = (100, 100), include_drone: bool = False
) -> Array:
    """Render depth images using raycasting.

    Note:
        Code has been adoped from
        https://github.com/Andrew-Luo1/jax_shac/blob/main/vision/2dof_ball.ipynb
    """
    return _render_depth(
        mjx_data=sim.mjx_data,
        mjx_model=sim.mjx_model,
        camera=camera,
        resolution=resolution,
        include_drone=include_drone,
    )