bevy_mujoco 0.0.3

Use MuJoCo physics and MCJF / URDF scenes in bevy
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

use bevy::{
    prelude::*,
    render::{mesh::Indices, render_resource::PrimitiveTopology},
};
use mujoco_rust::{Body, Geom, GeomType};
use trees::Tree;

use crate::mujoco_shape;

#[derive(Deref, DerefMut)]
pub struct BodyTree(pub Tree<Body>);

fn collect_children(parent_leaf: &mut Tree<Body>, bodies: &[Body]) {
    let parent_id = parent_leaf.data().id;
    let children: Vec<Body> = bodies
        .iter()
        .filter(|child| {
            // this will make world plane and bodies different entities
            // otherwise this would be a tree with one root — world plane
            child.parent_id == parent_id && child.id != child.parent_id && child.parent_id != 0
        })
        .cloned()
        .collect();

    for child in children {
        let mut leaf = Tree::<Body>::new(child.clone());
        collect_children(&mut leaf, bodies);
        parent_leaf.push_back(leaf);
    }
}

/// Returns tree of bodies
pub fn body_tree(bodies: &[Body]) -> Vec<BodyTree> {
    let mut trees: Vec<BodyTree> = vec![];

    let root_bodies = bodies.iter().filter(|body| body.parent_id == 0);

    for body in root_bodies {
        let mut root_leaf: BodyTree = BodyTree(Tree::new(body.clone()));
        collect_children(&mut root_leaf.0, bodies);

        trees.push(root_leaf);
    }

    trees
}

pub fn mujoco_mesh_2_bevy(mj_mesh: mujoco_rust::Mesh) -> Mesh {
    let mut mesh = Mesh::new(PrimitiveTopology::TriangleList);
    mesh.set_indices(Some(Indices::U32(mj_mesh.indices)));
    mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, mj_mesh.vertices);
    mesh.insert_attribute(Mesh::ATTRIBUTE_NORMAL, mj_mesh.normals);
    mesh
}

pub fn geom_mesh(geom: &Geom) -> Mesh {
    let size = &mut [geom.size.x, geom.size.y, geom.size.z];
    size.swap(1, 2);

    match geom.geom_type {
        GeomType::PLANE => {
            let plane_size = if size[0] > 0.0 {
                size[0]
            } else {
                // MuJoCo size 0 means infinite
                1e6
            };
            Mesh::from(shape::Plane {
                size: plane_size as f32,
            })
        }
        GeomType::BOX => Mesh::from(shape::Box::new(
            size[0] as f32,
            size[1] as f32,
            size[2] as f32,
        )),
        GeomType::SPHERE => Mesh::from(shape::Icosphere {
            radius: size[0] as f32,
            ..default()
        }),
        GeomType::CAPSULE => Mesh::from(shape::Capsule {
            radius: size[0] as f32,
            depth: (size[2] * 2.0) as f32,
            ..default()
        }),
        GeomType::ELLIPSOID => todo!(),
        GeomType::CYLINDER => Mesh::from(mujoco_shape::Cylinder {
            radius: size[0] as f32,
            height: (size[2] * 2.0) as f32,
            ..default()
        }),

        GeomType::MESH => mujoco_mesh_2_bevy(geom.mesh.clone().unwrap()),
        // --- NOT IMPLEMENTED ---
        _ => todo!(),
    }
}

pub fn geom_rotation(geom: &Geom) -> Quat {
    match geom.geom_type {
        GeomType::MESH => Quat::from_xyzw(
            geom.quat[1] as f32,
            geom.quat[2] as f32,
            geom.quat[3] as f32,
            geom.quat[0] as f32,
        ),
        _ => Quat::from_xyzw(
            geom.quat[1] as f32,
            geom.quat[3] as f32,
            geom.quat[2] as f32,
            -geom.quat[0] as f32,
        ),
    }
}

/// bevy and mujoco treat object frame differently, this function converts
pub fn geom_correction(geom: &Geom) -> Vec3 {
    let size = &mut [geom.size.x, geom.size.y, geom.size.z];
    size.swap(1, 2);
    match geom.geom_type {
        GeomType::BOX => Vec3::new(0.0, (size[1] * 2.0) as f32, 0.0),
        GeomType::CAPSULE => Vec3::new(0.0, (size[1] * 2.0) as f32, 0.0),
        GeomType::CYLINDER => Vec3::new(0.0, (size[2] * 2.0) as f32, 0.0),
        _ => Vec3::ZERO,
    }
}

pub fn geom_translation(geom: &Geom) -> Vec3 {
    Vec3::new(geom.pos[0] as f32, geom.pos[1] as f32, geom.pos[2] as f32) - geom_correction(geom)
}

pub fn geom_transform(geom: &Geom) -> Transform {
    Transform {
        translation: geom_translation(geom),
        rotation: geom_rotation(geom),
        ..default()
    }
}

pub fn body_rotation(body: &Body) -> Quat {
    Quat::from_xyzw(
        body.quat[1] as f32,
        body.quat[2] as f32,
        body.quat[3] as f32,
        body.quat[0] as f32,
    )
}

pub fn body_translation(body: &Body) -> Vec3 {
    Vec3::new(body.pos[0] as f32, body.pos[1] as f32, body.pos[2] as f32)
}

pub fn body_transform(body: &Body) -> Transform {
    Transform {
        translation: body_translation(body),
        rotation: body_rotation(body),
        ..default()
    }
}