Skip to content

env

env

Modules:

Name Description
abstract_sensors
arena
camera_manager

Camera management for MolmoSpaces environments.
data_views
env
mj_extensions
object_manager
rby1_sensors
sensors
sensors_cameras

abstract_sensors

Classes:

Name Description
Sensor

Represents a sensor that provides data from the environment to agent.
SensorSuite

Represents a set of sensors, with each sensor being identified through a

Attributes:

Name Type Description
SpaceDict

SpaceDict module-attribute 

SpaceDict = Dict

Sensor 

Sensor(uuid: str, observation_space: Space, **kwargs: Any)

Bases: ABC

Represents a sensor that provides data from the environment to agent. The user of this class needs to implement the get_observation method and the user is also required to set the below attributes:

Attributes

uuid : universally unique id. observation_space : gym.Space object corresponding to observation of sensor. is_dict : whether the observation is a dictionary str_max_len : maximum length of the string representation of the encoded dictionary, if is_dict is True

Methods:

Name Description
get_observation

Returns observations from the environment (or task).
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/abstract_sensors.py 
34
35
36
def __init__(self, uuid: str, observation_space: gym.Space, **kwargs: Any) -> None:
    self.uuid = uuid
    self.observation_space = observation_space
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation abstractmethod 
get_observation(env, task, *args: Any, **kwargs: Any) -> Any

Returns observations from the environment (or task).

Parameters

env : The environment the sensor is used upon. task : (Optionally) a Task from which the sensor should get data.

Returns

Current observation for Sensor.

Source code in molmo_spaces/env/abstract_sensors.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
@abstractmethod
def get_observation(self, env, task, *args: Any, **kwargs: Any) -> Any:
    """Returns observations from the environment (or task).

    # Parameters

    env : The environment the sensor is used upon.
    task : (Optionally) a Task from which the sensor should get data.

    # Returns

    Current observation for Sensor.
    """
    raise NotImplementedError()
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

SensorSuite 

SensorSuite(sensors: Sequence[Sensor])

Represents a set of sensors, with each sensor being identified through a unique id.

Attributes
list containing sensors for the environment, uuid of each

sensor must be unique.

Initializer.

Parameters

param sensors: the sensors that will be included in the suite.

Methods:

Name Description
get

Return sensor with the given uuid.
get_observations

Get all observations corresponding to the sensors in the suite.

Attributes:

Name Type Description
observation_spaces Dict
sensors dict[str, Sensor]
Source code in molmo_spaces/env/abstract_sensors.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def __init__(self, sensors: Sequence[Sensor]) -> None:
    """Initializer.

    # Parameters

    param sensors: the sensors that will be included in the suite.
    """
    self.sensors = OrderedDict()
    spaces: OrderedDict[str, gym.Space] = OrderedDict()
    for sensor in sensors:
        assert sensor.uuid not in self.sensors, f"'{sensor.uuid}' is duplicated sensor uuid"
        self.sensors[sensor.uuid] = sensor
        spaces[sensor.uuid] = sensor.observation_space
    self.observation_spaces = SpaceDict(spaces=spaces)
observation_spaces instance-attribute 
observation_spaces: Dict = SpaceDict(spaces=spaces)
sensors instance-attribute 
sensors: dict[str, Sensor] = OrderedDict()
get 
get(uuid: str) -> Sensor

Return sensor with the given uuid.

Parameters

uuid : The unique id of the sensor

Returns

The sensor with unique id uuid.

Source code in molmo_spaces/env/abstract_sensors.py 
86
87
88
89
90
91
92
93
94
95
96
97
def get(self, uuid: str) -> Sensor:
    """Return sensor with the given `uuid`.

    # Parameters

    uuid : The unique id of the sensor

    # Returns

    The sensor with unique id `uuid`.
    """
    return self.sensors[uuid]
get_observations 
get_observations(env, task, **kwargs: Any) -> dict[str, Any]

Get all observations corresponding to the sensors in the suite.

Parameters

env : The environment from which to get the observation. task : (Optionally) the task from which to get the observation.

Returns

Data from all sensors packaged inside a Dict.

Source code in molmo_spaces/env/abstract_sensors.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def get_observations(self, env, task, **kwargs: Any) -> dict[str, Any]:
    """Get all observations corresponding to the sensors in the suite.

    # Parameters

    env : The environment from which to get the observation.
    task : (Optionally) the task from which to get the observation.

    # Returns

    Data from all sensors packaged inside a Dict.
    """
    return {
        uuid: sensor.get_observation(env=env, task=task, **kwargs)  # type: ignore
        for uuid, sensor in self.sensors.items()
    }

arena

Modules:

Name Description
arena_utils
bathroom
cabinet
drawer
kitchen
procthor_types
randomization
scene_tweaks

arena_utils

Functions:

Name Description
fix_exclude_contact_floor_with_fridges
fix_move_objects_within_inner_sites_up_abit
fix_remove_all_toasters
fix_remove_objects_within_inner_sites
get_all_bodies_with_joints_as_mlspaces_objects

Get all bodies with joints as MlSpacesObject instances.
load_env_with_objects
load_env_with_objects_with_tweaks
modify_mjmodel_thor_articulated

Attributes:

Name Type Description
DEFAULT_Z_OFFSET_OBJS_WITHIN_SITES
args
data
house_path
iTHOR_CATEGORIES
new_model MjModel | None
old_model MjModel | None
parser
t_start
DEFAULT_Z_OFFSET_OBJS_WITHIN_SITES module-attribute 
DEFAULT_Z_OFFSET_OBJS_WITHIN_SITES = 0.025
args module-attribute 
args = parse_args()
data module-attribute 
data = MjData(new_model)
house_path module-attribute 
house_path = Path(house)
iTHOR_CATEGORIES module-attribute 
iTHOR_CATEGORIES = ['Cabinet', 'Drawer', 'ShowerDoor', 'Oven', 'Dishwasher', 'StoveKnob']
new_model module-attribute 
new_model: MjModel | None = None
old_model module-attribute 
old_model: MjModel | None = None
parser module-attribute 
parser = ArgumentParser()
t_start module-attribute 
t_start = time
fix_exclude_contact_floor_with_fridges 
fix_exclude_contact_floor_with_fridges(spec: MjSpec) -> None
Source code in molmo_spaces/env/arena/arena_utils.py 
73
74
75
76
77
78
79
80
81
82
83
84
def fix_exclude_contact_floor_with_fridges(spec: mj.MjSpec) -> None:
    floor_handle = spec.body("floor")
    if not floor_handle:
        return

    fridges_bodies: list[mj.MjsBody] = []
    body_spec: mj.MjsBody = spec.worldbody.first_body()
    while body_spec:
        if "fridge" in body_spec.name.lower():
            spec.add_exclude(bodyname1="floor", bodyname2=body_spec.name)
            fridges_bodies.append(body_spec)
        body_spec = spec.worldbody.next_body(body_spec)
fix_move_objects_within_inner_sites_up_abit 
fix_move_objects_within_inner_sites_up_abit(spec: MjSpec, z_offset: float = DEFAULT_Z_OFFSET_OBJS_WITHIN_SITES) -> None
Source code in molmo_spaces/env/arena/arena_utils.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def fix_move_objects_within_inner_sites_up_abit(
    spec: mj.MjSpec, z_offset: float = DEFAULT_Z_OFFSET_OBJS_WITHIN_SITES
) -> None:
    model: mj.MjModel = spec.compile()
    data: mj.MjData = mj.MjData(model)
    mj.mj_forward(model, data)

    # Collect all root bodies from the model
    root_bodies: list[mj.MjsBody] = []
    root_bodies_z0: list[float] = []
    current_body_spec = spec.worldbody.first_body()
    while current_body_spec:
        root_bodies.append(current_body_spec)
        root_bodies_z0.append(current_body_spec.pos[2].item())
        current_body_spec = spec.worldbody.next_body(current_body_spec)

    for idx in range(len(root_bodies)):
        body_spec: mj.MjsBody = root_bodies[idx]
        body_id = model.body(body_spec.name).id
        is_within_site, _ = is_body_within_any_site(model, data, body_id)
        if is_within_site:
            body_spec.pos[2] = root_bodies_z0[idx] + z_offset
fix_remove_all_toasters 
fix_remove_all_toasters(spec: MjSpec) -> None
Source code in molmo_spaces/env/arena/arena_utils.py 
87
88
89
90
91
92
93
94
95
96
def fix_remove_all_toasters(spec: mj.MjSpec) -> None:
    toasters_handles: list[mj.MjsBody] = []
    root_body: mj.MjsBody = spec.worldbody.first_body()
    while root_body is not None:
        if "toaster" in root_body.name.lower():
            toasters_handles.append(root_body)
        root_body = spec.worldbody.next_body(root_body)

    for toaster_body in toasters_handles:
        spec.delete(toaster_body)
fix_remove_objects_within_inner_sites 
fix_remove_objects_within_inner_sites(spec: MjSpec) -> list[str]
Source code in molmo_spaces/env/arena/arena_utils.py 
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
def fix_remove_objects_within_inner_sites(spec: mj.MjSpec) -> list[str]:
    model: mj.MjModel = spec.compile()
    data: mj.MjData = mj.MjData(model)
    mj.mj_forward(model, data)

    # Collect all root bodies from the model
    root_bodies: list[mj.MjsBody] = []
    current_body_spec = spec.worldbody.first_body()
    while current_body_spec:
        root_bodies.append(current_body_spec)
        current_body_spec = spec.worldbody.next_body(current_body_spec)

    bodies_deleted = []
    for idx in range(len(root_bodies)):
        body_spec: mj.MjsBody = root_bodies[idx]
        body_id = model.body(body_spec.name).id
        for site_id in range(model.nsite):
            if is_body_com_within_box_site(site_id, body_id, model, data):
                in_free_space, _, _ = is_body_within_site_in_freespace(
                    site_id, body_id, model, data
                )
                if not in_free_space:  # it's actually inside a drawer or similar
                    bodies_deleted.append(body_spec.name)
                    spec.delete(body_spec)
                    break
    return bodies_deleted
get_all_bodies_with_joints_as_mlspaces_objects 
get_all_bodies_with_joints_as_mlspaces_objects(model: MjModel, data: MjData) -> list[MlSpacesObject]

Get all bodies with joints as MlSpacesObject instances.

This function finds all bodies in the model that have joints (movable bodies) and creates MlSpacesObject instances for them. Bodies without valid names or that fail to create MlSpacesObject instances are skipped.

Parameters:

Name Type Description Default
model MjModel

MuJoCo model

 required
data MjData

MuJoCo data

 required

Returns:

Type Description
list[MlSpacesObject]

List of MlSpacesObject instances for all bodies with joints that could be
list[MlSpacesObject]

successfully created. Bodies that fail to create MlSpacesObject instances
list[MlSpacesObject]

are silently skipped.
Source code in molmo_spaces/env/arena/arena_utils.py 
 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
def get_all_bodies_with_joints_as_mlspaces_objects(
    model: mj.MjModel, data: mj.MjData
) -> list[MlSpacesObject]:
    """
    Get all bodies with joints as MlSpacesObject instances.

    This function finds all bodies in the model that have joints (movable bodies)
    and creates MlSpacesObject instances for them. Bodies without valid names or
    that fail to create MlSpacesObject instances are skipped.

    Args:
        model: MuJoCo model
        data: MuJoCo data

    Returns:
        List of MlSpacesObject instances for all bodies with joints that could be
        successfully created. Bodies that fail to create MlSpacesObject instances
        are silently skipped.
    """
    mlspaces_objects = []
    bodies_with_joints = []

    # Iterate through all bodies (skip body 0 which is worldbody)
    for body_id in range(1, model.nbody):
        # Check if this body has a joint
        jnt_adr = model.body_jntadr[body_id]
        if jnt_adr >= 0:
            # This body has a joint, check if it's a free joint or any non-fixed joint
            # In MuJoCo, joints are either FREE, BALL, HINGE, SLIDE (no fixed joint type)
            # If a body has a joint, it's movable
            joint_type = model.jnt_type[jnt_adr]
            # Get body name
            name_adr = model.name_bodyadr[body_id]
            if name_adr >= 0:
                name_bytes = model.names[name_adr:]
                body_name = name_bytes.split(b"\x00")[0].decode("utf-8")
                if body_name:
                    bodies_with_joints.append((body_id, body_name, joint_type))

    # Create MlSpacesObject instances for bodies with joints
    for _, body_name, _ in bodies_with_joints:
        with contextlib.suppress(Exception):
            obj = MlSpacesObject(object_name=body_name, data=data)
            mlspaces_objects.append(obj)

    return mlspaces_objects
load_env_with_objects 
load_env_with_objects(xml_path: str) -> tuple[MjModel, dict[str, list[MlSpacesArticulationObject]]]
Source code in molmo_spaces/env/arena/arena_utils.py 
194
195
196
197
198
199
def load_env_with_objects(
    xml_path: str,
) -> tuple[mj.MjModel, dict[str, list[MlSpacesArticulationObject]]]:
    model = mj.MjModel.from_xml_path(xml_path)
    data = mj.MjData(model)
    return model, modify_mjmodel_thor_articulated(model, data)
load_env_with_objects_with_tweaks 
load_env_with_objects_with_tweaks(xml_path: str, remove_objects_within_inner_sites: bool = False, move_objects_within_sites_up_abit: bool = False, remove_all_toasters: bool = False) -> tuple[MjModel, dict[str, list[MlSpacesArticulationObject]], dict]
Source code in molmo_spaces/env/arena/arena_utils.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
def load_env_with_objects_with_tweaks(
    xml_path: str,
    remove_objects_within_inner_sites: bool = False,
    move_objects_within_sites_up_abit: bool = False,
    remove_all_toasters: bool = False,
    # exclude_floor_contact_with_fridges: bool = False,
) -> tuple[mj.MjModel, dict[str, list[MlSpacesArticulationObject]], dict]:
    tweaks = {"bodies_deleted": []}

    spec = mj.MjSpec.from_file(xml_path)

    if remove_objects_within_inner_sites:
        tweaks["bodies_deleted"].extend(fix_remove_objects_within_inner_sites(spec))
    if move_objects_within_sites_up_abit:
        fix_move_objects_within_inner_sites_up_abit(spec, DEFAULT_Z_OFFSET_OBJS_WITHIN_SITES)
    if remove_all_toasters:
        fix_remove_all_toasters(spec)
    # if exclude_floor_contact_with_fridges:
    #     fix_exclude_contact_floor_with_fridges(spec)

    model = spec.compile()

    # filepath = Path(xml_path)
    # filepath_new = filepath.parent / f"{filepath.stem}_new.xml"
    # with open(filepath_new, "w") as fhandle:
    #     fhandle.write(spec.to_xml())

    data = mj.MjData(model)
    body_name2id = {model.body(i).name: i for i in range(0, model.nbody)}

    # get root bodies
    root_bodies_dict = {
        "Cabinet": [],
        "Drawer": [],
        "ShowerDoor": [],
        "Oven": [],
        "Dishwasher": [],
        "StoveKnob": [],
    }

    root_bodies = set()
    for i in range(0, model.nbody):
        rootid = model.body(i).rootid.item()
        root_body_name = model.body(rootid).name
        root_bodies.add(root_body_name)

    for root_body_name in root_bodies:
        category = next(
            (i for i in iTHOR_CATEGORIES if root_body_name.lower().startswith(i.lower())), None
        )
        if category == "Cabinet":
            cabinet = Cabinet(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(cabinet)
        elif category == "Drawer":
            drawer = Drawer(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(drawer)
        elif category == "ShowerDoor":
            shower_door = ShowerDoor(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(shower_door)
        elif category == "Oven":
            oven = Oven(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(oven)
        elif category == "Dishwasher":
            dishwasher = Dishwasher(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(dishwasher)
        elif category == "StoveKnob":
            stove_knob = Stoveknob(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(stove_knob)

    return model, root_bodies_dict, tweaks
modify_mjmodel_thor_articulated 
modify_mjmodel_thor_articulated(model: MjModel, data) -> dict[str, list[MlSpacesArticulationObject]]
Source code in molmo_spaces/env/arena/arena_utils.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def modify_mjmodel_thor_articulated(
    model: mj.MjModel, data
) -> dict[str, list[MlSpacesArticulationObject]]:
    body_name2id = {model.body(i).name: i for i in range(0, model.nbody)}

    # get root bodies
    root_bodies_dict = {
        "Cabinet": [],
        "Drawer": [],
        "ShowerDoor": [],
        "Oven": [],
        "Dishwasher": [],
        "StoveKnob": [],
    }

    root_bodies = set()
    for i in range(0, model.nbody):
        rootid = model.body(i).rootid.item()
        root_body_name = model.body(rootid).name
        root_bodies.add(root_body_name)

    for root_body_name in root_bodies:
        category = next(
            (i for i in iTHOR_CATEGORIES if root_body_name.lower().startswith(i.lower())), None
        )
        if category == "Cabinet":
            cabinet = Cabinet(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(cabinet)
        elif category == "Drawer":
            drawer = Drawer(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(drawer)
        elif category == "ShowerDoor":
            shower_door = ShowerDoor(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(shower_door)
        elif category == "Oven":
            oven = Oven(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(oven)
        elif category == "Dishwasher":
            dishwasher = Dishwasher(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(dishwasher)
        elif category == "StoveKnob":
            stove_knob = Stoveknob(root_body_name, data, body_name2id)
            root_bodies_dict[category].append(stove_knob)

    return root_bodies_dict

bathroom

Classes:

Name Description
ShowerDoor

Attributes:

Name Type Description
SHOWER_DOOR
SHOWER_DOOR module-attribute 
SHOWER_DOOR = RandomizePhysicalProperties()
ShowerDoor 
ShowerDoor(object_name: str, data: MjData, body_name2id: dict[str, int])

Bases: MlSpacesArticulationObject

Methods:

Name Description
body_name2id
body_parent_id
build_children_lists
find_top_object_body_id
get_ancestors
get_descendants
get_direct_children
get_friction
get_geom_infos

Get geom information for this object.
get_geom_type_name
get_joint_anchor_position
get_joint_armature
get_joint_axis
get_joint_body_orientation
get_joint_damping
get_joint_frictionloss
get_joint_leaf_body_position
get_joint_position
get_joint_qpos_adr
get_joint_range
get_joint_stiffness
get_joint_type
get_mass
get_top_level_bodies

Return bodies whose parent is the world body.
is_child_name_of
is_object_in_gripper
is_object_picked_up
set_joint_position

Attributes:

Name Type Description
aabb_center ndarray
aabb_size ndarray
body_id int
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).
bvh_root
center_of_mass
geom_ids

Get all geom IDs belonging to this object (lazy, cached).
joint_id2name
joint_id2qpos_adr
joint_ids
joint_names
mj_data
mj_model
name str
njoints int
object_id int
object_root_id
pose ndarray
position
quat
Source code in molmo_spaces/env/arena/bathroom.py 
23
24
25
26
27
28
29
30
31
def __init__(self, object_name: str, data: mujoco.MjData, body_name2id: dict[str, int]) -> None:
    super().__init__(object_name, data)
    self._set_friction(SHOWER_DOOR.get_random_value("friction"))
    # Note: Density cannot be modified at runtime in MuJoCo
    # self._set_density(SHOWER_DOOR.get_random_value("density"))
    self._set_joint_stiffness(SHOWER_DOOR.get_random_value("joint_stiffness"))
    self._set_joint_damping(SHOWER_DOOR.get_random_value("joint_damping"))
    self._set_joint_armature(SHOWER_DOOR.get_random_value("joint_armature"))
    self._set_joint_frictionloss(SHOWER_DOOR.get_random_value("joint_frictionloss"))
aabb_center property 
aabb_center: ndarray
aabb_size property 
aabb_size: ndarray
body_id cached property 
body_id: int
body_ids cached property 
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).

bvh_root property 
bvh_root
center_of_mass property 
center_of_mass
geom_ids cached property 
geom_ids

Get all geom IDs belonging to this object (lazy, cached).

joint_id2name instance-attribute 
joint_id2name = {}
joint_id2qpos_adr instance-attribute 
joint_id2qpos_adr = {}
joint_ids instance-attribute 
joint_ids = []
joint_names instance-attribute 
joint_names = []
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
njoints property 
njoints: int
object_id property 
object_id: int
object_root_id cached property 
object_root_id
pose property writable 
pose: ndarray
position property 
position
quat property 
quat
body_name2id staticmethod 
body_name2id(model: MjModel) -> dict[str, int]
Source code in molmo_spaces/env/data_views.py 
467
468
469
@staticmethod
def body_name2id(model: mujoco.MjModel) -> dict[str, int]:
    return {model.body(i).name: i for i in range(model.nbody)}
body_parent_id staticmethod 
body_parent_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
387
388
389
390
391
392
393
@staticmethod
def body_parent_id(model: mujoco.MjModel, body_id: int) -> int:
    try:
        return int(model.body_parentid[body_id])
    except Exception as e:
        print(f"Error getting body parent id: {e}")
        return int(model.body(body_id).parentid)
build_children_lists staticmethod 
build_children_lists(model: MjModel) -> list[list[int]]
Source code in molmo_spaces/env/data_views.py 
395
396
397
398
399
400
401
402
@staticmethod
def build_children_lists(model: mujoco.MjModel) -> list[list[int]]:
    children: list[list[int]] = [[] for _ in range(model.nbody)]
    for child_id in range(model.nbody):
        pid = MlSpacesObject.body_parent_id(model, child_id)
        if pid >= 0:
            children[pid].append(child_id)
    return children
find_top_object_body_id staticmethod 
find_top_object_body_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
@staticmethod
def find_top_object_body_id(model: mujoco.MjModel, body_id: int) -> int:
    cur = body_id
    visited = set()
    while True:
        if cur in visited:
            break
        visited.add(cur)
        pid = MlSpacesObject.body_parent_id(model, cur)
        if pid < 0:
            break
        pname = model.body(pid).name
        cname = model.body(cur).name
        if MlSpacesObject.is_child_name_of(pname, cname):
            cur = pid
        else:
            break
    return int(cur)
get_ancestors staticmethod 
get_ancestors(model: MjModel, body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
404
405
406
407
408
409
410
411
412
413
@staticmethod
def get_ancestors(model: mujoco.MjModel, body_id: int) -> list[int]:
    ancestors: list[int] = []
    visited = {body_id}
    pid = MlSpacesObject.body_parent_id(model, body_id)
    while pid >= 0 and pid not in visited:
        ancestors.append(pid)
        visited.add(pid)
        pid = MlSpacesObject.body_parent_id(model, pid)
    return ancestors
get_descendants staticmethod 
get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
415
416
417
418
419
420
421
422
423
424
@staticmethod
def get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]:
    stack = [body_id]
    desc: list[int] = []
    while stack:
        cur = stack.pop()
        for c in children_lists[cur]:
            desc.append(int(c))
            stack.append(c)
    return desc
get_direct_children staticmethod 
get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
426
427
428
@staticmethod
def get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]:
    return list(children_lists[body_id])
get_friction 
get_friction()
Source code in molmo_spaces/env/data_views.py 
381
382
def get_friction(self):
    return [self.mj_model.geom(geom_id).friction for geom_id in self.geom_ids]
get_geom_infos 
get_geom_infos(include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]

Get geom information for this object.

Parameters:

Name Type Description Default
include_descendants bool

If True, includes geoms from all descendant bodies. If False, only includes geoms directly attached to this object's body.

 True
max_geoms int | None

Maximum number of geoms to return. If None, returns all.

 2048

Returns:

Type Description
list[dict[str, object]]

List of dicts with geom info: id, name, position, size, type, type_name
Source code in molmo_spaces/env/data_views.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def get_geom_infos(
    self, include_descendants: bool = True, max_geoms: int | None = 2048
) -> list[dict[str, object]]:
    """Get geom information for this object.

    Args:
        include_descendants: If True, includes geoms from all descendant bodies.
            If False, only includes geoms directly attached to this object's body.
        max_geoms: Maximum number of geoms to return. If None, returns all.

    Returns:
        List of dicts with geom info: id, name, position, size, type, type_name
    """
    # Build set of body IDs to include
    if include_descendants:
        body_ids = set(self.body_ids)
    else:
        body_ids = {self.object_id}

    # Find geoms belonging to these bodies
    geoms: list[dict[str, object]] = []
    count = 0
    for geom_id in range(self.mj_model.ngeom):
        if max_geoms is not None and count >= max_geoms:
            break
        if int(self.mj_model.geom_bodyid[geom_id]) in body_ids:
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id))
            pos = self.mj_data.geom_xpos[geom_id].copy()
            size = self.mj_model.geom_size[geom_id].copy()
            gtype = int(self.mj_model.geom_type[geom_id])
            geoms.append(
                {
                    "id": int(geom_id),
                    "name": name,
                    "position": pos,
                    "size": size,
                    "type": gtype,
                    "type_name": MlSpacesObject.get_geom_type_name(gtype),
                }
            )
            count += 1

    return geoms
get_geom_type_name staticmethod 
get_geom_type_name(geom_type: int) -> str
Source code in molmo_spaces/env/data_views.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
@staticmethod
def get_geom_type_name(geom_type: int) -> str:
    names = {
        0: "PLANE",
        1: "HFIELD",
        2: "SPHERE",
        3: "CAPSULE",
        4: "ELLIPSOID",
        5: "CYLINDER",
        6: "BOX",
        7: "MESH",
    }
    return names.get(int(geom_type), f"UNKNOWN({geom_type})")
get_joint_anchor_position 
get_joint_anchor_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
550
551
552
553
554
555
556
557
558
559
def get_joint_anchor_position(self, i: int) -> np.ndarray:
    # Get the position in world frame of the joint anchor point
    joint_id = self.joint_ids[i]
    body_id = self.mj_model.jnt_bodyid[joint_id]
    local_anchor = self.mj_model.jnt_pos[joint_id]
    body_pos = self.mj_data.xpos[body_id]
    body_quat = self.mj_data.xquat[body_id]
    body_rot = R.from_quat(body_quat, scalar_first=True).as_matrix()
    world_anchor = body_pos + body_rot @ local_anchor
    return world_anchor
get_joint_armature 
get_joint_armature()
Source code in molmo_spaces/env/data_views.py 
570
571
def get_joint_armature(self):
    return [self.mj_model.joint(joint_id).armature for joint_id in self.joint_ids]
get_joint_axis 
get_joint_axis(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
526
527
def get_joint_axis(self, i: int) -> np.ndarray:
    return self.mj_model.joint(self.joint_ids[i]).axis
get_joint_body_orientation 
get_joint_body_orientation(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
545
546
547
548
def get_joint_body_orientation(self, i: int) -> np.ndarray:
    # orientation of the joint body not the root orientation of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    return self.mj_data.xmat[body_id].copy().reshape(3, 3)
get_joint_damping 
get_joint_damping()
Source code in molmo_spaces/env/data_views.py 
567
568
def get_joint_damping(self):
    return [self.mj_model.joint(joint_id).damping for joint_id in self.joint_ids]
get_joint_frictionloss 
get_joint_frictionloss()
Source code in molmo_spaces/env/data_views.py 
561
562
def get_joint_frictionloss(self):
    return [self.mj_model.joint(joint_id).frictionloss for joint_id in self.joint_ids]
get_joint_leaf_body_position 
get_joint_leaf_body_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
535
536
537
538
539
540
541
542
543
def get_joint_leaf_body_position(self, i: int) -> np.ndarray:
    # position of the joint body not the root center of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    # find child body of this joint body
    for child_body_id in range(self.mj_model.nbody):
        if self.mj_model.body(child_body_id).parentid[0] == body_id:
            body_id = child_body_id
            break
    return self.mj_data.xpos[body_id].copy()
get_joint_position 
get_joint_position(i: int) -> float
Source code in molmo_spaces/env/data_views.py 
520
521
def get_joint_position(self, i: int) -> float:
    return self.mj_data.qpos[self.get_joint_qpos_adr(i)].copy()
get_joint_qpos_adr 
get_joint_qpos_adr(i: int) -> int
Source code in molmo_spaces/env/data_views.py 
532
533
def get_joint_qpos_adr(self, i: int) -> int:
    return int(self.mj_model.joint(self.joint_ids[i]).qposadr[0])
get_joint_range 
get_joint_range(i: int) -> tuple[float, float]
Source code in molmo_spaces/env/data_views.py 
523
524
def get_joint_range(self, i: int) -> tuple[float, float]:
    return self.mj_model.joint(self.joint_ids[i]).range
get_joint_stiffness 
get_joint_stiffness()
Source code in molmo_spaces/env/data_views.py 
564
565
def get_joint_stiffness(self):
    return [self.mj_model.joint(joint_id).stiffness for joint_id in self.joint_ids]
get_joint_type 
get_joint_type(i: int) -> mjtJoint
Source code in molmo_spaces/env/data_views.py 
529
530
def get_joint_type(self, i: int) -> mujoco.mjtJoint:
    return self.mj_model.joint(self.joint_ids[i]).type
get_mass 
get_mass()
Source code in molmo_spaces/env/data_views.py 
384
385
def get_mass(self):
    return [self.mj_model.body_mass[body_id] for body_id in self.body_ids]
get_top_level_bodies staticmethod 
get_top_level_bodies(model: MjModel) -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/data_views.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
@staticmethod
def get_top_level_bodies(model: mujoco.MjModel) -> list[int]:
    """Return bodies whose parent is the world body."""
    # Identify world body id: the one with parent < 0
    try:
        world_id = next(
            b for b in range(model.nbody) if MlSpacesObject.body_parent_id(model, b) < 0
        )
    except StopIteration:
        world_id = 0
    tops: list[int] = []
    for b in range(model.nbody):
        try:
            pid = MlSpacesObject.body_parent_id(model, b)
            if pid == world_id:
                tops.append(b)
        except Exception as e:
            print(f"Error getting top-level bodies: {e}")
            continue
    return tops
is_child_name_of staticmethod 
is_child_name_of(parent_name: str, child_name: str) -> bool
Source code in molmo_spaces/env/data_views.py 
430
431
432
@staticmethod
def is_child_name_of(parent_name: str, child_name: str) -> bool:
    return child_name.startswith(parent_name + "_")
is_object_in_gripper 
is_object_in_gripper() -> None
Source code in molmo_spaces/env/data_views.py 
277
278
def is_object_in_gripper(self) -> None:
    pass
is_object_picked_up 
is_object_picked_up() -> None
Source code in molmo_spaces/env/data_views.py 
274
275
def is_object_picked_up(self) -> None:
    pass
set_joint_position 
set_joint_position(i: int, position: float) -> None
Source code in molmo_spaces/env/data_views.py 
573
574
575
576
def set_joint_position(self, i: int, position: float) -> None:
    if isinstance(position, np.ndarray):
        position = position.item()
    self.mj_data.qpos[self.get_joint_qpos_adr(i)] = position

cabinet

Classes:

Name Description
Cabinet

Attributes:

Name Type Description
CABINET
CABINET module-attribute 
CABINET = RandomizePhysicalProperties()
Cabinet 
Cabinet(object_name: str, data: MjData, body_name2id: dict[str, int])

Bases: MlSpacesArticulationObject

Methods:

Name Description
body_name2id
body_parent_id
build_children_lists
find_top_object_body_id
get_ancestors
get_descendants
get_direct_children
get_friction
get_geom_infos

Get geom information for this object.
get_geom_type_name
get_joint_anchor_position
get_joint_armature
get_joint_axis
get_joint_body_orientation
get_joint_damping
get_joint_frictionloss
get_joint_leaf_body_position
get_joint_position
get_joint_qpos_adr
get_joint_range
get_joint_stiffness
get_joint_type
get_mass
get_top_level_bodies

Return bodies whose parent is the world body.
is_child_name_of
is_object_in_gripper
is_object_picked_up
set_joint_position

Attributes:

Name Type Description
aabb_center ndarray
aabb_size ndarray
body_id int
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).
bvh_root
center_of_mass
geom_ids

Get all geom IDs belonging to this object (lazy, cached).
joint_id2name
joint_id2qpos_adr
joint_ids
joint_names
mj_data
mj_model
name str
njoints int
object_id int
object_root_id
pose ndarray
position
quat
Source code in molmo_spaces/env/arena/cabinet.py 
23
24
25
26
27
28
29
30
31
32
def __init__(self, object_name: str, data: mujoco.MjData, body_name2id: dict[str, int]) -> None:
    super().__init__(object_name, data)
    self._set_friction(CABINET.get_random_value("friction"))
    # Note: Density cannot be modified at runtime in MuJoCo
    # self._set_density(CABINET.get_random_value("density"))
    self._set_max_mass(1)
    self._set_joint_stiffness(CABINET.get_random_value("joint_stiffness"))
    self._set_joint_damping(CABINET.get_random_value("joint_damping"))
    self._set_joint_armature(CABINET.get_random_value("joint_armature"))
    self._set_joint_frictionloss(CABINET.get_random_value("joint_frictionloss"))
aabb_center property 
aabb_center: ndarray
aabb_size property 
aabb_size: ndarray
body_id cached property 
body_id: int
body_ids cached property 
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).

bvh_root property 
bvh_root
center_of_mass property 
center_of_mass
geom_ids cached property 
geom_ids

Get all geom IDs belonging to this object (lazy, cached).

joint_id2name instance-attribute 
joint_id2name = {}
joint_id2qpos_adr instance-attribute 
joint_id2qpos_adr = {}
joint_ids instance-attribute 
joint_ids = []
joint_names instance-attribute 
joint_names = []
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
njoints property 
njoints: int
object_id property 
object_id: int
object_root_id cached property 
object_root_id
pose property writable 
pose: ndarray
position property 
position
quat property 
quat
body_name2id staticmethod 
body_name2id(model: MjModel) -> dict[str, int]
Source code in molmo_spaces/env/data_views.py 
467
468
469
@staticmethod
def body_name2id(model: mujoco.MjModel) -> dict[str, int]:
    return {model.body(i).name: i for i in range(model.nbody)}
body_parent_id staticmethod 
body_parent_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
387
388
389
390
391
392
393
@staticmethod
def body_parent_id(model: mujoco.MjModel, body_id: int) -> int:
    try:
        return int(model.body_parentid[body_id])
    except Exception as e:
        print(f"Error getting body parent id: {e}")
        return int(model.body(body_id).parentid)
build_children_lists staticmethod 
build_children_lists(model: MjModel) -> list[list[int]]
Source code in molmo_spaces/env/data_views.py 
395
396
397
398
399
400
401
402
@staticmethod
def build_children_lists(model: mujoco.MjModel) -> list[list[int]]:
    children: list[list[int]] = [[] for _ in range(model.nbody)]
    for child_id in range(model.nbody):
        pid = MlSpacesObject.body_parent_id(model, child_id)
        if pid >= 0:
            children[pid].append(child_id)
    return children
find_top_object_body_id staticmethod 
find_top_object_body_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
@staticmethod
def find_top_object_body_id(model: mujoco.MjModel, body_id: int) -> int:
    cur = body_id
    visited = set()
    while True:
        if cur in visited:
            break
        visited.add(cur)
        pid = MlSpacesObject.body_parent_id(model, cur)
        if pid < 0:
            break
        pname = model.body(pid).name
        cname = model.body(cur).name
        if MlSpacesObject.is_child_name_of(pname, cname):
            cur = pid
        else:
            break
    return int(cur)
get_ancestors staticmethod 
get_ancestors(model: MjModel, body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
404
405
406
407
408
409
410
411
412
413
@staticmethod
def get_ancestors(model: mujoco.MjModel, body_id: int) -> list[int]:
    ancestors: list[int] = []
    visited = {body_id}
    pid = MlSpacesObject.body_parent_id(model, body_id)
    while pid >= 0 and pid not in visited:
        ancestors.append(pid)
        visited.add(pid)
        pid = MlSpacesObject.body_parent_id(model, pid)
    return ancestors
get_descendants staticmethod 
get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
415
416
417
418
419
420
421
422
423
424
@staticmethod
def get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]:
    stack = [body_id]
    desc: list[int] = []
    while stack:
        cur = stack.pop()
        for c in children_lists[cur]:
            desc.append(int(c))
            stack.append(c)
    return desc
get_direct_children staticmethod 
get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
426
427
428
@staticmethod
def get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]:
    return list(children_lists[body_id])
get_friction 
get_friction()
Source code in molmo_spaces/env/data_views.py 
381
382
def get_friction(self):
    return [self.mj_model.geom(geom_id).friction for geom_id in self.geom_ids]
get_geom_infos 
get_geom_infos(include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]

Get geom information for this object.

Parameters:

Name Type Description Default
include_descendants bool

If True, includes geoms from all descendant bodies. If False, only includes geoms directly attached to this object's body.

 True
max_geoms int | None

Maximum number of geoms to return. If None, returns all.

 2048

Returns:

Type Description
list[dict[str, object]]

List of dicts with geom info: id, name, position, size, type, type_name
Source code in molmo_spaces/env/data_views.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def get_geom_infos(
    self, include_descendants: bool = True, max_geoms: int | None = 2048
) -> list[dict[str, object]]:
    """Get geom information for this object.

    Args:
        include_descendants: If True, includes geoms from all descendant bodies.
            If False, only includes geoms directly attached to this object's body.
        max_geoms: Maximum number of geoms to return. If None, returns all.

    Returns:
        List of dicts with geom info: id, name, position, size, type, type_name
    """
    # Build set of body IDs to include
    if include_descendants:
        body_ids = set(self.body_ids)
    else:
        body_ids = {self.object_id}

    # Find geoms belonging to these bodies
    geoms: list[dict[str, object]] = []
    count = 0
    for geom_id in range(self.mj_model.ngeom):
        if max_geoms is not None and count >= max_geoms:
            break
        if int(self.mj_model.geom_bodyid[geom_id]) in body_ids:
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id))
            pos = self.mj_data.geom_xpos[geom_id].copy()
            size = self.mj_model.geom_size[geom_id].copy()
            gtype = int(self.mj_model.geom_type[geom_id])
            geoms.append(
                {
                    "id": int(geom_id),
                    "name": name,
                    "position": pos,
                    "size": size,
                    "type": gtype,
                    "type_name": MlSpacesObject.get_geom_type_name(gtype),
                }
            )
            count += 1

    return geoms
get_geom_type_name staticmethod 
get_geom_type_name(geom_type: int) -> str
Source code in molmo_spaces/env/data_views.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
@staticmethod
def get_geom_type_name(geom_type: int) -> str:
    names = {
        0: "PLANE",
        1: "HFIELD",
        2: "SPHERE",
        3: "CAPSULE",
        4: "ELLIPSOID",
        5: "CYLINDER",
        6: "BOX",
        7: "MESH",
    }
    return names.get(int(geom_type), f"UNKNOWN({geom_type})")
get_joint_anchor_position 
get_joint_anchor_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
550
551
552
553
554
555
556
557
558
559
def get_joint_anchor_position(self, i: int) -> np.ndarray:
    # Get the position in world frame of the joint anchor point
    joint_id = self.joint_ids[i]
    body_id = self.mj_model.jnt_bodyid[joint_id]
    local_anchor = self.mj_model.jnt_pos[joint_id]
    body_pos = self.mj_data.xpos[body_id]
    body_quat = self.mj_data.xquat[body_id]
    body_rot = R.from_quat(body_quat, scalar_first=True).as_matrix()
    world_anchor = body_pos + body_rot @ local_anchor
    return world_anchor
get_joint_armature 
get_joint_armature()
Source code in molmo_spaces/env/data_views.py 
570
571
def get_joint_armature(self):
    return [self.mj_model.joint(joint_id).armature for joint_id in self.joint_ids]
get_joint_axis 
get_joint_axis(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
526
527
def get_joint_axis(self, i: int) -> np.ndarray:
    return self.mj_model.joint(self.joint_ids[i]).axis
get_joint_body_orientation 
get_joint_body_orientation(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
545
546
547
548
def get_joint_body_orientation(self, i: int) -> np.ndarray:
    # orientation of the joint body not the root orientation of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    return self.mj_data.xmat[body_id].copy().reshape(3, 3)
get_joint_damping 
get_joint_damping()
Source code in molmo_spaces/env/data_views.py 
567
568
def get_joint_damping(self):
    return [self.mj_model.joint(joint_id).damping for joint_id in self.joint_ids]
get_joint_frictionloss 
get_joint_frictionloss()
Source code in molmo_spaces/env/data_views.py 
561
562
def get_joint_frictionloss(self):
    return [self.mj_model.joint(joint_id).frictionloss for joint_id in self.joint_ids]
get_joint_leaf_body_position 
get_joint_leaf_body_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
535
536
537
538
539
540
541
542
543
def get_joint_leaf_body_position(self, i: int) -> np.ndarray:
    # position of the joint body not the root center of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    # find child body of this joint body
    for child_body_id in range(self.mj_model.nbody):
        if self.mj_model.body(child_body_id).parentid[0] == body_id:
            body_id = child_body_id
            break
    return self.mj_data.xpos[body_id].copy()
get_joint_position 
get_joint_position(i: int) -> float
Source code in molmo_spaces/env/data_views.py 
520
521
def get_joint_position(self, i: int) -> float:
    return self.mj_data.qpos[self.get_joint_qpos_adr(i)].copy()
get_joint_qpos_adr 
get_joint_qpos_adr(i: int) -> int
Source code in molmo_spaces/env/data_views.py 
532
533
def get_joint_qpos_adr(self, i: int) -> int:
    return int(self.mj_model.joint(self.joint_ids[i]).qposadr[0])
get_joint_range 
get_joint_range(i: int) -> tuple[float, float]
Source code in molmo_spaces/env/data_views.py 
523
524
def get_joint_range(self, i: int) -> tuple[float, float]:
    return self.mj_model.joint(self.joint_ids[i]).range
get_joint_stiffness 
get_joint_stiffness()
Source code in molmo_spaces/env/data_views.py 
564
565
def get_joint_stiffness(self):
    return [self.mj_model.joint(joint_id).stiffness for joint_id in self.joint_ids]
get_joint_type 
get_joint_type(i: int) -> mjtJoint
Source code in molmo_spaces/env/data_views.py 
529
530
def get_joint_type(self, i: int) -> mujoco.mjtJoint:
    return self.mj_model.joint(self.joint_ids[i]).type
get_mass 
get_mass()
Source code in molmo_spaces/env/data_views.py 
384
385
def get_mass(self):
    return [self.mj_model.body_mass[body_id] for body_id in self.body_ids]
get_top_level_bodies staticmethod 
get_top_level_bodies(model: MjModel) -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/data_views.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
@staticmethod
def get_top_level_bodies(model: mujoco.MjModel) -> list[int]:
    """Return bodies whose parent is the world body."""
    # Identify world body id: the one with parent < 0
    try:
        world_id = next(
            b for b in range(model.nbody) if MlSpacesObject.body_parent_id(model, b) < 0
        )
    except StopIteration:
        world_id = 0
    tops: list[int] = []
    for b in range(model.nbody):
        try:
            pid = MlSpacesObject.body_parent_id(model, b)
            if pid == world_id:
                tops.append(b)
        except Exception as e:
            print(f"Error getting top-level bodies: {e}")
            continue
    return tops
is_child_name_of staticmethod 
is_child_name_of(parent_name: str, child_name: str) -> bool
Source code in molmo_spaces/env/data_views.py 
430
431
432
@staticmethod
def is_child_name_of(parent_name: str, child_name: str) -> bool:
    return child_name.startswith(parent_name + "_")
is_object_in_gripper 
is_object_in_gripper() -> None
Source code in molmo_spaces/env/data_views.py 
277
278
def is_object_in_gripper(self) -> None:
    pass
is_object_picked_up 
is_object_picked_up() -> None
Source code in molmo_spaces/env/data_views.py 
274
275
def is_object_picked_up(self) -> None:
    pass
set_joint_position 
set_joint_position(i: int, position: float) -> None
Source code in molmo_spaces/env/data_views.py 
573
574
575
576
def set_joint_position(self, i: int, position: float) -> None:
    if isinstance(position, np.ndarray):
        position = position.item()
    self.mj_data.qpos[self.get_joint_qpos_adr(i)] = position

drawer

Classes:

Name Description
Drawer

Attributes:

Name Type Description
DRAWER
DRAWER module-attribute 
DRAWER = RandomizePhysicalProperties()
Drawer 
Drawer(object_name: str, data: MjData, body_name2id: dict[str, int])

Bases: MlSpacesArticulationObject

Methods:

Name Description
body_name2id
body_parent_id
build_children_lists
find_top_object_body_id
get_ancestors
get_descendants
get_direct_children
get_friction
get_geom_infos

Get geom information for this object.
get_geom_type_name
get_joint_anchor_position
get_joint_armature
get_joint_axis
get_joint_body_orientation
get_joint_damping
get_joint_frictionloss
get_joint_leaf_body_position
get_joint_position
get_joint_qpos_adr
get_joint_range
get_joint_stiffness
get_joint_type
get_mass
get_top_level_bodies

Return bodies whose parent is the world body.
is_child_name_of
is_object_in_gripper
is_object_picked_up
set_joint_position

Attributes:

Name Type Description
aabb_center ndarray
aabb_size ndarray
body_id int
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).
bvh_root
center_of_mass
geom_ids

Get all geom IDs belonging to this object (lazy, cached).
joint_id2name
joint_id2qpos_adr
joint_ids
joint_names
mj_data
mj_model
name str
njoints int
object_id int
object_root_id
pose ndarray
position
quat
Source code in molmo_spaces/env/arena/drawer.py 
23
24
25
26
27
28
29
30
31
32
def __init__(self, object_name: str, data: mujoco.MjData, body_name2id: dict[str, int]) -> None:
    super().__init__(object_name, data)
    self._set_friction(DRAWER.get_random_value("friction"))
    # Note: Density cannot be modified at runtime in MuJoCo
    # self._set_density(DRAWER.get_random_value("density"))
    self._set_max_mass(1)
    self._set_joint_stiffness(DRAWER.get_random_value("joint_stiffness"))
    self._set_joint_damping(DRAWER.get_random_value("joint_damping"))
    self._set_joint_armature(DRAWER.get_random_value("joint_armature"))
    self._set_joint_frictionloss(DRAWER.get_random_value("joint_frictionloss"))
aabb_center property 
aabb_center: ndarray
aabb_size property 
aabb_size: ndarray
body_id cached property 
body_id: int
body_ids cached property 
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).

bvh_root property 
bvh_root
center_of_mass property 
center_of_mass
geom_ids cached property 
geom_ids

Get all geom IDs belonging to this object (lazy, cached).

joint_id2name instance-attribute 
joint_id2name = {}
joint_id2qpos_adr instance-attribute 
joint_id2qpos_adr = {}
joint_ids instance-attribute 
joint_ids = []
joint_names instance-attribute 
joint_names = []
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
njoints property 
njoints: int
object_id property 
object_id: int
object_root_id cached property 
object_root_id
pose property writable 
pose: ndarray
position property 
position
quat property 
quat
body_name2id staticmethod 
body_name2id(model: MjModel) -> dict[str, int]
Source code in molmo_spaces/env/data_views.py 
467
468
469
@staticmethod
def body_name2id(model: mujoco.MjModel) -> dict[str, int]:
    return {model.body(i).name: i for i in range(model.nbody)}
body_parent_id staticmethod 
body_parent_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
387
388
389
390
391
392
393
@staticmethod
def body_parent_id(model: mujoco.MjModel, body_id: int) -> int:
    try:
        return int(model.body_parentid[body_id])
    except Exception as e:
        print(f"Error getting body parent id: {e}")
        return int(model.body(body_id).parentid)
build_children_lists staticmethod 
build_children_lists(model: MjModel) -> list[list[int]]
Source code in molmo_spaces/env/data_views.py 
395
396
397
398
399
400
401
402
@staticmethod
def build_children_lists(model: mujoco.MjModel) -> list[list[int]]:
    children: list[list[int]] = [[] for _ in range(model.nbody)]
    for child_id in range(model.nbody):
        pid = MlSpacesObject.body_parent_id(model, child_id)
        if pid >= 0:
            children[pid].append(child_id)
    return children
find_top_object_body_id staticmethod 
find_top_object_body_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
@staticmethod
def find_top_object_body_id(model: mujoco.MjModel, body_id: int) -> int:
    cur = body_id
    visited = set()
    while True:
        if cur in visited:
            break
        visited.add(cur)
        pid = MlSpacesObject.body_parent_id(model, cur)
        if pid < 0:
            break
        pname = model.body(pid).name
        cname = model.body(cur).name
        if MlSpacesObject.is_child_name_of(pname, cname):
            cur = pid
        else:
            break
    return int(cur)
get_ancestors staticmethod 
get_ancestors(model: MjModel, body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
404
405
406
407
408
409
410
411
412
413
@staticmethod
def get_ancestors(model: mujoco.MjModel, body_id: int) -> list[int]:
    ancestors: list[int] = []
    visited = {body_id}
    pid = MlSpacesObject.body_parent_id(model, body_id)
    while pid >= 0 and pid not in visited:
        ancestors.append(pid)
        visited.add(pid)
        pid = MlSpacesObject.body_parent_id(model, pid)
    return ancestors
get_descendants staticmethod 
get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
415
416
417
418
419
420
421
422
423
424
@staticmethod
def get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]:
    stack = [body_id]
    desc: list[int] = []
    while stack:
        cur = stack.pop()
        for c in children_lists[cur]:
            desc.append(int(c))
            stack.append(c)
    return desc
get_direct_children staticmethod 
get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
426
427
428
@staticmethod
def get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]:
    return list(children_lists[body_id])
get_friction 
get_friction()
Source code in molmo_spaces/env/data_views.py 
381
382
def get_friction(self):
    return [self.mj_model.geom(geom_id).friction for geom_id in self.geom_ids]
get_geom_infos 
get_geom_infos(include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]

Get geom information for this object.

Parameters:

Name Type Description Default
include_descendants bool

If True, includes geoms from all descendant bodies. If False, only includes geoms directly attached to this object's body.

 True
max_geoms int | None

Maximum number of geoms to return. If None, returns all.

 2048

Returns:

Type Description
list[dict[str, object]]

List of dicts with geom info: id, name, position, size, type, type_name
Source code in molmo_spaces/env/data_views.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def get_geom_infos(
    self, include_descendants: bool = True, max_geoms: int | None = 2048
) -> list[dict[str, object]]:
    """Get geom information for this object.

    Args:
        include_descendants: If True, includes geoms from all descendant bodies.
            If False, only includes geoms directly attached to this object's body.
        max_geoms: Maximum number of geoms to return. If None, returns all.

    Returns:
        List of dicts with geom info: id, name, position, size, type, type_name
    """
    # Build set of body IDs to include
    if include_descendants:
        body_ids = set(self.body_ids)
    else:
        body_ids = {self.object_id}

    # Find geoms belonging to these bodies
    geoms: list[dict[str, object]] = []
    count = 0
    for geom_id in range(self.mj_model.ngeom):
        if max_geoms is not None and count >= max_geoms:
            break
        if int(self.mj_model.geom_bodyid[geom_id]) in body_ids:
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id))
            pos = self.mj_data.geom_xpos[geom_id].copy()
            size = self.mj_model.geom_size[geom_id].copy()
            gtype = int(self.mj_model.geom_type[geom_id])
            geoms.append(
                {
                    "id": int(geom_id),
                    "name": name,
                    "position": pos,
                    "size": size,
                    "type": gtype,
                    "type_name": MlSpacesObject.get_geom_type_name(gtype),
                }
            )
            count += 1

    return geoms
get_geom_type_name staticmethod 
get_geom_type_name(geom_type: int) -> str
Source code in molmo_spaces/env/data_views.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
@staticmethod
def get_geom_type_name(geom_type: int) -> str:
    names = {
        0: "PLANE",
        1: "HFIELD",
        2: "SPHERE",
        3: "CAPSULE",
        4: "ELLIPSOID",
        5: "CYLINDER",
        6: "BOX",
        7: "MESH",
    }
    return names.get(int(geom_type), f"UNKNOWN({geom_type})")
get_joint_anchor_position 
get_joint_anchor_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
550
551
552
553
554
555
556
557
558
559
def get_joint_anchor_position(self, i: int) -> np.ndarray:
    # Get the position in world frame of the joint anchor point
    joint_id = self.joint_ids[i]
    body_id = self.mj_model.jnt_bodyid[joint_id]
    local_anchor = self.mj_model.jnt_pos[joint_id]
    body_pos = self.mj_data.xpos[body_id]
    body_quat = self.mj_data.xquat[body_id]
    body_rot = R.from_quat(body_quat, scalar_first=True).as_matrix()
    world_anchor = body_pos + body_rot @ local_anchor
    return world_anchor
get_joint_armature 
get_joint_armature()
Source code in molmo_spaces/env/data_views.py 
570
571
def get_joint_armature(self):
    return [self.mj_model.joint(joint_id).armature for joint_id in self.joint_ids]
get_joint_axis 
get_joint_axis(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
526
527
def get_joint_axis(self, i: int) -> np.ndarray:
    return self.mj_model.joint(self.joint_ids[i]).axis
get_joint_body_orientation 
get_joint_body_orientation(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
545
546
547
548
def get_joint_body_orientation(self, i: int) -> np.ndarray:
    # orientation of the joint body not the root orientation of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    return self.mj_data.xmat[body_id].copy().reshape(3, 3)
get_joint_damping 
get_joint_damping()
Source code in molmo_spaces/env/data_views.py 
567
568
def get_joint_damping(self):
    return [self.mj_model.joint(joint_id).damping for joint_id in self.joint_ids]
get_joint_frictionloss 
get_joint_frictionloss()
Source code in molmo_spaces/env/data_views.py 
561
562
def get_joint_frictionloss(self):
    return [self.mj_model.joint(joint_id).frictionloss for joint_id in self.joint_ids]
get_joint_leaf_body_position 
get_joint_leaf_body_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
535
536
537
538
539
540
541
542
543
def get_joint_leaf_body_position(self, i: int) -> np.ndarray:
    # position of the joint body not the root center of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    # find child body of this joint body
    for child_body_id in range(self.mj_model.nbody):
        if self.mj_model.body(child_body_id).parentid[0] == body_id:
            body_id = child_body_id
            break
    return self.mj_data.xpos[body_id].copy()
get_joint_position 
get_joint_position(i: int) -> float
Source code in molmo_spaces/env/data_views.py 
520
521
def get_joint_position(self, i: int) -> float:
    return self.mj_data.qpos[self.get_joint_qpos_adr(i)].copy()
get_joint_qpos_adr 
get_joint_qpos_adr(i: int) -> int
Source code in molmo_spaces/env/data_views.py 
532
533
def get_joint_qpos_adr(self, i: int) -> int:
    return int(self.mj_model.joint(self.joint_ids[i]).qposadr[0])
get_joint_range 
get_joint_range(i: int) -> tuple[float, float]
Source code in molmo_spaces/env/data_views.py 
523
524
def get_joint_range(self, i: int) -> tuple[float, float]:
    return self.mj_model.joint(self.joint_ids[i]).range
get_joint_stiffness 
get_joint_stiffness()
Source code in molmo_spaces/env/data_views.py 
564
565
def get_joint_stiffness(self):
    return [self.mj_model.joint(joint_id).stiffness for joint_id in self.joint_ids]
get_joint_type 
get_joint_type(i: int) -> mjtJoint
Source code in molmo_spaces/env/data_views.py 
529
530
def get_joint_type(self, i: int) -> mujoco.mjtJoint:
    return self.mj_model.joint(self.joint_ids[i]).type
get_mass 
get_mass()
Source code in molmo_spaces/env/data_views.py 
384
385
def get_mass(self):
    return [self.mj_model.body_mass[body_id] for body_id in self.body_ids]
get_top_level_bodies staticmethod 
get_top_level_bodies(model: MjModel) -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/data_views.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
@staticmethod
def get_top_level_bodies(model: mujoco.MjModel) -> list[int]:
    """Return bodies whose parent is the world body."""
    # Identify world body id: the one with parent < 0
    try:
        world_id = next(
            b for b in range(model.nbody) if MlSpacesObject.body_parent_id(model, b) < 0
        )
    except StopIteration:
        world_id = 0
    tops: list[int] = []
    for b in range(model.nbody):
        try:
            pid = MlSpacesObject.body_parent_id(model, b)
            if pid == world_id:
                tops.append(b)
        except Exception as e:
            print(f"Error getting top-level bodies: {e}")
            continue
    return tops
is_child_name_of staticmethod 
is_child_name_of(parent_name: str, child_name: str) -> bool
Source code in molmo_spaces/env/data_views.py 
430
431
432
@staticmethod
def is_child_name_of(parent_name: str, child_name: str) -> bool:
    return child_name.startswith(parent_name + "_")
is_object_in_gripper 
is_object_in_gripper() -> None
Source code in molmo_spaces/env/data_views.py 
277
278
def is_object_in_gripper(self) -> None:
    pass
is_object_picked_up 
is_object_picked_up() -> None
Source code in molmo_spaces/env/data_views.py 
274
275
def is_object_picked_up(self) -> None:
    pass
set_joint_position 
set_joint_position(i: int, position: float) -> None
Source code in molmo_spaces/env/data_views.py 
573
574
575
576
def set_joint_position(self, i: int, position: float) -> None:
    if isinstance(position, np.ndarray):
        position = position.item()
    self.mj_data.qpos[self.get_joint_qpos_adr(i)] = position

kitchen

Classes:

Name Description
Dishwasher
Oven
Stoveknob

Attributes:

Name Type Description
DISHWASHER
OVEN
STOVEKNOB
DISHWASHER module-attribute 
DISHWASHER = RandomizePhysicalProperties()
OVEN module-attribute 
OVEN = RandomizePhysicalProperties()
STOVEKNOB module-attribute 
STOVEKNOB = RandomizePhysicalProperties()
Dishwasher 
Dishwasher(object_name: str, data: MjData, body_name2id: dict[str, int])

Bases: MlSpacesArticulationObject

Methods:

Name Description
body_name2id
body_parent_id
build_children_lists
find_top_object_body_id
get_ancestors
get_descendants
get_direct_children
get_friction
get_geom_infos

Get geom information for this object.
get_geom_type_name
get_joint_anchor_position
get_joint_armature
get_joint_axis
get_joint_body_orientation
get_joint_damping
get_joint_frictionloss
get_joint_leaf_body_position
get_joint_position
get_joint_qpos_adr
get_joint_range
get_joint_stiffness
get_joint_type
get_mass
get_top_level_bodies

Return bodies whose parent is the world body.
is_child_name_of
is_object_in_gripper
is_object_picked_up
set_joint_position

Attributes:

Name Type Description
aabb_center ndarray
aabb_size ndarray
body_id int
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).
bvh_root
center_of_mass
geom_ids

Get all geom IDs belonging to this object (lazy, cached).
joint_id2name
joint_id2qpos_adr
joint_ids
joint_names
mj_data
mj_model
name str
njoints int
object_id int
object_root_id
pose ndarray
position
quat
Source code in molmo_spaces/env/arena/kitchen.py 
70
71
72
73
74
75
76
77
78
def __init__(self, object_name: str, data: mujoco.MjData, body_name2id: dict[str, int]) -> None:
    super().__init__(object_name, data)
    self._set_friction(DISHWASHER.get_random_value("friction"))
    # self._set_density(DISHWASHER.get_random_value("density"))
    self._set_mass(DISHWASHER.get_random_value("mass"))
    self._set_joint_stiffness(DISHWASHER.get_random_value("joint_stiffness"))
    self._set_joint_damping(DISHWASHER.get_random_value("joint_damping"))
    self._set_joint_armature(DISHWASHER.get_random_value("joint_armature"))
    self._set_joint_frictionloss(DISHWASHER.get_random_value("joint_frictionloss"))
aabb_center property 
aabb_center: ndarray
aabb_size property 
aabb_size: ndarray
body_id cached property 
body_id: int
body_ids cached property 
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).

bvh_root property 
bvh_root
center_of_mass property 
center_of_mass
geom_ids cached property 
geom_ids

Get all geom IDs belonging to this object (lazy, cached).

joint_id2name instance-attribute 
joint_id2name = {}
joint_id2qpos_adr instance-attribute 
joint_id2qpos_adr = {}
joint_ids instance-attribute 
joint_ids = []
joint_names instance-attribute 
joint_names = []
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
njoints property 
njoints: int
object_id property 
object_id: int
object_root_id cached property 
object_root_id
pose property writable 
pose: ndarray
position property 
position
quat property 
quat
body_name2id staticmethod 
body_name2id(model: MjModel) -> dict[str, int]
Source code in molmo_spaces/env/data_views.py 
467
468
469
@staticmethod
def body_name2id(model: mujoco.MjModel) -> dict[str, int]:
    return {model.body(i).name: i for i in range(model.nbody)}
body_parent_id staticmethod 
body_parent_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
387
388
389
390
391
392
393
@staticmethod
def body_parent_id(model: mujoco.MjModel, body_id: int) -> int:
    try:
        return int(model.body_parentid[body_id])
    except Exception as e:
        print(f"Error getting body parent id: {e}")
        return int(model.body(body_id).parentid)
build_children_lists staticmethod 
build_children_lists(model: MjModel) -> list[list[int]]
Source code in molmo_spaces/env/data_views.py 
395
396
397
398
399
400
401
402
@staticmethod
def build_children_lists(model: mujoco.MjModel) -> list[list[int]]:
    children: list[list[int]] = [[] for _ in range(model.nbody)]
    for child_id in range(model.nbody):
        pid = MlSpacesObject.body_parent_id(model, child_id)
        if pid >= 0:
            children[pid].append(child_id)
    return children
find_top_object_body_id staticmethod 
find_top_object_body_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
@staticmethod
def find_top_object_body_id(model: mujoco.MjModel, body_id: int) -> int:
    cur = body_id
    visited = set()
    while True:
        if cur in visited:
            break
        visited.add(cur)
        pid = MlSpacesObject.body_parent_id(model, cur)
        if pid < 0:
            break
        pname = model.body(pid).name
        cname = model.body(cur).name
        if MlSpacesObject.is_child_name_of(pname, cname):
            cur = pid
        else:
            break
    return int(cur)
get_ancestors staticmethod 
get_ancestors(model: MjModel, body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
404
405
406
407
408
409
410
411
412
413
@staticmethod
def get_ancestors(model: mujoco.MjModel, body_id: int) -> list[int]:
    ancestors: list[int] = []
    visited = {body_id}
    pid = MlSpacesObject.body_parent_id(model, body_id)
    while pid >= 0 and pid not in visited:
        ancestors.append(pid)
        visited.add(pid)
        pid = MlSpacesObject.body_parent_id(model, pid)
    return ancestors
get_descendants staticmethod 
get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
415
416
417
418
419
420
421
422
423
424
@staticmethod
def get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]:
    stack = [body_id]
    desc: list[int] = []
    while stack:
        cur = stack.pop()
        for c in children_lists[cur]:
            desc.append(int(c))
            stack.append(c)
    return desc
get_direct_children staticmethod 
get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
426
427
428
@staticmethod
def get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]:
    return list(children_lists[body_id])
get_friction 
get_friction()
Source code in molmo_spaces/env/data_views.py 
381
382
def get_friction(self):
    return [self.mj_model.geom(geom_id).friction for geom_id in self.geom_ids]
get_geom_infos 
get_geom_infos(include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]

Get geom information for this object.

Parameters:

Name Type Description Default
include_descendants bool

If True, includes geoms from all descendant bodies. If False, only includes geoms directly attached to this object's body.

 True
max_geoms int | None

Maximum number of geoms to return. If None, returns all.

 2048

Returns:

Type Description
list[dict[str, object]]

List of dicts with geom info: id, name, position, size, type, type_name
Source code in molmo_spaces/env/data_views.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def get_geom_infos(
    self, include_descendants: bool = True, max_geoms: int | None = 2048
) -> list[dict[str, object]]:
    """Get geom information for this object.

    Args:
        include_descendants: If True, includes geoms from all descendant bodies.
            If False, only includes geoms directly attached to this object's body.
        max_geoms: Maximum number of geoms to return. If None, returns all.

    Returns:
        List of dicts with geom info: id, name, position, size, type, type_name
    """
    # Build set of body IDs to include
    if include_descendants:
        body_ids = set(self.body_ids)
    else:
        body_ids = {self.object_id}

    # Find geoms belonging to these bodies
    geoms: list[dict[str, object]] = []
    count = 0
    for geom_id in range(self.mj_model.ngeom):
        if max_geoms is not None and count >= max_geoms:
            break
        if int(self.mj_model.geom_bodyid[geom_id]) in body_ids:
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id))
            pos = self.mj_data.geom_xpos[geom_id].copy()
            size = self.mj_model.geom_size[geom_id].copy()
            gtype = int(self.mj_model.geom_type[geom_id])
            geoms.append(
                {
                    "id": int(geom_id),
                    "name": name,
                    "position": pos,
                    "size": size,
                    "type": gtype,
                    "type_name": MlSpacesObject.get_geom_type_name(gtype),
                }
            )
            count += 1

    return geoms
get_geom_type_name staticmethod 
get_geom_type_name(geom_type: int) -> str
Source code in molmo_spaces/env/data_views.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
@staticmethod
def get_geom_type_name(geom_type: int) -> str:
    names = {
        0: "PLANE",
        1: "HFIELD",
        2: "SPHERE",
        3: "CAPSULE",
        4: "ELLIPSOID",
        5: "CYLINDER",
        6: "BOX",
        7: "MESH",
    }
    return names.get(int(geom_type), f"UNKNOWN({geom_type})")
get_joint_anchor_position 
get_joint_anchor_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
550
551
552
553
554
555
556
557
558
559
def get_joint_anchor_position(self, i: int) -> np.ndarray:
    # Get the position in world frame of the joint anchor point
    joint_id = self.joint_ids[i]
    body_id = self.mj_model.jnt_bodyid[joint_id]
    local_anchor = self.mj_model.jnt_pos[joint_id]
    body_pos = self.mj_data.xpos[body_id]
    body_quat = self.mj_data.xquat[body_id]
    body_rot = R.from_quat(body_quat, scalar_first=True).as_matrix()
    world_anchor = body_pos + body_rot @ local_anchor
    return world_anchor
get_joint_armature 
get_joint_armature()
Source code in molmo_spaces/env/data_views.py 
570
571
def get_joint_armature(self):
    return [self.mj_model.joint(joint_id).armature for joint_id in self.joint_ids]
get_joint_axis 
get_joint_axis(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
526
527
def get_joint_axis(self, i: int) -> np.ndarray:
    return self.mj_model.joint(self.joint_ids[i]).axis
get_joint_body_orientation 
get_joint_body_orientation(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
545
546
547
548
def get_joint_body_orientation(self, i: int) -> np.ndarray:
    # orientation of the joint body not the root orientation of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    return self.mj_data.xmat[body_id].copy().reshape(3, 3)
get_joint_damping 
get_joint_damping()
Source code in molmo_spaces/env/data_views.py 
567
568
def get_joint_damping(self):
    return [self.mj_model.joint(joint_id).damping for joint_id in self.joint_ids]
get_joint_frictionloss 
get_joint_frictionloss()
Source code in molmo_spaces/env/data_views.py 
561
562
def get_joint_frictionloss(self):
    return [self.mj_model.joint(joint_id).frictionloss for joint_id in self.joint_ids]
get_joint_leaf_body_position 
get_joint_leaf_body_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
535
536
537
538
539
540
541
542
543
def get_joint_leaf_body_position(self, i: int) -> np.ndarray:
    # position of the joint body not the root center of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    # find child body of this joint body
    for child_body_id in range(self.mj_model.nbody):
        if self.mj_model.body(child_body_id).parentid[0] == body_id:
            body_id = child_body_id
            break
    return self.mj_data.xpos[body_id].copy()
get_joint_position 
get_joint_position(i: int) -> float
Source code in molmo_spaces/env/data_views.py 
520
521
def get_joint_position(self, i: int) -> float:
    return self.mj_data.qpos[self.get_joint_qpos_adr(i)].copy()
get_joint_qpos_adr 
get_joint_qpos_adr(i: int) -> int
Source code in molmo_spaces/env/data_views.py 
532
533
def get_joint_qpos_adr(self, i: int) -> int:
    return int(self.mj_model.joint(self.joint_ids[i]).qposadr[0])
get_joint_range 
get_joint_range(i: int) -> tuple[float, float]
Source code in molmo_spaces/env/data_views.py 
523
524
def get_joint_range(self, i: int) -> tuple[float, float]:
    return self.mj_model.joint(self.joint_ids[i]).range
get_joint_stiffness 
get_joint_stiffness()
Source code in molmo_spaces/env/data_views.py 
564
565
def get_joint_stiffness(self):
    return [self.mj_model.joint(joint_id).stiffness for joint_id in self.joint_ids]
get_joint_type 
get_joint_type(i: int) -> mjtJoint
Source code in molmo_spaces/env/data_views.py 
529
530
def get_joint_type(self, i: int) -> mujoco.mjtJoint:
    return self.mj_model.joint(self.joint_ids[i]).type
get_mass 
get_mass()
Source code in molmo_spaces/env/data_views.py 
384
385
def get_mass(self):
    return [self.mj_model.body_mass[body_id] for body_id in self.body_ids]
get_top_level_bodies staticmethod 
get_top_level_bodies(model: MjModel) -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/data_views.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
@staticmethod
def get_top_level_bodies(model: mujoco.MjModel) -> list[int]:
    """Return bodies whose parent is the world body."""
    # Identify world body id: the one with parent < 0
    try:
        world_id = next(
            b for b in range(model.nbody) if MlSpacesObject.body_parent_id(model, b) < 0
        )
    except StopIteration:
        world_id = 0
    tops: list[int] = []
    for b in range(model.nbody):
        try:
            pid = MlSpacesObject.body_parent_id(model, b)
            if pid == world_id:
                tops.append(b)
        except Exception as e:
            print(f"Error getting top-level bodies: {e}")
            continue
    return tops
is_child_name_of staticmethod 
is_child_name_of(parent_name: str, child_name: str) -> bool
Source code in molmo_spaces/env/data_views.py 
430
431
432
@staticmethod
def is_child_name_of(parent_name: str, child_name: str) -> bool:
    return child_name.startswith(parent_name + "_")
is_object_in_gripper 
is_object_in_gripper() -> None
Source code in molmo_spaces/env/data_views.py 
277
278
def is_object_in_gripper(self) -> None:
    pass
is_object_picked_up 
is_object_picked_up() -> None
Source code in molmo_spaces/env/data_views.py 
274
275
def is_object_picked_up(self) -> None:
    pass
set_joint_position 
set_joint_position(i: int, position: float) -> None
Source code in molmo_spaces/env/data_views.py 
573
574
575
576
def set_joint_position(self, i: int, position: float) -> None:
    if isinstance(position, np.ndarray):
        position = position.item()
    self.mj_data.qpos[self.get_joint_qpos_adr(i)] = position
Oven 
Oven(object_name: str, data: MjData, body_name2id: dict[str, int])

Bases: MlSpacesArticulationObject

Methods:

Name Description
body_name2id
body_parent_id
build_children_lists
find_top_object_body_id
get_ancestors
get_descendants
get_direct_children
get_friction
get_geom_infos

Get geom information for this object.
get_geom_type_name
get_joint_anchor_position
get_joint_armature
get_joint_axis
get_joint_body_orientation
get_joint_damping
get_joint_frictionloss
get_joint_leaf_body_position
get_joint_position
get_joint_qpos_adr
get_joint_range
get_joint_stiffness
get_joint_type
get_mass
get_top_level_bodies

Return bodies whose parent is the world body.
is_child_name_of
is_object_in_gripper
is_object_picked_up
set_joint_position

Attributes:

Name Type Description
aabb_center ndarray
aabb_size ndarray
body_id int
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).
bvh_root
center_of_mass
geom_ids

Get all geom IDs belonging to this object (lazy, cached).
joint_id2name
joint_id2qpos_adr
joint_ids
joint_names
mj_data
mj_model
name str
njoints int
object_id int
object_root_id
pose ndarray
position
quat
Source code in molmo_spaces/env/arena/kitchen.py 
57
58
59
60
61
62
63
64
65
66
def __init__(self, object_name: str, data: mujoco.MjData, body_name2id: dict[str, int]) -> None:
    super().__init__(object_name, data)
    self._set_friction(OVEN.get_random_value("friction"))
    # Note: Density cannot be modified at runtime in MuJoCo
    # self._set_density(OVEN.get_random_value("density"))
    self._set_mass(OVEN.get_random_value("mass"))
    self._set_joint_stiffness(OVEN.get_random_value("joint_stiffness"))
    self._set_joint_damping(OVEN.get_random_value("joint_damping"))
    self._set_joint_armature(OVEN.get_random_value("joint_armature"))
    self._set_joint_frictionloss(OVEN.get_random_value("joint_frictionloss"))
aabb_center property 
aabb_center: ndarray
aabb_size property 
aabb_size: ndarray
body_id cached property 
body_id: int
body_ids cached property 
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).

bvh_root property 
bvh_root
center_of_mass property 
center_of_mass
geom_ids cached property 
geom_ids

Get all geom IDs belonging to this object (lazy, cached).

joint_id2name instance-attribute 
joint_id2name = {}
joint_id2qpos_adr instance-attribute 
joint_id2qpos_adr = {}
joint_ids instance-attribute 
joint_ids = []
joint_names instance-attribute 
joint_names = []
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
njoints property 
njoints: int
object_id property 
object_id: int
object_root_id cached property 
object_root_id
pose property writable 
pose: ndarray
position property 
position
quat property 
quat
body_name2id staticmethod 
body_name2id(model: MjModel) -> dict[str, int]
Source code in molmo_spaces/env/data_views.py 
467
468
469
@staticmethod
def body_name2id(model: mujoco.MjModel) -> dict[str, int]:
    return {model.body(i).name: i for i in range(model.nbody)}
body_parent_id staticmethod 
body_parent_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
387
388
389
390
391
392
393
@staticmethod
def body_parent_id(model: mujoco.MjModel, body_id: int) -> int:
    try:
        return int(model.body_parentid[body_id])
    except Exception as e:
        print(f"Error getting body parent id: {e}")
        return int(model.body(body_id).parentid)
build_children_lists staticmethod 
build_children_lists(model: MjModel) -> list[list[int]]
Source code in molmo_spaces/env/data_views.py 
395
396
397
398
399
400
401
402
@staticmethod
def build_children_lists(model: mujoco.MjModel) -> list[list[int]]:
    children: list[list[int]] = [[] for _ in range(model.nbody)]
    for child_id in range(model.nbody):
        pid = MlSpacesObject.body_parent_id(model, child_id)
        if pid >= 0:
            children[pid].append(child_id)
    return children
find_top_object_body_id staticmethod 
find_top_object_body_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
@staticmethod
def find_top_object_body_id(model: mujoco.MjModel, body_id: int) -> int:
    cur = body_id
    visited = set()
    while True:
        if cur in visited:
            break
        visited.add(cur)
        pid = MlSpacesObject.body_parent_id(model, cur)
        if pid < 0:
            break
        pname = model.body(pid).name
        cname = model.body(cur).name
        if MlSpacesObject.is_child_name_of(pname, cname):
            cur = pid
        else:
            break
    return int(cur)
get_ancestors staticmethod 
get_ancestors(model: MjModel, body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
404
405
406
407
408
409
410
411
412
413
@staticmethod
def get_ancestors(model: mujoco.MjModel, body_id: int) -> list[int]:
    ancestors: list[int] = []
    visited = {body_id}
    pid = MlSpacesObject.body_parent_id(model, body_id)
    while pid >= 0 and pid not in visited:
        ancestors.append(pid)
        visited.add(pid)
        pid = MlSpacesObject.body_parent_id(model, pid)
    return ancestors
get_descendants staticmethod 
get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
415
416
417
418
419
420
421
422
423
424
@staticmethod
def get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]:
    stack = [body_id]
    desc: list[int] = []
    while stack:
        cur = stack.pop()
        for c in children_lists[cur]:
            desc.append(int(c))
            stack.append(c)
    return desc
get_direct_children staticmethod 
get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
426
427
428
@staticmethod
def get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]:
    return list(children_lists[body_id])
get_friction 
get_friction()
Source code in molmo_spaces/env/data_views.py 
381
382
def get_friction(self):
    return [self.mj_model.geom(geom_id).friction for geom_id in self.geom_ids]
get_geom_infos 
get_geom_infos(include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]

Get geom information for this object.

Parameters:

Name Type Description Default
include_descendants bool

If True, includes geoms from all descendant bodies. If False, only includes geoms directly attached to this object's body.

 True
max_geoms int | None

Maximum number of geoms to return. If None, returns all.

 2048

Returns:

Type Description
list[dict[str, object]]

List of dicts with geom info: id, name, position, size, type, type_name
Source code in molmo_spaces/env/data_views.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def get_geom_infos(
    self, include_descendants: bool = True, max_geoms: int | None = 2048
) -> list[dict[str, object]]:
    """Get geom information for this object.

    Args:
        include_descendants: If True, includes geoms from all descendant bodies.
            If False, only includes geoms directly attached to this object's body.
        max_geoms: Maximum number of geoms to return. If None, returns all.

    Returns:
        List of dicts with geom info: id, name, position, size, type, type_name
    """
    # Build set of body IDs to include
    if include_descendants:
        body_ids = set(self.body_ids)
    else:
        body_ids = {self.object_id}

    # Find geoms belonging to these bodies
    geoms: list[dict[str, object]] = []
    count = 0
    for geom_id in range(self.mj_model.ngeom):
        if max_geoms is not None and count >= max_geoms:
            break
        if int(self.mj_model.geom_bodyid[geom_id]) in body_ids:
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id))
            pos = self.mj_data.geom_xpos[geom_id].copy()
            size = self.mj_model.geom_size[geom_id].copy()
            gtype = int(self.mj_model.geom_type[geom_id])
            geoms.append(
                {
                    "id": int(geom_id),
                    "name": name,
                    "position": pos,
                    "size": size,
                    "type": gtype,
                    "type_name": MlSpacesObject.get_geom_type_name(gtype),
                }
            )
            count += 1

    return geoms
get_geom_type_name staticmethod 
get_geom_type_name(geom_type: int) -> str
Source code in molmo_spaces/env/data_views.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
@staticmethod
def get_geom_type_name(geom_type: int) -> str:
    names = {
        0: "PLANE",
        1: "HFIELD",
        2: "SPHERE",
        3: "CAPSULE",
        4: "ELLIPSOID",
        5: "CYLINDER",
        6: "BOX",
        7: "MESH",
    }
    return names.get(int(geom_type), f"UNKNOWN({geom_type})")
get_joint_anchor_position 
get_joint_anchor_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
550
551
552
553
554
555
556
557
558
559
def get_joint_anchor_position(self, i: int) -> np.ndarray:
    # Get the position in world frame of the joint anchor point
    joint_id = self.joint_ids[i]
    body_id = self.mj_model.jnt_bodyid[joint_id]
    local_anchor = self.mj_model.jnt_pos[joint_id]
    body_pos = self.mj_data.xpos[body_id]
    body_quat = self.mj_data.xquat[body_id]
    body_rot = R.from_quat(body_quat, scalar_first=True).as_matrix()
    world_anchor = body_pos + body_rot @ local_anchor
    return world_anchor
get_joint_armature 
get_joint_armature()
Source code in molmo_spaces/env/data_views.py 
570
571
def get_joint_armature(self):
    return [self.mj_model.joint(joint_id).armature for joint_id in self.joint_ids]
get_joint_axis 
get_joint_axis(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
526
527
def get_joint_axis(self, i: int) -> np.ndarray:
    return self.mj_model.joint(self.joint_ids[i]).axis
get_joint_body_orientation 
get_joint_body_orientation(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
545
546
547
548
def get_joint_body_orientation(self, i: int) -> np.ndarray:
    # orientation of the joint body not the root orientation of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    return self.mj_data.xmat[body_id].copy().reshape(3, 3)
get_joint_damping 
get_joint_damping()
Source code in molmo_spaces/env/data_views.py 
567
568
def get_joint_damping(self):
    return [self.mj_model.joint(joint_id).damping for joint_id in self.joint_ids]
get_joint_frictionloss 
get_joint_frictionloss()
Source code in molmo_spaces/env/data_views.py 
561
562
def get_joint_frictionloss(self):
    return [self.mj_model.joint(joint_id).frictionloss for joint_id in self.joint_ids]
get_joint_leaf_body_position 
get_joint_leaf_body_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
535
536
537
538
539
540
541
542
543
def get_joint_leaf_body_position(self, i: int) -> np.ndarray:
    # position of the joint body not the root center of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    # find child body of this joint body
    for child_body_id in range(self.mj_model.nbody):
        if self.mj_model.body(child_body_id).parentid[0] == body_id:
            body_id = child_body_id
            break
    return self.mj_data.xpos[body_id].copy()
get_joint_position 
get_joint_position(i: int) -> float
Source code in molmo_spaces/env/data_views.py 
520
521
def get_joint_position(self, i: int) -> float:
    return self.mj_data.qpos[self.get_joint_qpos_adr(i)].copy()
get_joint_qpos_adr 
get_joint_qpos_adr(i: int) -> int
Source code in molmo_spaces/env/data_views.py 
532
533
def get_joint_qpos_adr(self, i: int) -> int:
    return int(self.mj_model.joint(self.joint_ids[i]).qposadr[0])
get_joint_range 
get_joint_range(i: int) -> tuple[float, float]
Source code in molmo_spaces/env/data_views.py 
523
524
def get_joint_range(self, i: int) -> tuple[float, float]:
    return self.mj_model.joint(self.joint_ids[i]).range
get_joint_stiffness 
get_joint_stiffness()
Source code in molmo_spaces/env/data_views.py 
564
565
def get_joint_stiffness(self):
    return [self.mj_model.joint(joint_id).stiffness for joint_id in self.joint_ids]
get_joint_type 
get_joint_type(i: int) -> mjtJoint
Source code in molmo_spaces/env/data_views.py 
529
530
def get_joint_type(self, i: int) -> mujoco.mjtJoint:
    return self.mj_model.joint(self.joint_ids[i]).type
get_mass 
get_mass()
Source code in molmo_spaces/env/data_views.py 
384
385
def get_mass(self):
    return [self.mj_model.body_mass[body_id] for body_id in self.body_ids]
get_top_level_bodies staticmethod 
get_top_level_bodies(model: MjModel) -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/data_views.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
@staticmethod
def get_top_level_bodies(model: mujoco.MjModel) -> list[int]:
    """Return bodies whose parent is the world body."""
    # Identify world body id: the one with parent < 0
    try:
        world_id = next(
            b for b in range(model.nbody) if MlSpacesObject.body_parent_id(model, b) < 0
        )
    except StopIteration:
        world_id = 0
    tops: list[int] = []
    for b in range(model.nbody):
        try:
            pid = MlSpacesObject.body_parent_id(model, b)
            if pid == world_id:
                tops.append(b)
        except Exception as e:
            print(f"Error getting top-level bodies: {e}")
            continue
    return tops
is_child_name_of staticmethod 
is_child_name_of(parent_name: str, child_name: str) -> bool
Source code in molmo_spaces/env/data_views.py 
430
431
432
@staticmethod
def is_child_name_of(parent_name: str, child_name: str) -> bool:
    return child_name.startswith(parent_name + "_")
is_object_in_gripper 
is_object_in_gripper() -> None
Source code in molmo_spaces/env/data_views.py 
277
278
def is_object_in_gripper(self) -> None:
    pass
is_object_picked_up 
is_object_picked_up() -> None
Source code in molmo_spaces/env/data_views.py 
274
275
def is_object_picked_up(self) -> None:
    pass
set_joint_position 
set_joint_position(i: int, position: float) -> None
Source code in molmo_spaces/env/data_views.py 
573
574
575
576
def set_joint_position(self, i: int, position: float) -> None:
    if isinstance(position, np.ndarray):
        position = position.item()
    self.mj_data.qpos[self.get_joint_qpos_adr(i)] = position
Stoveknob 
Stoveknob(object_name: str, data: MjData, body_name2id: dict[str, int])

Bases: MlSpacesArticulationObject

Methods:

Name Description
body_name2id
body_parent_id
build_children_lists
find_top_object_body_id
get_ancestors
get_descendants
get_direct_children
get_friction
get_geom_infos

Get geom information for this object.
get_geom_type_name
get_joint_anchor_position
get_joint_armature
get_joint_axis
get_joint_body_orientation
get_joint_damping
get_joint_frictionloss
get_joint_leaf_body_position
get_joint_position
get_joint_qpos_adr
get_joint_range
get_joint_stiffness
get_joint_type
get_mass
get_top_level_bodies

Return bodies whose parent is the world body.
is_child_name_of
is_object_in_gripper
is_object_picked_up
set_joint_position

Attributes:

Name Type Description
aabb_center ndarray
aabb_size ndarray
body_id int
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).
bvh_root
center_of_mass
geom_ids

Get all geom IDs belonging to this object (lazy, cached).
joint_id2name
joint_id2qpos_adr
joint_ids
joint_names
mj_data
mj_model
name str
njoints int
object_id int
object_root_id
pose ndarray
position
quat
Source code in molmo_spaces/env/arena/kitchen.py 
82
83
84
85
86
87
88
89
def __init__(self, object_name: str, data: mujoco.MjData, body_name2id: dict[str, int]) -> None:
    super().__init__(object_name, data)
    self._set_friction(STOVEKNOB.get_random_value("friction"))
    # self._set_density(STOVEKNOB.get_random_value("density"))
    self._set_joint_stiffness(STOVEKNOB.get_random_value("joint_stiffness"))
    self._set_joint_damping(STOVEKNOB.get_random_value("joint_damping"))
    self._set_joint_armature(STOVEKNOB.get_random_value("joint_armature"))
    self._set_joint_frictionloss(STOVEKNOB.get_random_value("joint_frictionloss"))
aabb_center property 
aabb_center: ndarray
aabb_size property 
aabb_size: ndarray
body_id cached property 
body_id: int
body_ids cached property 
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).

bvh_root property 
bvh_root
center_of_mass property 
center_of_mass
geom_ids cached property 
geom_ids

Get all geom IDs belonging to this object (lazy, cached).

joint_id2name instance-attribute 
joint_id2name = {}
joint_id2qpos_adr instance-attribute 
joint_id2qpos_adr = {}
joint_ids instance-attribute 
joint_ids = []
joint_names instance-attribute 
joint_names = []
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
njoints property 
njoints: int
object_id property 
object_id: int
object_root_id cached property 
object_root_id
pose property writable 
pose: ndarray
position property 
position
quat property 
quat
body_name2id staticmethod 
body_name2id(model: MjModel) -> dict[str, int]
Source code in molmo_spaces/env/data_views.py 
467
468
469
@staticmethod
def body_name2id(model: mujoco.MjModel) -> dict[str, int]:
    return {model.body(i).name: i for i in range(model.nbody)}
body_parent_id staticmethod 
body_parent_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
387
388
389
390
391
392
393
@staticmethod
def body_parent_id(model: mujoco.MjModel, body_id: int) -> int:
    try:
        return int(model.body_parentid[body_id])
    except Exception as e:
        print(f"Error getting body parent id: {e}")
        return int(model.body(body_id).parentid)
build_children_lists staticmethod 
build_children_lists(model: MjModel) -> list[list[int]]
Source code in molmo_spaces/env/data_views.py 
395
396
397
398
399
400
401
402
@staticmethod
def build_children_lists(model: mujoco.MjModel) -> list[list[int]]:
    children: list[list[int]] = [[] for _ in range(model.nbody)]
    for child_id in range(model.nbody):
        pid = MlSpacesObject.body_parent_id(model, child_id)
        if pid >= 0:
            children[pid].append(child_id)
    return children
find_top_object_body_id staticmethod 
find_top_object_body_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
@staticmethod
def find_top_object_body_id(model: mujoco.MjModel, body_id: int) -> int:
    cur = body_id
    visited = set()
    while True:
        if cur in visited:
            break
        visited.add(cur)
        pid = MlSpacesObject.body_parent_id(model, cur)
        if pid < 0:
            break
        pname = model.body(pid).name
        cname = model.body(cur).name
        if MlSpacesObject.is_child_name_of(pname, cname):
            cur = pid
        else:
            break
    return int(cur)
get_ancestors staticmethod 
get_ancestors(model: MjModel, body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
404
405
406
407
408
409
410
411
412
413
@staticmethod
def get_ancestors(model: mujoco.MjModel, body_id: int) -> list[int]:
    ancestors: list[int] = []
    visited = {body_id}
    pid = MlSpacesObject.body_parent_id(model, body_id)
    while pid >= 0 and pid not in visited:
        ancestors.append(pid)
        visited.add(pid)
        pid = MlSpacesObject.body_parent_id(model, pid)
    return ancestors
get_descendants staticmethod 
get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
415
416
417
418
419
420
421
422
423
424
@staticmethod
def get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]:
    stack = [body_id]
    desc: list[int] = []
    while stack:
        cur = stack.pop()
        for c in children_lists[cur]:
            desc.append(int(c))
            stack.append(c)
    return desc
get_direct_children staticmethod 
get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
426
427
428
@staticmethod
def get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]:
    return list(children_lists[body_id])
get_friction 
get_friction()
Source code in molmo_spaces/env/data_views.py 
381
382
def get_friction(self):
    return [self.mj_model.geom(geom_id).friction for geom_id in self.geom_ids]
get_geom_infos 
get_geom_infos(include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]

Get geom information for this object.

Parameters:

Name Type Description Default
include_descendants bool

If True, includes geoms from all descendant bodies. If False, only includes geoms directly attached to this object's body.

 True
max_geoms int | None

Maximum number of geoms to return. If None, returns all.

 2048

Returns:

Type Description
list[dict[str, object]]

List of dicts with geom info: id, name, position, size, type, type_name
Source code in molmo_spaces/env/data_views.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def get_geom_infos(
    self, include_descendants: bool = True, max_geoms: int | None = 2048
) -> list[dict[str, object]]:
    """Get geom information for this object.

    Args:
        include_descendants: If True, includes geoms from all descendant bodies.
            If False, only includes geoms directly attached to this object's body.
        max_geoms: Maximum number of geoms to return. If None, returns all.

    Returns:
        List of dicts with geom info: id, name, position, size, type, type_name
    """
    # Build set of body IDs to include
    if include_descendants:
        body_ids = set(self.body_ids)
    else:
        body_ids = {self.object_id}

    # Find geoms belonging to these bodies
    geoms: list[dict[str, object]] = []
    count = 0
    for geom_id in range(self.mj_model.ngeom):
        if max_geoms is not None and count >= max_geoms:
            break
        if int(self.mj_model.geom_bodyid[geom_id]) in body_ids:
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id))
            pos = self.mj_data.geom_xpos[geom_id].copy()
            size = self.mj_model.geom_size[geom_id].copy()
            gtype = int(self.mj_model.geom_type[geom_id])
            geoms.append(
                {
                    "id": int(geom_id),
                    "name": name,
                    "position": pos,
                    "size": size,
                    "type": gtype,
                    "type_name": MlSpacesObject.get_geom_type_name(gtype),
                }
            )
            count += 1

    return geoms
get_geom_type_name staticmethod 
get_geom_type_name(geom_type: int) -> str
Source code in molmo_spaces/env/data_views.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
@staticmethod
def get_geom_type_name(geom_type: int) -> str:
    names = {
        0: "PLANE",
        1: "HFIELD",
        2: "SPHERE",
        3: "CAPSULE",
        4: "ELLIPSOID",
        5: "CYLINDER",
        6: "BOX",
        7: "MESH",
    }
    return names.get(int(geom_type), f"UNKNOWN({geom_type})")
get_joint_anchor_position 
get_joint_anchor_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
550
551
552
553
554
555
556
557
558
559
def get_joint_anchor_position(self, i: int) -> np.ndarray:
    # Get the position in world frame of the joint anchor point
    joint_id = self.joint_ids[i]
    body_id = self.mj_model.jnt_bodyid[joint_id]
    local_anchor = self.mj_model.jnt_pos[joint_id]
    body_pos = self.mj_data.xpos[body_id]
    body_quat = self.mj_data.xquat[body_id]
    body_rot = R.from_quat(body_quat, scalar_first=True).as_matrix()
    world_anchor = body_pos + body_rot @ local_anchor
    return world_anchor
get_joint_armature 
get_joint_armature()
Source code in molmo_spaces/env/data_views.py 
570
571
def get_joint_armature(self):
    return [self.mj_model.joint(joint_id).armature for joint_id in self.joint_ids]
get_joint_axis 
get_joint_axis(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
526
527
def get_joint_axis(self, i: int) -> np.ndarray:
    return self.mj_model.joint(self.joint_ids[i]).axis
get_joint_body_orientation 
get_joint_body_orientation(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
545
546
547
548
def get_joint_body_orientation(self, i: int) -> np.ndarray:
    # orientation of the joint body not the root orientation of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    return self.mj_data.xmat[body_id].copy().reshape(3, 3)
get_joint_damping 
get_joint_damping()
Source code in molmo_spaces/env/data_views.py 
567
568
def get_joint_damping(self):
    return [self.mj_model.joint(joint_id).damping for joint_id in self.joint_ids]
get_joint_frictionloss 
get_joint_frictionloss()
Source code in molmo_spaces/env/data_views.py 
561
562
def get_joint_frictionloss(self):
    return [self.mj_model.joint(joint_id).frictionloss for joint_id in self.joint_ids]
get_joint_leaf_body_position 
get_joint_leaf_body_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
535
536
537
538
539
540
541
542
543
def get_joint_leaf_body_position(self, i: int) -> np.ndarray:
    # position of the joint body not the root center of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    # find child body of this joint body
    for child_body_id in range(self.mj_model.nbody):
        if self.mj_model.body(child_body_id).parentid[0] == body_id:
            body_id = child_body_id
            break
    return self.mj_data.xpos[body_id].copy()
get_joint_position 
get_joint_position(i: int) -> float
Source code in molmo_spaces/env/data_views.py 
520
521
def get_joint_position(self, i: int) -> float:
    return self.mj_data.qpos[self.get_joint_qpos_adr(i)].copy()
get_joint_qpos_adr 
get_joint_qpos_adr(i: int) -> int
Source code in molmo_spaces/env/data_views.py 
532
533
def get_joint_qpos_adr(self, i: int) -> int:
    return int(self.mj_model.joint(self.joint_ids[i]).qposadr[0])
get_joint_range 
get_joint_range(i: int) -> tuple[float, float]
Source code in molmo_spaces/env/data_views.py 
523
524
def get_joint_range(self, i: int) -> tuple[float, float]:
    return self.mj_model.joint(self.joint_ids[i]).range
get_joint_stiffness 
get_joint_stiffness()
Source code in molmo_spaces/env/data_views.py 
564
565
def get_joint_stiffness(self):
    return [self.mj_model.joint(joint_id).stiffness for joint_id in self.joint_ids]
get_joint_type 
get_joint_type(i: int) -> mjtJoint
Source code in molmo_spaces/env/data_views.py 
529
530
def get_joint_type(self, i: int) -> mujoco.mjtJoint:
    return self.mj_model.joint(self.joint_ids[i]).type
get_mass 
get_mass()
Source code in molmo_spaces/env/data_views.py 
384
385
def get_mass(self):
    return [self.mj_model.body_mass[body_id] for body_id in self.body_ids]
get_top_level_bodies staticmethod 
get_top_level_bodies(model: MjModel) -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/data_views.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
@staticmethod
def get_top_level_bodies(model: mujoco.MjModel) -> list[int]:
    """Return bodies whose parent is the world body."""
    # Identify world body id: the one with parent < 0
    try:
        world_id = next(
            b for b in range(model.nbody) if MlSpacesObject.body_parent_id(model, b) < 0
        )
    except StopIteration:
        world_id = 0
    tops: list[int] = []
    for b in range(model.nbody):
        try:
            pid = MlSpacesObject.body_parent_id(model, b)
            if pid == world_id:
                tops.append(b)
        except Exception as e:
            print(f"Error getting top-level bodies: {e}")
            continue
    return tops
is_child_name_of staticmethod 
is_child_name_of(parent_name: str, child_name: str) -> bool
Source code in molmo_spaces/env/data_views.py 
430
431
432
@staticmethod
def is_child_name_of(parent_name: str, child_name: str) -> bool:
    return child_name.startswith(parent_name + "_")
is_object_in_gripper 
is_object_in_gripper() -> None
Source code in molmo_spaces/env/data_views.py 
277
278
def is_object_in_gripper(self) -> None:
    pass
is_object_picked_up 
is_object_picked_up() -> None
Source code in molmo_spaces/env/data_views.py 
274
275
def is_object_picked_up(self) -> None:
    pass
set_joint_position 
set_joint_position(i: int, position: float) -> None
Source code in molmo_spaces/env/data_views.py 
573
574
575
576
def set_joint_position(self, i: int, position: float) -> None:
    if isinstance(position, np.ndarray):
        position = position.item()
    self.mj_data.qpos[self.get_joint_qpos_adr(i)] = position

procthor_types

Classes:

Name Description
PhysicalProperty
RandomizePhysicalProperties
PhysicalProperty 
PhysicalProperty(enabled: bool = False, range: tuple[float, float] = (1.0, 3.0))

Attributes:

Name Type Description
enabled
range
Source code in molmo_spaces/env/arena/procthor_types.py 
5
6
7
def __init__(self, enabled: bool = False, range: tuple[float, float] = (1.0, 3.0)) -> None:
    self.enabled = enabled
    self.range = range
enabled instance-attribute 
enabled = enabled
range instance-attribute 
range = range
RandomizePhysicalProperties 
RandomizePhysicalProperties(property_name: str = None, range: tuple[float, float] = None)

Methods:

Name Description
enable_property

Enable randomization for a specific property.
get_random_value

Get a random value for a specific property.
get_range

Get the randomization range for a specific property.
is_enabled

Check if randomization is enabled for a specific property.
set_range

Set the randomization range for a specific property.

Attributes:

Name Type Description
properties
randomize_density bool
randomize_friction bool
randomize_joint_armature bool
randomize_joint_damping bool
randomize_joint_frictionloss bool
randomize_joint_limited bool
randomize_joint_stiffness bool
randomize_mass bool
range_density tuple[float, float]
range_friction tuple[float, float]
range_joint_armature tuple[float, float]
range_joint_damping tuple[float, float]
range_joint_frictionloss tuple[float, float]
range_joint_limited tuple[float, float]
range_joint_stiffness tuple[float, float]
range_mass tuple[float, float]
Source code in molmo_spaces/env/arena/procthor_types.py 
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
def __init__(self, property_name: str = None, range: tuple[float, float] = None) -> None:
    self.properties = {
        "friction": PhysicalProperty(),
        "density": PhysicalProperty(),
        "mass": PhysicalProperty(),
        "joint": {
            "stiffness": PhysicalProperty(),
            "damping": PhysicalProperty(),
            "limited": PhysicalProperty(),
            "armature": PhysicalProperty(),
            "frictionloss": PhysicalProperty(),
        },
    }

    # Initialize with specific property if provided
    if property_name is not None:
        if property_name.startswith("joint_"):
            joint_prop = property_name.replace("joint_", "")
            if joint_prop in self.properties["joint"]:
                self.properties["joint"][joint_prop].enabled = True
                if range is not None:
                    self.properties["joint"][joint_prop].range = range
        elif property_name in self.properties:
            self.properties[property_name].enabled = True
            if range is not None:
                self.properties[property_name].range = range
properties instance-attribute 
properties = {'friction': PhysicalProperty(), 'density': PhysicalProperty(), 'mass': PhysicalProperty(), 'joint': {'stiffness': PhysicalProperty(), 'damping': PhysicalProperty(), 'limited': PhysicalProperty(), 'armature': PhysicalProperty(), 'frictionloss': PhysicalProperty()}}
randomize_density property writable 
randomize_density: bool
randomize_friction property writable 
randomize_friction: bool
randomize_joint_armature property writable 
randomize_joint_armature: bool
randomize_joint_damping property writable 
randomize_joint_damping: bool
randomize_joint_frictionloss property writable 
randomize_joint_frictionloss: bool
randomize_joint_limited property writable 
randomize_joint_limited: bool
randomize_joint_stiffness property writable 
randomize_joint_stiffness: bool
randomize_mass property writable 
randomize_mass: bool
range_density property writable 
range_density: tuple[float, float]
range_friction property writable 
range_friction: tuple[float, float]
range_joint_armature property writable 
range_joint_armature: tuple[float, float]
range_joint_damping property writable 
range_joint_damping: tuple[float, float]
range_joint_frictionloss property writable 
range_joint_frictionloss: tuple[float, float]
range_joint_limited property writable 
range_joint_limited: tuple[float, float]
range_joint_stiffness property writable 
range_joint_stiffness: tuple[float, float]
range_mass property writable 
range_mass: tuple[float, float]
enable_property 
enable_property(property_name: str, joint_property: str = None) -> None

Enable randomization for a specific property.

Parameters:

Name Type Description Default
property_name str

Name of the property to enable ('friction', 'density', 'mass', 'joint')

 required
joint_property str

If property_name is 'joint', specify which joint property

 None
Source code in molmo_spaces/env/arena/procthor_types.py 
169
170
171
172
173
174
175
176
177
178
179
def enable_property(self, property_name: str, joint_property: str = None) -> None:
    """Enable randomization for a specific property.

    Args:
        property_name: Name of the property to enable ('friction', 'density', 'mass', 'joint')
        joint_property: If property_name is 'joint', specify which joint property
    """
    if property_name == "joint" and joint_property:
        self.properties["joint"][joint_property].enabled = True
    elif property_name in self.properties:
        self.properties[property_name].enabled = True
get_random_value 
get_random_value(property_name: str) -> float

Get a random value for a specific property.

Parameters:

Name Type Description Default
property_name str

Name of the property to get random value for

 required
joint_property

If property_name is 'joint', specify which joint property

 required

Returns:

Name Type Description
float float

A random value for the property
Source code in molmo_spaces/env/arena/procthor_types.py 
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def get_random_value(self, property_name: str) -> float:
    """Get a random value for a specific property.

    Args:
        property_name: Name of the property to get random value for
        joint_property: If property_name is 'joint', specify which joint property

    Returns:
        float: A random value for the property
    """
    if property_name.startswith("joint_"):
        joint_property = property_name.replace("joint_", "")
        return np.random.uniform(
            self.properties["joint"][joint_property].range[0],
            self.properties["joint"][joint_property].range[1],
        )
    return np.random.uniform(
        self.properties[property_name].range[0], self.properties[property_name].range[1]
    )
get_range 
get_range(property_name: str, joint_property: str = None) -> tuple[float, float]

Get the randomization range for a specific property.

Parameters:

Name Type Description Default
property_name str

Name of the property to get range for

 required
joint_property str

If property_name is 'joint', specify which joint property

 None

Returns:

Type Description
tuple[float, float]

Tuple[float, float]: The (min, max) range for the property
Source code in molmo_spaces/env/arena/procthor_types.py 
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def get_range(self, property_name: str, joint_property: str = None) -> tuple[float, float]:
    """Get the randomization range for a specific property.

    Args:
        property_name: Name of the property to get range for
        joint_property: If property_name is 'joint', specify which joint property

    Returns:
        Tuple[float, float]: The (min, max) range for the property
    """
    if property_name == "joint" and joint_property:
        return self.properties["joint"][joint_property].range
    return (
        self.properties[property_name].range if property_name in self.properties else (1.0, 3.0)
    )
is_enabled 
is_enabled(property_name: str, joint_property: str = None) -> bool

Check if randomization is enabled for a specific property.

Parameters:

Name Type Description Default
property_name str

Name of the property to check

 required
joint_property str

If property_name is 'joint', specify which joint property

 None

Returns:

Name Type Description
bool bool

True if randomization is enabled for the property
Source code in molmo_spaces/env/arena/procthor_types.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
def is_enabled(self, property_name: str, joint_property: str = None) -> bool:
    """Check if randomization is enabled for a specific property.

    Args:
        property_name: Name of the property to check
        joint_property: If property_name is 'joint', specify which joint property

    Returns:
        bool: True if randomization is enabled for the property
    """
    if property_name == "joint" and joint_property:
        return self.properties["joint"][joint_property].enabled
    return self.properties[property_name].enabled if property_name in self.properties else False
set_range 
set_range(property_name: str, range: tuple[float, float], joint_property: str = None) -> None

Set the randomization range for a specific property.

Parameters:

Name Type Description Default
property_name str

Name of the property to set range for

 required
range tuple[float, float]

Tuple of (min, max) values

 required
joint_property str

If property_name is 'joint', specify which joint property

 None
Source code in molmo_spaces/env/arena/procthor_types.py 
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def set_range(
    self, property_name: str, range: tuple[float, float], joint_property: str = None
) -> None:
    """Set the randomization range for a specific property.

    Args:
        property_name: Name of the property to set range for
        range: Tuple of (min, max) values
        joint_property: If property_name is 'joint', specify which joint property
    """
    if property_name == "joint" and joint_property:
        self.properties["joint"][joint_property].range = range
    elif property_name in self.properties:
        self.properties[property_name].range = range

randomization

Modules:

Name Description
dynamics
lighting
test_randomizers

Env wrapper to randomize the scene.
texture
dynamics

Classes:

Name Description
DynamicsRandomizer

Randomizer for dynamics properties of MlSpacesObject instances.
DynamicsRandomizer 
DynamicsRandomizer(random_state: RandomState | None = None, randomize_friction: bool = True, randomize_mass: bool = True, randomize_inertia: bool = True, friction_perturbation_ratio: float = 0.1, mass_perturbation_ratio: float = 0.1, inertia_perturbation_ratio: float = 0.1)

Randomizer for dynamics properties of MlSpacesObject instances.

Randomizes object-level properties: friction (of geoms), mass, and inertia. Note: Density cannot be modified at runtime in MuJoCo, only mass can be changed.

Parameters:

Name Type Description Default
random_state RandomState | None

Random state for reproducibility. If None, uses global numpy random state.

 None
randomize_friction bool

If True, randomizes geom friction

 True
randomize_mass bool

If True, randomizes object mass

 True
randomize_inertia bool

If True, randomizes object inertia

 True
friction_perturbation_ratio float

Relative magnitude of friction randomization

 0.1
mass_perturbation_ratio float

Relative magnitude of mass randomization

 0.1
inertia_perturbation_ratio float

Relative magnitude of inertia randomization

 0.1

Methods:

Name Description
randomize_object

Randomize dynamics properties of a single MlSpacesObject.
randomize_objects

Randomize dynamics properties of multiple MlSpacesObject instances.
restore_object

Restore default values for a single object.
restore_objects

Restore default values for multiple objects.

Attributes:

Name Type Description
friction_perturbation_ratio
inertia_perturbation_ratio
mass_perturbation_ratio
random_state
randomize_friction
randomize_inertia
randomize_mass
Source code in molmo_spaces/env/arena/randomization/dynamics.py 
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
def __init__(
    self,
    random_state: np.random.RandomState | None = None,
    randomize_friction: bool = True,
    randomize_mass: bool = True,
    randomize_inertia: bool = True,
    friction_perturbation_ratio: float = 0.1,
    mass_perturbation_ratio: float = 0.1,
    inertia_perturbation_ratio: float = 0.1,
):
    if random_state is None:
        self.random_state = np.random
    else:
        self.random_state = random_state

    self.randomize_friction = randomize_friction
    self.randomize_mass = randomize_mass
    self.randomize_inertia = randomize_inertia

    # Validate that perturbation ratios are < 1 to ensure positive values after perturbation
    if friction_perturbation_ratio >= 1.0:
        raise ValueError(
            f"friction_perturbation_ratio must be < 1.0, got {friction_perturbation_ratio}"
        )
    if mass_perturbation_ratio >= 1.0:
        raise ValueError(
            f"mass_perturbation_ratio must be < 1.0, got {mass_perturbation_ratio}"
        )
    if inertia_perturbation_ratio >= 1.0:
        raise ValueError(
            f"inertia_perturbation_ratio must be < 1.0, got {inertia_perturbation_ratio}"
        )

    self.friction_perturbation_ratio = friction_perturbation_ratio
    self.mass_perturbation_ratio = mass_perturbation_ratio
    self.inertia_perturbation_ratio = inertia_perturbation_ratio

    # Will be populated when randomize_objects is called
    # Keyed by object_id (int) instead of name since objects may not have names
    self._defaults: dict[int, dict] = {}
friction_perturbation_ratio instance-attribute 
friction_perturbation_ratio = friction_perturbation_ratio
inertia_perturbation_ratio instance-attribute 
inertia_perturbation_ratio = inertia_perturbation_ratio
mass_perturbation_ratio instance-attribute 
mass_perturbation_ratio = mass_perturbation_ratio
random_state instance-attribute 
random_state = random
randomize_friction instance-attribute 
randomize_friction = randomize_friction
randomize_inertia instance-attribute 
randomize_inertia = randomize_inertia
randomize_mass instance-attribute 
randomize_mass = randomize_mass
randomize_object 
randomize_object(obj: MlSpacesObject) -> None

Randomize dynamics properties of a single MlSpacesObject.

Parameters:

Name Type Description Default
obj MlSpacesObject

MlSpacesObject instance to randomize

 required
Source code in molmo_spaces/env/arena/randomization/dynamics.py 
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
def randomize_object(self, obj: "MlSpacesObject") -> None:
    """
    Randomize dynamics properties of a single MlSpacesObject.

    Args:
        obj: MlSpacesObject instance to randomize
    """
    from molmo_spaces.env.data_views import MlSpacesObject

    if not isinstance(obj, MlSpacesObject):
        raise TypeError(f"Expected MlSpacesObject, got {type(obj)}")

    # Save defaults if not already saved
    object_id = obj.object_id
    if object_id not in self._defaults:
        self._save_object_defaults(obj)

    model = obj.mj_model
    defaults = self._defaults[object_id]

    # Randomize mass for all descendant bodies (maintaining proportional distribution)
    if self.randomize_mass:
        # Use object's body_ids property which uses descendant_bodies()
        body_ids = obj.body_ids
        body_masses = defaults.get("body_masses", {object_id: defaults["mass"]})
        total_mass = defaults["mass"]

        if total_mass > 0:  # Only randomize if object has non-zero mass
            # Apply perturbation to total mass
            # Since mass_perturbation_ratio < 1, (1.0 + perturbation) > 0, so new_total_mass > 0
            perturbation = self.random_state.uniform(
                -self.mass_perturbation_ratio, self.mass_perturbation_ratio
            )
            new_total_mass = total_mass * (1.0 + perturbation)

            # Distribute the new total mass proportionally across all descendant bodies
            if len(body_ids) == 1:
                # Single body: use object's _set_mass method
                obj._set_mass(new_total_mass)
            else:
                # Multiple bodies: maintain proportional distribution across all descendant bodies
                mass_ratio = new_total_mass / total_mass
                for bid in body_ids:
                    if bid in body_masses and body_masses[bid] > 0:
                        new_body_mass = body_masses[bid] * mass_ratio
                        model.body_mass[bid] = new_body_mass

    # Randomize inertia for all bodies (root and all children)
    if self.randomize_inertia:
        # Use object's body_ids property which uses descendant_bodies()
        body_ids = obj.body_ids
        body_inertias = defaults.get("body_inertias", {object_id: defaults["inertia"]})

        for bid in body_ids:
            if bid in body_inertias:
                current_inertia = body_inertias[bid]
                if np.any(current_inertia > 0):  # Only randomize if inertia exists
                    # Since inertia_perturbation_ratio < 1, (1.0 + perturbation) > 0, so new_inertia > 0
                    perturbation = self.random_state.uniform(
                        -self.inertia_perturbation_ratio,
                        self.inertia_perturbation_ratio,
                        size=3,
                    )
                    new_inertia = current_inertia * (1.0 + perturbation)
                    model.body_inertia[bid] = new_inertia

    # Randomize friction using object's _set_friction method
    # Compute average friction value across all geoms and apply perturbation
    if self.randomize_friction:
        if defaults["geom_frictions"]:
            # Get average default friction (using first component - sliding friction)
            avg_default_friction = np.mean(
                [
                    friction[0]
                    for friction in defaults["geom_frictions"].values()
                    if np.any(friction > 0)
                ]
            )
            if avg_default_friction > 0:
                # Apply perturbation to average friction
                # Since friction_perturbation_ratio < 1, (1.0 + perturbation) > 0, so new_friction > 0
                perturbation = self.random_state.uniform(
                    -self.friction_perturbation_ratio,
                    self.friction_perturbation_ratio,
                )
                new_friction = avg_default_friction * (1.0 + perturbation)

                # Ensure _geom_ids exists for _set_friction method (geom_ids is a property)
                if not hasattr(obj, "_geom_ids") or obj._geom_ids is None:
                    object_root_id = model.body(object_id).rootid[0]
                    obj._geom_ids = []
                    for geom_id in range(model.ngeom):
                        geom_body_id = model.geom(geom_id).bodyid.item()
                        geom_root_id = model.body(geom_body_id).rootid[0]
                        if geom_root_id == object_root_id:
                            obj._geom_ids.append(geom_id)

                # Use object's _set_friction method which sets friction for all geoms
                obj._set_friction(new_friction)
randomize_objects 
randomize_objects(objects: list[MlSpacesObject]) -> None

Randomize dynamics properties of multiple MlSpacesObject instances.

Parameters:

Name Type Description Default
objects list[MlSpacesObject]

List of MlSpacesObject instances to randomize

 required
Source code in molmo_spaces/env/arena/randomization/dynamics.py 
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
def randomize_objects(self, objects: list["MlSpacesObject"]) -> None:
    """
    Randomize dynamics properties of multiple MlSpacesObject instances.

    Args:
        objects: List of MlSpacesObject instances to randomize
    """
    # Save defaults for all objects first
    for obj in objects:
        if obj.object_id not in self._defaults:
            self._save_object_defaults(obj)

    # Then randomize each object
    for obj in objects:
        self.randomize_object(obj)

    # Forward pass to propagate changes
    if objects:
        # All objects should share the same model/data
        model = objects[0].mj_model
        data = objects[0].mj_data
        mujoco.mj_forward(model, data)
restore_object 
restore_object(obj: MlSpacesObject) -> None

Restore default values for a single object.

Parameters:

Name Type Description Default
obj MlSpacesObject

MlSpacesObject instance to restore

 required
Source code in molmo_spaces/env/arena/randomization/dynamics.py 
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
def restore_object(self, obj: "MlSpacesObject") -> None:
    """
    Restore default values for a single object.

    Args:
        obj: MlSpacesObject instance to restore
    """
    object_id = obj.object_id
    if object_id not in self._defaults:
        return  # No defaults saved for this object

    model = obj.mj_model
    defaults = self._defaults[object_id]

    # Restore mass using object's _set_mass method
    # body_ids property uses descendant_bodies() internally
    if "body_masses" in defaults:
        # Sum up all body masses to get total mass
        total_mass = sum(defaults["body_masses"].values())
        obj._set_mass(total_mass)
    else:
        # Fallback to old behavior if body_masses not saved
        obj._set_mass(defaults["mass"])

    # Restore inertia for all bodies (root and all children)
    # Note: No object method for inertia, so use direct access
    # Use object's body_ids property which uses descendant_bodies()
    body_ids = obj.body_ids
    if "body_inertias" in defaults:
        for body_id in body_ids:
            if body_id in defaults["body_inertias"]:
                model.body_inertia[body_id] = defaults["body_inertias"][body_id]
    else:
        # Fallback to old behavior if body_inertias not saved
        model.body_inertia[object_id] = defaults["inertia"]

    # Restore friction using object's _set_friction method
    # Compute average default friction to restore
    if defaults["geom_frictions"]:
        avg_default_friction = np.mean(
            [
                friction[0]
                for friction in defaults["geom_frictions"].values()
                if np.any(friction > 0)
            ]
        )
        if avg_default_friction > 0:
            # Ensure _geom_ids exists for _set_friction method (geom_ids is a property)
            if not hasattr(obj, "_geom_ids") or obj._geom_ids is None:
                object_root_id = model.body(object_id).rootid[0]
                obj._geom_ids = []
                for geom_id in range(model.ngeom):
                    geom_body_id = model.geom(geom_id).bodyid.item()
                    geom_root_id = model.body(geom_body_id).rootid[0]
                    if geom_root_id == object_root_id:
                        obj._geom_ids.append(geom_id)

            obj._set_friction(avg_default_friction)

    # Forward pass
    mujoco.mj_forward(model, obj.mj_data)
restore_objects 
restore_objects(objects: list[MlSpacesObject]) -> None

Restore default values for multiple objects.

Parameters:

Name Type Description Default
objects list[MlSpacesObject]

List of MlSpacesObject instances to restore

 required
Source code in molmo_spaces/env/arena/randomization/dynamics.py 
295
296
297
298
299
300
301
302
303
def restore_objects(self, objects: list["MlSpacesObject"]) -> None:
    """
    Restore default values for multiple objects.

    Args:
        objects: List of MlSpacesObject instances to restore
    """
    for obj in objects:
        self.restore_object(obj)
lighting

Classes:

Name Description
LightingRandomizer

Randomizer for lighting properties in MuJoCo simulations.
LightingRandomizer 
LightingRandomizer(model: MjModel, random_state: RandomState | None = None, light_names: list[str] | None = None, randomize_position: bool = True, randomize_direction: bool = True, randomize_specular: bool = True, randomize_ambient: bool = True, randomize_diffuse: bool = True, randomize_active: bool = True, position_perturbation_size: float = 0.1, direction_perturbation_size: float = 0.35, specular_perturbation_size: float = 0.1, ambient_perturbation_size: float = 0.1, diffuse_perturbation_size: float = 0.1)

Randomizer for lighting properties in MuJoCo simulations.

Based on the mujoco-py LightingModder implementation, adapted to work with MjModel and MjData directly (instead of MjSim).

Parameters:

Name Type Description Default
model MjModel

MuJoCo model

 required
random_state RandomState | None

Random state for reproducibility. If None, uses global numpy random state.

 None
light_names list[str] | None

List of light names to randomize. If None, randomizes all lights in the model.

 None
randomize_position bool

If True, randomizes light position

 True
randomize_direction bool

If True, randomizes light direction

 True
randomize_specular bool

If True, randomizes specular color

 True
randomize_ambient bool

If True, randomizes ambient color

 True
randomize_diffuse bool

If True, randomizes diffuse color

 True
randomize_active bool

If True, randomizes whether light is active

 True
position_perturbation_size float

Magnitude of position randomization

 0.1
direction_perturbation_size float

Magnitude of direction randomization in radians

 0.35
specular_perturbation_size float

Magnitude of specular color randomization

 0.1
ambient_perturbation_size float

Magnitude of ambient color randomization

 0.1
diffuse_perturbation_size float

Magnitude of diffuse color randomization

 0.1
Note

MjData should be passed to the randomize() method, not to init.

Methods:

Name Description
get_active

Get active state of a specific light.
get_ambient

Get ambient color of a specific light.
get_diffuse

Get diffuse color of a specific light.
get_dir

Get direction of a specific light.
get_pos

Get position of a specific light.
get_specular

Get specular color of a specific light.
randomize

Randomize all enabled light properties.
restore_defaults

Restore saved default light parameter values.
save_defaults

Save default light parameter values from the current model state.
set_active

Set active state of a specific light.
set_ambient

Set ambient color of a specific light.
set_diffuse

Set diffuse color of a specific light.
set_dir

Set direction of a specific light.
set_pos

Set position of a specific light.
set_specular

Set specular color of a specific light.
update_model

Update the model reference.

Attributes:

Name Type Description
ambient_perturbation_size
diffuse_perturbation_size
direction_perturbation_size
light_ids
model
position_perturbation_size
random_state
randomize_active
randomize_ambient
randomize_diffuse
randomize_direction
randomize_position
randomize_specular
specular_perturbation_size
Source code in molmo_spaces/env/arena/randomization/lighting.py 
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
def __init__(
    self,
    model: MjModel,
    random_state: np.random.RandomState | None = None,
    light_names: list[str] | None = None,
    randomize_position: bool = True,
    randomize_direction: bool = True,
    randomize_specular: bool = True,
    randomize_ambient: bool = True,
    randomize_diffuse: bool = True,
    randomize_active: bool = True,
    position_perturbation_size: float = 0.1,
    direction_perturbation_size: float = 0.35,  # ~20 degrees
    specular_perturbation_size: float = 0.1,
    ambient_perturbation_size: float = 0.1,
    diffuse_perturbation_size: float = 0.1,
):
    self.model = model

    if random_state is None:
        self.random_state = np.random
    else:
        self.random_state = random_state

    # Get light IDs from model (use IDs directly since lights may not have names)
    if light_names is None:
        # Use all light IDs (0 to nlight-1)
        self.light_ids = list(range(model.nlight))
    else:
        # Convert light names to IDs
        self.light_ids = []
        for name in light_names:
            light_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, name)
            if light_id >= 0:
                self.light_ids.append(light_id)

    # Debug: print detected lights
    if model.nlight == 0:
        print(f"   Warning: No lights in model (model.nlight={model.nlight})")
    else:
        print(
            f"   Found {len(self.light_ids)} lights (model.nlight={model.nlight}): IDs {self.light_ids}"
        )

    self.randomize_position = randomize_position
    self.randomize_direction = randomize_direction
    self.randomize_specular = randomize_specular
    self.randomize_ambient = randomize_ambient
    self.randomize_diffuse = randomize_diffuse
    self.randomize_active = randomize_active

    self.position_perturbation_size = position_perturbation_size
    self.direction_perturbation_size = direction_perturbation_size
    self.specular_perturbation_size = specular_perturbation_size
    self.ambient_perturbation_size = ambient_perturbation_size
    self.diffuse_perturbation_size = diffuse_perturbation_size

    # Enable shadow casting for all lights by default (required for shadows to appear)
    # This ensures shadows are visible even when lighting randomization is disabled
    for light_id in self.light_ids:
        if hasattr(self.model, "light_castshadow"):
            self.model.light_castshadow[light_id] = 1

    self.save_defaults()
ambient_perturbation_size instance-attribute 
ambient_perturbation_size = ambient_perturbation_size
diffuse_perturbation_size instance-attribute 
diffuse_perturbation_size = diffuse_perturbation_size
direction_perturbation_size instance-attribute 
direction_perturbation_size = direction_perturbation_size
light_ids instance-attribute 
light_ids = list(range(nlight))
model instance-attribute 
model = model
position_perturbation_size instance-attribute 
position_perturbation_size = position_perturbation_size
random_state instance-attribute 
random_state = random
randomize_active instance-attribute 
randomize_active = randomize_active
randomize_ambient instance-attribute 
randomize_ambient = randomize_ambient
randomize_diffuse instance-attribute 
randomize_diffuse = randomize_diffuse
randomize_direction instance-attribute 
randomize_direction = randomize_direction
randomize_position instance-attribute 
randomize_position = randomize_position
randomize_specular instance-attribute 
randomize_specular = randomize_specular
specular_perturbation_size instance-attribute 
specular_perturbation_size = specular_perturbation_size
get_active 
get_active(light_id: int) -> int

Get active state of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required

Returns:

Name Type Description
int int

1 if active, 0 if inactive
Source code in molmo_spaces/env/arena/randomization/lighting.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
def get_active(self, light_id: int) -> int:
    """
    Get active state of a specific light.

    Args:
        light_id (int): ID of the light

    Returns:
        int: 1 if active, 0 if inactive
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    return int(self.model.light_active[light_id])
get_ambient 
get_ambient(light_id: int) -> ndarray

Get ambient color of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required

Returns:

Type Description
ndarray

np.ndarray: (r, g, b) ambient color
Source code in molmo_spaces/env/arena/randomization/lighting.py 
418
419
420
421
422
423
424
425
426
427
428
429
430
def get_ambient(self, light_id: int) -> np.ndarray:
    """
    Get ambient color of a specific light.

    Args:
        light_id (int): ID of the light

    Returns:
        np.ndarray: (r, g, b) ambient color
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    return np.array(self.model.light_ambient[light_id])
get_diffuse 
get_diffuse(light_id: int) -> ndarray

Get diffuse color of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required

Returns:

Type Description
ndarray

np.ndarray: (r, g, b) diffuse color
Source code in molmo_spaces/env/arena/randomization/lighting.py 
447
448
449
450
451
452
453
454
455
456
457
458
459
def get_diffuse(self, light_id: int) -> np.ndarray:
    """
    Get diffuse color of a specific light.

    Args:
        light_id (int): ID of the light

    Returns:
        np.ndarray: (r, g, b) diffuse color
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    return np.array(self.model.light_diffuse[light_id])
get_dir 
get_dir(light_id: int) -> ndarray

Get direction of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required

Returns:

Type Description
ndarray

np.ndarray: (x, y, z) direction vector
Source code in molmo_spaces/env/arena/randomization/lighting.py 
330
331
332
333
334
335
336
337
338
339
340
341
342
def get_dir(self, light_id: int) -> np.ndarray:
    """
    Get direction of a specific light.

    Args:
        light_id (int): ID of the light

    Returns:
        np.ndarray: (x, y, z) direction vector
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    return np.array(self.model.light_dir[light_id])
get_pos 
get_pos(light_id: int) -> ndarray

Get position of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required

Returns:

Type Description
ndarray

np.ndarray: (x, y, z) position
Source code in molmo_spaces/env/arena/randomization/lighting.py 
301
302
303
304
305
306
307
308
309
310
311
312
313
def get_pos(self, light_id: int) -> np.ndarray:
    """
    Get position of a specific light.

    Args:
        light_id (int): ID of the light

    Returns:
        np.ndarray: (x, y, z) position
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    return np.array(self.model.light_pos[light_id])
get_specular 
get_specular(light_id: int) -> ndarray

Get specular color of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required

Returns:

Type Description
ndarray

np.ndarray: (r, g, b) specular color
Source code in molmo_spaces/env/arena/randomization/lighting.py 
389
390
391
392
393
394
395
396
397
398
399
400
401
def get_specular(self, light_id: int) -> np.ndarray:
    """
    Get specular color of a specific light.

    Args:
        light_id (int): ID of the light

    Returns:
        np.ndarray: (r, g, b) specular color
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    return np.array(self.model.light_specular[light_id])
randomize 
randomize(data: MjData | None = None)

Randomize all enabled light properties.

Parameters:

Name Type Description Default
data MjData | None

MuJoCo data for forward pass. If None, forward pass is skipped.

 None
Source code in molmo_spaces/env/arena/randomization/lighting.py 
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def randomize(self, data: MjData | None = None):
    """
    Randomize all enabled light properties.

    Args:
        data (MjData | None): MuJoCo data for forward pass. If None, forward pass is skipped.
    """
    # Track which lights are active before randomization
    active_lights_before = []
    for light_id in self.light_ids:
        if self.model.light_active[light_id] > 0:
            active_lights_before.append(light_id)

    for light_id in self.light_ids:
        if self.randomize_position:
            self._randomize_position(light_id)

        if self.randomize_direction:
            self._randomize_direction(light_id)

        if self.randomize_specular:
            self._randomize_specular(light_id)

        if self.randomize_ambient:
            self._randomize_ambient(light_id)

        if self.randomize_diffuse:
            self._randomize_diffuse(light_id)

        if self.randomize_active:
            self._randomize_active(light_id)

        # Enable shadow casting for all lights (required for shadows to appear)
        if hasattr(self.model, "light_castshadow"):
            self.model.light_castshadow[light_id] = 1

    # Ensure at least one light is active for shadows to be visible
    # If randomize_active turned off all lights, re-enable at least one
    active_lights_after = []
    for light_id in self.light_ids:
        if self.model.light_active[light_id] > 0:
            active_lights_after.append(light_id)

    if len(active_lights_after) == 0 and len(self.light_ids) > 0:
        # All lights were turned off - re-enable at least one (prefer the first one that was active before)
        if active_lights_before:
            # Re-enable the first light that was active before
            self.model.light_active[active_lights_before[0]] = 1
        else:
            # If no lights were active before, enable the first light
            self.model.light_active[self.light_ids[0]] = 1

    # Forward pass to propagate changes
    if data is not None:
        mujoco.mj_forward(self.model, data)
restore_defaults 
restore_defaults()

Restore saved default light parameter values.

Source code in molmo_spaces/env/arena/randomization/lighting.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def restore_defaults(self):
    """
    Restore saved default light parameter values.
    """
    for light_id in self.light_ids:
        self.set_pos(light_id, self._defaults[light_id]["pos"])
        self.set_dir(light_id, self._defaults[light_id]["dir"])
        self.set_specular(light_id, self._defaults[light_id]["specular"])
        self.set_ambient(light_id, self._defaults[light_id]["ambient"])
        self.set_diffuse(light_id, self._defaults[light_id]["diffuse"])
        self.set_active(light_id, self._defaults[light_id]["active"])
        # Restore castshadow if it was saved
        if hasattr(self.model, "light_castshadow") and "castshadow" in self._defaults[light_id]:
            self.model.light_castshadow[light_id] = self._defaults[light_id]["castshadow"]
save_defaults 
save_defaults()

Save default light parameter values from the current model state.

Source code in molmo_spaces/env/arena/randomization/lighting.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def save_defaults(self):
    """
    Save default light parameter values from the current model state.
    """
    self._defaults = {light_id: {} for light_id in self.light_ids}
    for light_id in self.light_ids:
        self._defaults[light_id]["pos"] = np.array(self.model.light_pos[light_id])
        self._defaults[light_id]["dir"] = np.array(self.model.light_dir[light_id])
        self._defaults[light_id]["specular"] = np.array(self.model.light_specular[light_id])
        self._defaults[light_id]["ambient"] = np.array(self.model.light_ambient[light_id])
        self._defaults[light_id]["diffuse"] = np.array(self.model.light_diffuse[light_id])
        self._defaults[light_id]["active"] = int(self.model.light_active[light_id])
        # Save castshadow if available (enables shadow casting)
        if hasattr(self.model, "light_castshadow"):
            self._defaults[light_id]["castshadow"] = int(self.model.light_castshadow[light_id])
set_active 
set_active(light_id: int, value: int)

Set active state of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required
value int

1 for active, 0 for inactive

 required
Source code in molmo_spaces/env/arena/randomization/lighting.py 
377
378
379
380
381
382
383
384
385
386
387
def set_active(self, light_id: int, value: int):
    """
    Set active state of a specific light.

    Args:
        light_id (int): ID of the light
        value (int): 1 for active, 0 for inactive
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    self.model.light_active[light_id] = value
set_ambient 
set_ambient(light_id: int, value: ndarray)

Set ambient color of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required
value ndarray

(r, g, b) ambient color

 required
Source code in molmo_spaces/env/arena/randomization/lighting.py 
432
433
434
435
436
437
438
439
440
441
442
443
444
445
def set_ambient(self, light_id: int, value: np.ndarray):
    """
    Set ambient color of a specific light.

    Args:
        light_id (int): ID of the light
        value (np.ndarray): (r, g, b) ambient color
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    value = np.asarray(value)
    if value.shape != (3,):
        raise ValueError(f"Expected 3-dim value, got shape {value.shape}")
    self.model.light_ambient[light_id] = np.clip(value, 0.0, 1.0)
set_diffuse 
set_diffuse(light_id: int, value: ndarray)

Set diffuse color of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required
value ndarray

(r, g, b) diffuse color

 required
Source code in molmo_spaces/env/arena/randomization/lighting.py 
461
462
463
464
465
466
467
468
469
470
471
472
473
474
def set_diffuse(self, light_id: int, value: np.ndarray):
    """
    Set diffuse color of a specific light.

    Args:
        light_id (int): ID of the light
        value (np.ndarray): (r, g, b) diffuse color
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    value = np.asarray(value)
    if value.shape != (3,):
        raise ValueError(f"Expected 3-dim value, got shape {value.shape}")
    self.model.light_diffuse[light_id] = np.clip(value, 0.0, 1.0)
set_dir 
set_dir(light_id: int, value: ndarray)

Set direction of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required
value ndarray

(x, y, z) direction vector

 required
Source code in molmo_spaces/env/arena/randomization/lighting.py 
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def set_dir(self, light_id: int, value: np.ndarray):
    """
    Set direction of a specific light.

    Args:
        light_id (int): ID of the light
        value (np.ndarray): (x, y, z) direction vector
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    value = np.asarray(value)
    if value.shape != (3,):
        raise ValueError(f"Expected 3-dim value, got shape {value.shape}")
    # Normalize direction vector
    norm = np.linalg.norm(value)
    if norm > 0:
        value = value / norm
    self.model.light_dir[light_id] = value
set_pos 
set_pos(light_id: int, value: ndarray)

Set position of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required
value ndarray

(x, y, z) position

 required
Source code in molmo_spaces/env/arena/randomization/lighting.py 
315
316
317
318
319
320
321
322
323
324
325
326
327
328
def set_pos(self, light_id: int, value: np.ndarray):
    """
    Set position of a specific light.

    Args:
        light_id (int): ID of the light
        value (np.ndarray): (x, y, z) position
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    value = np.asarray(value)
    if value.shape != (3,):
        raise ValueError(f"Expected 3-dim value, got shape {value.shape}")
    self.model.light_pos[light_id] = value
set_specular 
set_specular(light_id: int, value: ndarray)

Set specular color of a specific light.

Parameters:

Name Type Description Default
light_id int

ID of the light

 required
value ndarray

(r, g, b) specular color

 required
Source code in molmo_spaces/env/arena/randomization/lighting.py 
403
404
405
406
407
408
409
410
411
412
413
414
415
416
def set_specular(self, light_id: int, value: np.ndarray):
    """
    Set specular color of a specific light.

    Args:
        light_id (int): ID of the light
        value (np.ndarray): (r, g, b) specular color
    """
    if light_id < 0 or light_id >= self.model.nlight:
        raise ValueError(f"Invalid light ID: {light_id}")
    value = np.asarray(value)
    if value.shape != (3,):
        raise ValueError(f"Expected 3-dim value, got shape {value.shape}")
    self.model.light_specular[light_id] = np.clip(value, 0.0, 1.0)
update_model 
update_model(model: MjModel)

Update the model reference.

Parameters:

Name Type Description Default
model MjModel

New MuJoCo model

 required
Source code in molmo_spaces/env/arena/randomization/lighting.py 
476
477
478
479
480
481
482
483
484
def update_model(self, model: MjModel):
    """
    Update the model reference.

    Args:
        model (MjModel): New MuJoCo model
    """
    self.model = model
    self.save_defaults()
test_randomizers

Env wrapper to randomize the scene.

TEST RUNTIME RANDOMIZATION
  • lightings (randomize intensity and color and position)
  • textures (switch among preloaded textures and randomize material properties)
  • dynamics (object mass, inertia, friction, density, etc.)

Functions:

Name Description
test

Test function to test the randomizers and visualize using MuJoCo passive viewer.
test 
test(scene_path: str, texture_paths: list[str] | None = None, relaunch_viewer_on_randomize: bool = True) -> None

Test function to test the randomizers and visualize using MuJoCo passive viewer.

Parameters:

Name Type Description Default
scene_path str

Path to scene XML file. Model and data will be loaded from this file.

 required
texture_paths list[str] | None

Optional list of texture file paths for texture randomization. If None, uses textures already loaded in the scene XML.

 None
relaunch_viewer_on_randomize bool

If True, relaunches the viewer after each randomization to ensure model property changes (textures, lights) are visible. Default is True.

 True

Loads the scene, applies randomization, and visualizes the results.

Source code in molmo_spaces/env/arena/randomization/test_randomizers.py 
 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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
def test(
    scene_path: str,
    texture_paths: list[str] | None = None,
    relaunch_viewer_on_randomize: bool = True,
) -> None:
    """
    Test function to test the randomizers and visualize using MuJoCo passive viewer.

    Args:
        scene_path: Path to scene XML file. Model and data will be loaded from this file.
        texture_paths: Optional list of texture file paths for texture randomization.
            If None, uses textures already loaded in the scene XML.
        relaunch_viewer_on_randomize: If True, relaunches the viewer after each randomization
            to ensure model property changes (textures, lights) are visible. Default is True.

    Loads the scene, applies randomization, and visualizes the results.
    """
    import os
    import time

    # Load model and data from scene_path
    if not os.path.exists(scene_path):
        raise FileNotFoundError(f"Scene path does not exist: {scene_path}")

    spec = mujoco.MjSpec.from_file(scene_path)
    setup_empty_materials(spec)
    model = spec.compile()

    print(f"Loading model from: {scene_path}")
    # model = MjModel.from_xml_path(scene_path)
    data = MjData(model)
    mujoco.mj_forward(model, data)
    print(f"Loaded model: {model.ngeom} geoms, {model.nlight} lights, {model.ntex} textures")

    print("=" * 60)
    print("Testing Scene Randomizers")
    print("=" * 60)

    # Initialize randomizers
    print("\n1. Initializing randomizers...")
    lighting_randomizer = LightingRandomizer(
        model=model,
        randomize_position=True,
        randomize_direction=True,
        randomize_specular=True,
        randomize_ambient=True,
        randomize_diffuse=True,
        randomize_active=True,
        position_perturbation_size=0.2,
        direction_perturbation_size=0.5,
    )

    # Many materials in MuJoCo scenes have textures even if mat_texid isn't accessible
    # If we have textures in the model, we should enable randomization
    # Default: use textures from model (texture_paths=None)
    # Only use external files if explicitly provided
    # Enable if we found textures via mat_texid, OR if we have textures and materials (fallback)
    enable_texture_randomization = model.ntex > 0 or (
        texture_paths is not None and len(texture_paths) > 0
    )

    # Load scene_metadata from scene_path
    from molmo_spaces.utils.scene_metadata_utils import get_scene_metadata

    scene_metadata = get_scene_metadata(scene_path)
    if scene_metadata is None:
        print(f"   Warning: Could not load scene metadata from {scene_path}")

    texture_randomizer = TextureRandomizer(
        model=model,
        randomize_geom_rgba=True,
        randomize_material_rgba=True,
        randomize_material_specular=True,
        randomize_material_shininess=True,
        randomize_texture=enable_texture_randomization,
        texture_paths=texture_paths,  # None = use model textures (default)
        scene_metadata=scene_metadata,
        rgba_perturbation_size=0.2,
    )

    if enable_texture_randomization:
        if texture_paths:
            print(f"   Texture randomization: using {len(texture_paths)} external texture files")
        else:
            # When using model textures, texture_bitmaps is empty (we use texture_ids for on-demand extraction)
            num_texture_ids = (
                len(texture_randomizer.texture_ids)
                if hasattr(texture_randomizer, "texture_ids")
                else 0
            )
            print(
                f"   Texture randomization: using {num_texture_ids} textures from model (default, on-demand extraction)"
            )
            if num_texture_ids == 0 and model.ntex > 0:
                print(
                    f"   Warning: Could not find 2D textures in model despite {model.ntex} textures existing"
                )

    dynamics_randomizer = DynamicsRandomizer(
        randomize_friction=True,
        randomize_mass=True,
        randomize_inertia=True,
        mass_perturbation_ratio=0.2,
        friction_perturbation_ratio=0.2,
        inertia_perturbation_ratio=0.2,
    )
    print("   ✓ Randomizers initialized")

    # Create MlSpacesObject instances for dynamics randomization
    # Find all bodies with free joints or any non-fixed joints
    from molmo_spaces.env.arena.arena_utils import get_all_bodies_with_joints_as_mlspaces_objects

    test_objects = get_all_bodies_with_joints_as_mlspaces_objects(model, data)

    # Count total bodies with joints for logging
    bodies_with_joints_count = 0
    for body_id in range(1, model.nbody):
        jnt_adr = model.body_jntadr[body_id]
        if jnt_adr >= 0:
            bodies_with_joints_count += 1

    print(
        f"   ✓ Found {bodies_with_joints_count} bodies with joints ({len(test_objects)} successfully created as MlSpacesObject)"
    )
    if bodies_with_joints_count > 0 and len(test_objects) == 0:
        print(
            f"   ⚠ Warning: Found {bodies_with_joints_count} bodies with joints but couldn't create MlSpacesObject instances"
        )

    # Apply initial randomization before launching visualizer
    print("\n2. Applying initial randomization...")
    start_time = time.perf_counter()
    lighting_randomizer.randomize(data)
    lighting_time = time.perf_counter() - start_time
    print(f"   ✓ Lighting randomized ({lighting_time * 1000:.2f} ms)")

    start_time = time.perf_counter()
    if texture_randomizer.randomize_texture:
        # texture_randomizer.randomize(data)
        texture_randomizer.randomize_by_category(data)
        texture_time = time.perf_counter() - start_time
        print(f"   ✓ Textures randomized ({texture_time * 1000:.2f} ms)")
    else:
        texture_time = 0.0
        print("   ⚠ Texture randomization is disabled")

    dynamics_time = 0.0
    if test_objects:
        start_time = time.perf_counter()
        dynamics_randomizer.randomize_objects(test_objects)
        dynamics_time = time.perf_counter() - start_time
        print(f"   ✓ Dynamics randomized ({dynamics_time * 1000:.2f} ms)")
    else:
        print("   ⚠ No objects found for dynamics randomization")

    total_time = lighting_time + texture_time + dynamics_time
    print(f"   Total randomization time: {total_time * 1000:.2f} ms")

    # Ensure forward pass is done after initial randomization
    mujoco.mj_forward(model, data)

    print("\n3. Starting visualization...")
    if relaunch_viewer_on_randomize:
        print("   Viewer will be relaunched after each randomization")
    print("   Press ESC or close window to exit")

    # Launch passive viewer
    viewer = None
    step_count = 0
    should_relaunch = False

    while True:
        # Launch or relaunch viewer
        if viewer is None or should_relaunch:
            if viewer is not None:
                with contextlib.suppress(Exception):
                    viewer.close()
                viewer = None
                time.sleep(0.2)  # Allow viewer to fully close

            viewer = mujoco.viewer.launch_passive(model, data)
            viewer.cam.distance = 5.0
            viewer.cam.azimuth = 45.0
            viewer.cam.elevation = -30.0
            viewer.cam.lookat[:] = np.array([0.0, 0.0, 0.5])
            should_relaunch = False

        # Check if viewer is still running
        if viewer is None or not viewer.is_running():
            break

        # Step simulation
        mujoco.mj_step(model, data)

        # Sync viewer
        viewer.sync()

        # Re-randomize every 500 steps (including step 0 for immediate randomization)
        if step_count % 500 == 0:
            if step_count > 0:
                print(f"\n   Re-randomizing at step {step_count}...")

            # Capture light properties before randomization
            light_before = {}
            if lighting_randomizer.light_ids:
                for light_id in lighting_randomizer.light_ids:
                    light_before[light_id] = {
                        "pos": np.array(lighting_randomizer.get_pos(light_id)),
                        "dir": np.array(lighting_randomizer.get_dir(light_id)),
                        "specular": np.array(lighting_randomizer.get_specular(light_id)),
                        "ambient": np.array(lighting_randomizer.get_ambient(light_id)),
                        "diffuse": np.array(lighting_randomizer.get_diffuse(light_id)),
                        "active": lighting_randomizer.get_active(light_id),
                    }

            # Capture dynamics properties before randomization
            dynamics_before = {}
            if test_objects:
                for obj in test_objects:
                    model = obj.mj_model
                    object_id = obj.object_id
                    object_root_id = model.body(object_id).rootid[0]

                    # Get all bodies belonging to this object
                    body_ids = [object_id]
                    for body_id in range(model.nbody):
                        if body_id != object_id:
                            body_root_id = model.body(body_id).rootid[0]
                            if body_root_id == object_root_id:
                                body_ids.append(body_id)

                    # Get total mass of the object (including all descendant bodies)
                    total_mass = float(model.body_subtreemass[object_id])
                    # Inertia is typically on the root body
                    inertia = np.array(model.body_inertia[object_id])

                    # Get friction for all geoms
                    geom_frictions = {}
                    for geom_id in range(model.ngeom):
                        geom_body_id = model.geom(geom_id).bodyid.item()
                        geom_root_id = model.body(geom_body_id).rootid[0]
                        if geom_root_id == object_root_id:
                            geom_frictions[geom_id] = np.array(model.geom_friction[geom_id])

                    dynamics_before[obj.name] = {
                        "mass": total_mass,
                        "inertia": inertia,
                        "body_ids": body_ids,
                        "body_masses": {bid: float(model.body_mass[bid]) for bid in body_ids},
                        "geom_frictions": geom_frictions,
                    }

            start_time = time.perf_counter()
            lighting_randomizer.randomize(data)
            lighting_time = time.perf_counter() - start_time

            start_time = time.perf_counter()
            if texture_randomizer.randomize_texture:
                # texture_randomizer.randomize(data)
                texture_randomizer.randomize_by_category(data)
                texture_time = time.perf_counter() - start_time
            else:
                texture_time = 0.0

            dynamics_time = 0.0
            if test_objects:
                start_time = time.perf_counter()
                dynamics_randomizer.randomize_objects(test_objects)
                dynamics_time = time.perf_counter() - start_time

            total_time = lighting_time + texture_time + dynamics_time
            print(
                f"   Randomization complete: {total_time * 1000:.2f} ms "
                f"(lighting: {lighting_time * 1000:.2f} ms, "
                f"texture: {texture_time * 1000:.2f} ms, "
                f"dynamics: {dynamics_time * 1000:.2f} ms)"
            )

            # Print light property diffs
            if light_before and step_count > 0:
                print("\n   Light property changes:")
                for light_id, before in light_before.items():
                    after_pos = np.array(lighting_randomizer.get_pos(light_id))
                    after_dir = np.array(lighting_randomizer.get_dir(light_id))
                    after_specular = np.array(lighting_randomizer.get_specular(light_id))
                    after_ambient = np.array(lighting_randomizer.get_ambient(light_id))
                    after_diffuse = np.array(lighting_randomizer.get_diffuse(light_id))
                    after_active = lighting_randomizer.get_active(light_id)

                    pos_diff = after_pos - before["pos"]
                    dir_diff = after_dir - before["dir"]
                    specular_diff = after_specular - before["specular"]
                    ambient_diff = after_ambient - before["ambient"]
                    diffuse_diff = after_diffuse - before["diffuse"]
                    active_diff = after_active - before["active"]

                    print(f"     Light {light_id}:")
                    print(f"       pos diff: {pos_diff} (norm: {np.linalg.norm(pos_diff):.6f})")
                    print(f"       dir diff: {dir_diff} (norm: {np.linalg.norm(dir_diff):.6f})")
                    print(
                        f"       specular diff: {specular_diff} (norm: {np.linalg.norm(specular_diff):.6f})"
                    )
                    print(
                        f"       ambient diff: {ambient_diff} (norm: {np.linalg.norm(ambient_diff):.6f})"
                    )
                    print(
                        f"       diffuse diff: {diffuse_diff} (norm: {np.linalg.norm(diffuse_diff):.6f})"
                    )
                    print(f"       active diff: {active_diff}")

            # Print dynamics property diffs
            if dynamics_before and test_objects and step_count > 0:
                print("\n   Dynamics property changes:")
                for obj in test_objects:
                    if obj.name in dynamics_before:
                        model = obj.mj_model
                        object_id = obj.object_id
                        before = dynamics_before[obj.name]

                        # Get total mass of the object (including all descendant bodies)
                        body_ids = before.get("body_ids", [object_id])
                        after_mass = float(model.body_subtreemass[object_id])
                        after_inertia = np.array(model.body_inertia[object_id])

                        mass_diff = after_mass - before["mass"]
                        inertia_diff = after_inertia - before["inertia"]

                        print(f"     {obj.name}:")
                        print(
                            f"       mass diff: {mass_diff:.6f} ({before['mass']:.6f} -> {after_mass:.6f})"
                        )
                        if len(body_ids) > 1:
                            print(
                                f"         (object has {len(body_ids)} bodies, showing total mass)"
                            )
                        print(
                            f"       inertia diff: {inertia_diff} (norm: {np.linalg.norm(inertia_diff):.6f})"
                        )

                        # Get friction diffs
                        object_root_id = model.body(object_id).rootid[0]
                        friction_diffs = []
                        for geom_id in range(model.ngeom):
                            geom_body_id = model.geom(geom_id).bodyid.item()
                            geom_root_id = model.body(geom_body_id).rootid[0]
                            if geom_root_id == object_root_id:
                                if geom_id in before["geom_frictions"]:
                                    after_friction = np.array(model.geom_friction[geom_id])
                                    friction_diff = (
                                        after_friction - before["geom_frictions"][geom_id]
                                    )
                                    friction_diffs.append((geom_id, friction_diff))

                        if friction_diffs:
                            print("       friction diffs (geom_id: diff):")
                            for geom_id, friction_diff in friction_diffs[:5]:  # Show first 5
                                print(
                                    f"         geom_{geom_id}: {friction_diff} (norm: {np.linalg.norm(friction_diff):.6f})"
                                )
                            if len(friction_diffs) > 5:
                                print(f"         ... and {len(friction_diffs) - 5} more geoms")

            # Ensure forward pass is done after all randomizations
            mujoco.mj_forward(model, data)
            viewer.sync()

            # If relaunch_viewer_on_randomize is True, close and relaunch viewer to force refresh
            if relaunch_viewer_on_randomize and step_count > 0:
                if viewer is not None:
                    with contextlib.suppress(Exception):
                        viewer.close()
                    viewer = None
                should_relaunch = True
                time.sleep(0.1)
                continue

        step_count += 1

        # Small sleep to control frame rate
        time.sleep(0.01)

    # Clean up viewer if still open
    if viewer is not None:
        viewer.close()

    print("\n✓ Test completed successfully!")
texture

Classes:

Name Description
TextureRandomizer

Randomizer for geom colors, material properties, and textures in MuJoCo simulations.

Functions:

Name Description
assign_texture_to_material

Assign a texture to a material's RGB role.
create_empty_texture

Create an empty texture in the spec using the placeholder file.
create_placeholder_texture_file

Create a temporary placeholder texture file for MuJoCo 2D textures.
setup_empty_materials

Create a pool of empty materials and textures in the MjSpec that can be assigned to geoms at runtime.

Attributes:

Name Type Description
TEXTURE_FOLDER_PATH
TEXTURE_FOLDER_PATH module-attribute 
TEXTURE_FOLDER_PATH = ASSETS_DIR / 'objects' / 'thor' / 'Textures'
TextureRandomizer 
TextureRandomizer(model: MjModel, random_state: RandomState | None = None, geom_names: list[str] | None = None, randomize_geom_rgba: bool = True, randomize_material_rgba: bool = True, randomize_material_specular: bool = True, randomize_material_shininess: bool = True, randomize_texture: bool = True, texture_paths: list[str] | None = None, scene_metadata: dict | None = None, rgba_perturbation_size: float = 0.1, specular_perturbation_size: float = 0.1, shininess_perturbation_size: float = 0.1)

Randomizer for geom colors, material properties, and textures in MuJoCo simulations.

Can randomize: - Geom RGBA colors - Material properties (RGBA, specular, shininess) - Texture bitmaps from loaded texture files

Parameters:

Name Type Description Default
model MjModel

MuJoCo model

 required
random_state RandomState | None

Random state for reproducibility. If None, uses global numpy random state.

 None
geom_names list[str] | None

List of geom names to randomize. If None, randomizes all geoms in the model.

 None
randomize_geom_rgba bool

If True, randomizes geom RGBA colors

 True
randomize_material_rgba bool

If True, randomizes material RGBA colors

 True
randomize_material_specular bool

If True, randomizes material specular

 True
randomize_material_shininess bool

If True, randomizes material shininess

 True
randomize_texture bool

If True, randomizes texture bitmaps from loaded textures. Default behavior uses textures already loaded in the model XML.

 True
texture_paths list[str] | None

Optional list of paths to external texture image files (PNG, etc.). If None (default), uses textures already loaded in the model XML. If provided, loads textures from these external files instead.

 None
scene_metadata dict | None

Optional scene metadata for category-based texture randomization

 None
rgba_perturbation_size float

Magnitude of RGBA color randomization

 0.1
specular_perturbation_size float

Magnitude of specular randomization

 0.1
shininess_perturbation_size float

Magnitude of shininess randomization

 0.1
Note

MjData should be passed to the randomize() method, not to init.

Methods:

Name Description
randomize

Randomize all textures, colors, and material attributes for all geoms, regardless of category.
randomize_by_category

Randomize textures and colors by category.
randomize_object

Randomize colors and material attributes for a single MlSpacesObject.
save_defaults

Save default geom and material parameter values from the current model state.

Attributes:

Name Type Description
CAT_TO_TEXTURE
MAT_PER_CATEGORY
MAT_TO_TEXTURE
geom_names
mat_dict
material_database_filename
materials_to_texture_filename
model
random_state
randomize_geom_rgba
randomize_material_rgba
randomize_material_shininess
randomize_material_specular
randomize_texture
rgba_perturbation_size
scene_metadata
shininess_perturbation_size
specular_perturbation_size
texture_bitmaps list[ndarray]
texture_ids list[int]
texture_paths
Source code in molmo_spaces/env/arena/randomization/texture.py 
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
def __init__(
    self,
    model: MjModel,
    random_state: np.random.RandomState | None = None,
    geom_names: list[str] | None = None,
    randomize_geom_rgba: bool = True,
    randomize_material_rgba: bool = True,
    randomize_material_specular: bool = True,
    randomize_material_shininess: bool = True,
    randomize_texture: bool = True,
    texture_paths: list[str] | None = None,
    scene_metadata: dict | None = None,
    rgba_perturbation_size: float = 0.1,
    specular_perturbation_size: float = 0.1,
    shininess_perturbation_size: float = 0.1,
):
    self.model = model
    self.scene_metadata = scene_metadata
    self._empty_material_names: list[str] = []
    self._empty_material_ids: list[int] = []
    self._empty_texture_names: list[str] = []
    self._empty_texture_ids: list[int] = []
    self._empty_material_to_texture: dict[
        int, int
    ] = {}  # Map empty material ID -> empty texture ID
    self._next_empty_material_index = 0
    self._geom_to_empty_material: dict[int, int] = {}  # Map geom_id -> empty material ID
    self._geom_to_original_material: dict[
        int, int
    ] = {}  # Map geom_id -> original material ID (before assignment)
    self._geom_to_original_texture: dict[
        int, int
    ] = {}  # Map geom_id -> original texture ID (before assignment)

    # Find all empty materials and textures created for texture randomization
    self._find_empty_materials()

    if random_state is None:
        self.random_state = np.random
    else:
        self.random_state = random_state

    # Build mapping from body names to categories using scene_metadata
    self._body_name_to_category: dict[str, str] = {}
    if scene_metadata:
        objects = scene_metadata.get("objects", {})
        for obj_key, obj_data in objects.items():
            category = obj_data.get("category", "")
            name_map = obj_data.get("name_map", {})
            bodies = name_map.get("bodies", {})
            # Map both hash names and actual names to category
            for hash_name, actual_name in bodies.items():
                self._body_name_to_category[hash_name] = category
                self._body_name_to_category[actual_name] = category
            # Also map the object key itself
            self._body_name_to_category[obj_key] = category

    # Get geom names from model (only visual geoms: contype == 0)
    if geom_names is None:
        geom_names = []
        for i in range(model.ngeom):
            # Only include visual geoms (contype == 0 means no collision, visual only)
            if model.geom_contype[i] == 0 and model.geom_conaffinity[i] == 0:
                name_adr = model.name_geomadr[i]
                if name_adr >= 0:
                    name_bytes = model.names[name_adr:]
                    name = name_bytes.split(b"\x00")[0].decode("utf-8")
                    if name:
                        geom_names.append(name)

    self.geom_names = geom_names

    self.randomize_geom_rgba = randomize_geom_rgba
    self.randomize_material_rgba = randomize_material_rgba
    self.randomize_material_specular = randomize_material_specular
    self.randomize_material_shininess = randomize_material_shininess
    self.randomize_texture = randomize_texture

    self.rgba_perturbation_size = rgba_perturbation_size
    self.specular_perturbation_size = specular_perturbation_size
    self.shininess_perturbation_size = shininess_perturbation_size

    # Load texture files if provided, or extract existing textures from model (default)
    self.texture_bitmaps: list[np.ndarray] = []
    self.texture_ids: list[int] = []  # Texture IDs for on-demand extraction
    self.texture_paths = texture_paths  # Store for later use
    self._texture_cache: dict[
        int, np.ndarray
    ] = {}  # Cache for extracted textures (tex_id -> bitmap)
    self._texture_id_to_category: dict[
        int, str
    ] = {}  # Cache: texture_id -> category (for model textures)
    self._cat_texture_cache: dict[
        str, np.ndarray
    ] = {}  # Cache for category textures loaded from CAT_TO_TEXTURE
    if randomize_texture:
        if texture_paths:
            # Load textures from external files (user explicitly provided paths)
            self.texture_bitmaps = []
            for texture_path in texture_paths:
                bitmap = self._load_texture_from_path(texture_path, resolve_path=False)
                if bitmap is not None:
                    self.texture_bitmaps.append(bitmap)
        else:
            # Default: Extract existing textures from the model XML (on-demand)
            self._extract_model_textures()

    self.save_defaults()
CAT_TO_TEXTURE class-attribute instance-attribute 
CAT_TO_TEXTURE = {}
MAT_PER_CATEGORY class-attribute instance-attribute 
MAT_PER_CATEGORY = load(f)
MAT_TO_TEXTURE class-attribute instance-attribute 
MAT_TO_TEXTURE = load(f)
geom_names instance-attribute 
geom_names = geom_names
mat_dict class-attribute instance-attribute 
mat_dict = MAT_TO_TEXTURE[mat]
material_database_filename class-attribute instance-attribute 
material_database_filename = ASSETS_DIR / 'objects' / 'thor' / 'material-database.json'
materials_to_texture_filename class-attribute instance-attribute 
materials_to_texture_filename = ASSETS_DIR / 'objects' / 'thor' / 'material_to_textures.json'
model instance-attribute 
model = model
random_state instance-attribute 
random_state = random
randomize_geom_rgba instance-attribute 
randomize_geom_rgba = randomize_geom_rgba
randomize_material_rgba instance-attribute 
randomize_material_rgba = randomize_material_rgba
randomize_material_shininess instance-attribute 
randomize_material_shininess = randomize_material_shininess
randomize_material_specular instance-attribute 
randomize_material_specular = randomize_material_specular
randomize_texture instance-attribute 
randomize_texture = randomize_texture
rgba_perturbation_size instance-attribute 
rgba_perturbation_size = rgba_perturbation_size
scene_metadata instance-attribute 
scene_metadata = scene_metadata
shininess_perturbation_size instance-attribute 
shininess_perturbation_size = shininess_perturbation_size
specular_perturbation_size instance-attribute 
specular_perturbation_size = specular_perturbation_size
texture_bitmaps instance-attribute 
texture_bitmaps: list[ndarray] = []
texture_ids instance-attribute 
texture_ids: list[int] = []
texture_paths class-attribute instance-attribute 
texture_paths = texture_paths
randomize 
randomize(data: MjData | None = None) -> None

Randomize all textures, colors, and material attributes for all geoms, regardless of category. This method bypasses category filtering and applies full randomization to all geoms.

Parameters:

Name Type Description Default
data MjData | None

MuJoCo data for forward pass. If None, forward pass is skipped.

 None
Source code in molmo_spaces/env/arena/randomization/texture.py 
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
def randomize(self, data: MjData | None = None) -> None:
    """
    Randomize all textures, colors, and material attributes for all geoms, regardless of category.
    This method bypasses category filtering and applies full randomization to all geoms.

    Args:
        data (MjData | None): MuJoCo data for forward pass. If None, forward pass is skipped.
    """
    name_to_geom_id = self._build_name_to_geom_id()

    # Track randomization stats for debugging
    textures_randomized = 0
    colors_randomized = 0

    # Process all geoms
    for name in self.geom_names:
        geom_id = name_to_geom_id.get(name, -1)
        if geom_id < 0:
            continue

        mat_id = int(self.model.geom_matid[geom_id])

        # Randomize texture if available
        tex_id = self._get_texture_id_for_geom(geom_id)
        if tex_id >= 0:
            # Check if we have textures available (either from files or model)
            has_textures = (self.texture_paths and len(self.texture_bitmaps) > 0) or len(
                self.texture_ids
            ) > 0
            if has_textures and self.randomize_texture:
                # Randomize texture without category filtering
                self._randomize_texture_all(name, geom_id, mat_id)
                textures_randomized += 1

                # Always set RGB to white [1, 1, 1] but preserve original alpha when using textures
                # so texture renders at full intensity while maintaining transparency
                original_geom_alpha = self.model.geom_rgba[geom_id][3]
                self.model.geom_rgba[geom_id] = np.array([1.0, 1.0, 1.0, original_geom_alpha])
                if mat_id >= 0:
                    original_mat_alpha = self.model.mat_rgba[mat_id][3]
                    self.model.mat_rgba[mat_id] = np.array([1.0, 1.0, 1.0, original_mat_alpha])
            else:
                # No texture or texture randomization disabled, randomize colors and materials
                if self.randomize_geom_rgba:
                    self._randomize_geom_rgba_direct(geom_id)
                    colors_randomized += 1
                self._randomize_material_attributes(geom_id, mat_id)
        else:
            # No texture, randomize colors and materials
            if self.randomize_geom_rgba:
                self._randomize_geom_rgba_direct(geom_id)
                colors_randomized += 1
            self._randomize_material_attributes(geom_id, mat_id)

    # Forward pass to propagate changes
    # IMPORTANT: mj_forward must be called after texture changes for them to take effect
    if data is not None:
        mujoco.mj_forward(self.model, data)

    # Debug: Print randomization stats (only occasionally to avoid spam)
    if textures_randomized > 0 or colors_randomized > 0:
        import random

        if random.random() < 0.01:  # 1% chance to print
            print(
                f"   Debug: Randomized {textures_randomized} textures, "
                f"{colors_randomized} colors (all, no category filter)"
            )
randomize_by_category 
randomize_by_category(data: MjData | None = None)

Randomize textures and colors by category. For each geom, randomly picks a material from the appropriate category in MAT_PER_CATEGORY and applies its texture (if available) or color. Each geom gets a different random material.

Targets: floors, countertops, tabletops, doors (including door handles, drawers, cabinets)

Source code in molmo_spaces/env/arena/randomization/texture.py 
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
def randomize_by_category(self, data: MjData | None = None):
    """
    Randomize textures and colors by category.
    For each geom, randomly picks a material from the appropriate category in MAT_PER_CATEGORY
    and applies its texture (if available) or color. Each geom gets a different random material.

    Targets: floors, countertops, tabletops, doors (including door handles, drawers, cabinets)
    """
    name_to_geom_id = self._build_name_to_geom_id()

    # Track randomization stats for debugging
    textures_randomized = 0
    colors_randomized = 0

    # Group geoms by category for batch processing
    geoms_by_category: dict[str, list[tuple[str, int, int]]] = {}
    # Format: {category_key: [(name, geom_id, mat_id), ...]}

    # Process only geoms in our list (which are already filtered to visual geoms)
    for name in self.geom_names:
        geom_id = name_to_geom_id.get(name, -1)
        if geom_id < 0:
            continue

        # Get category from scene_metadata, fallback to name
        category = self._get_geom_category(geom_id)
        if not category:
            category = name

        # Check if this geom is in target categories
        if not self._is_target_category(category, name):
            continue  # Skip geoms not in target categories

        # Find matching category key from MAT_PER_CATEGORY
        category_lower = category.lower() if category else ""
        name_lower = name.lower()

        # Hard-coded category mapping for ProcTHOR and iTHOR
        # Maps scene category keywords to MAT_PER_CATEGORY keys
        type_map = {
            "room": "floor",
            "doorway": "doorway",
            "door": "doorway",  # Map "door" to "Doorway" in MAT_PER_CATEGORY
            "handle": "doorway",
            "drawer": "doorway",
            "cabinet": "doorway",
            "counter": "countertop",  # Map "counter" to "CounterTop" in MAT_PER_CATEGORY
            "island": "table",  # Map "island" to "CounterTop" in MAT_PER_CATEGORY
            "plane": "table",
            "mesh": "wall",
            "backsplash": "wall",
            "quad": "wall",
        }

        # First, try to map using type_map (check both category and name)
        # Check name first since it's more specific, then category
        mapped_category = None
        for type_key, type_value in type_map.items():
            # Check if keyword is in name (more reliable)
            if type_key.lower() in name_lower:
                mapped_category = type_value
                break
            # Also check category if available
            if category and type_key.lower() in category_lower:
                mapped_category = type_value
                break

        # If we found a mapping, use it; otherwise use the original category or name
        if mapped_category:
            search_category = mapped_category
        elif category:
            search_category = category_lower
        else:
            search_category = name_lower

        # Find matching category key from MAT_PER_CATEGORY
        matching_category_key = None
        search_category_lower = search_category.lower()

        # First try exact case-insensitive match
        for cat_key in self.MAT_PER_CATEGORY:
            if cat_key.lower() == search_category_lower:
                matching_category_key = cat_key
                break

        # If no exact match, try substring matching (bidirectional)
        if matching_category_key is None:
            for cat_key in self.MAT_PER_CATEGORY:
                cat_key_lower = cat_key.lower()
                # Check if search_category matches cat_key (bidirectional substring match)
                if (
                    cat_key_lower in search_category_lower
                    or search_category_lower in cat_key_lower
                ):
                    matching_category_key = cat_key
                    break

        if matching_category_key is None:
            continue  # No matching category in MAT_PER_CATEGORY

        mat_id = int(self.model.geom_matid[geom_id])

        # Group by category key
        if matching_category_key not in geoms_by_category:
            geoms_by_category[matching_category_key] = []
        geoms_by_category[matching_category_key].append((name, geom_id, mat_id))

    # For each category, randomly pick a different material for each geom
    for category_key, geoms in geoms_by_category.items():
        if not geoms:
            continue

        # Get materials for this category from MAT_PER_CATEGORY
        category_materials = self.MAT_PER_CATEGORY.get(category_key, [])
        if not category_materials:
            continue

        # Apply a randomly selected material to each geom (different for each)
        # IMPORTANT: Each geom gets its own independent random material selection
        # Only uses materials from MAT_PER_CATEGORY for this specific category
        selected_materials = []  # Track selected materials for debugging
        for name, geom_id, _ in geoms:
            # Assign an empty material to this geom first (ensures isolation from other geoms)
            empty_mat_id = self._assign_empty_material_to_geom(geom_id)
            # Get the new material ID after assignment
            new_mat_id = int(self.model.geom_matid[geom_id])

            # Verify the geom actually got the empty material assigned
            if new_mat_id != empty_mat_id:
                import logging

                log = logging.getLogger(__name__)
                log.warning(
                    f"Geom '{name}' (id={geom_id}) expected empty material {empty_mat_id} "
                    f"but got material {new_mat_id}"
                )

            # Randomly pick one material from this category's valid materials
            # Each geom gets its own random selection - this ensures variety
            # Only selects from materials in MAT_PER_CATEGORY for this category
            material_idx = self.random_state.randint(len(category_materials))
            selected_material = category_materials[material_idx]
            selected_materials.append(selected_material)

            texture_applied, color_applied = self._apply_material_to_geom(
                selected_material, geom_id, new_mat_id
            )
            if texture_applied:
                textures_randomized += 1
                # Verify texture is actually assigned to material
                actual_mat_id = int(self.model.geom_matid[geom_id])
                if actual_mat_id >= 0 and hasattr(self.model, "mat_texid"):
                    try:
                        if isinstance(self.model.mat_texid, np.ndarray):
                            if (
                                self.model.mat_texid.ndim == 2
                                and actual_mat_id < self.model.mat_texid.shape[0]
                            ):
                                assigned_tex_id = int(self.model.mat_texid[actual_mat_id, 0])
                                if assigned_tex_id >= 0:
                                    # Verify texture data is non-zero
                                    tex_adr = int(self.model.tex_adr[assigned_tex_id])
                                    tex_size = int(
                                        self.model.tex_height[assigned_tex_id]
                                        * self.model.tex_width[assigned_tex_id]
                                        * self.model.tex_nchannel[assigned_tex_id]
                                    )
                                    if tex_adr + tex_size <= len(self.model.tex_data):
                                        tex_data_sum = np.sum(
                                            np.abs(
                                                self.model.tex_data[
                                                    tex_adr : tex_adr + tex_size
                                                ]
                                            )
                                        )
                                        if tex_data_sum == 0:
                                            import logging

                                            log = logging.getLogger(__name__)
                                            log.warning(
                                                f"Geom {geom_id} has texture {assigned_tex_id} assigned to material {actual_mat_id}, "
                                                f"but texture data is all zeros!"
                                            )
                    except (IndexError, TypeError, ValueError):
                        pass
            elif color_applied:
                colors_randomized += 1
            # Note: If material has texture, color is set to white [1,1,1] to show texture
            # If material has no texture, color from albedo_rgba is applied

    # Forward pass to propagate changes
    if data is not None:
        mujoco.mj_forward(self.model, data)

    # Debug: Print randomization stats (only occasionally to avoid spam)
    if textures_randomized > 0 or colors_randomized > 0:
        import random

        if random.random() < 0.01:  # 1% chance to print
            print(
                f"   Debug: Randomized {textures_randomized} textures, "
                f"{colors_randomized} colors by category (material-based)"
            )
randomize_object 
randomize_object(thor_object: MlSpacesObject, data: MjData | None = None) -> None

Randomize colors and material attributes for a single MlSpacesObject. Only randomizes geoms that don't have textures.

Parameters:

Name Type Description Default
thor_object MlSpacesObject

MlSpacesObject instance to randomize

 required
data MjData | None

MuJoCo data for forward pass. If None, forward pass is skipped.

 None
Source code in molmo_spaces/env/arena/randomization/texture.py 
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
def randomize_object(self, thor_object: "MlSpacesObject", data: MjData | None = None) -> None:
    """
    Randomize colors and material attributes for a single MlSpacesObject.
    Only randomizes geoms that don't have textures.

    Args:
        thor_object: MlSpacesObject instance to randomize
        data: MuJoCo data for forward pass. If None, forward pass is skipped.
    """
    from molmo_spaces.env.data_views import MlSpacesObject

    if not isinstance(thor_object, MlSpacesObject):
        raise TypeError(f"Expected MlSpacesObject, got {type(thor_object)}")

    # Get all geoms for this object
    geom_infos = thor_object.get_geom_infos(include_descendants=True)
    if not geom_infos:
        return  # No geoms to randomize

    colors_randomized = 0

    # Process each geom
    for geom_info in geom_infos:
        geom_id = geom_info["id"]
        if geom_id < 0 or geom_id >= self.model.ngeom:
            continue

        # Check if geom has texture
        tex_id = self._get_texture_id_for_geom(geom_id)
        has_texture = tex_id >= 0

        # Only randomize if geom doesn't have texture
        if not has_texture:
            mat_id = int(self.model.geom_matid[geom_id])

            # Randomize colors and material attributes
            if self.randomize_geom_rgba:
                self._randomize_geom_rgba_direct(geom_id)
                colors_randomized += 1
            self._randomize_material_attributes(geom_id, mat_id)

    # Forward pass to propagate changes
    if data is not None:
        mujoco.mj_forward(self.model, data)

    # Debug: Print randomization stats (only occasionally to avoid spam)
    if colors_randomized > 0:
        import random

        if random.random() < 0.01:  # 1% chance to print
            print(
                f"   Debug: Randomized {colors_randomized} geoms (colors/material) "
                f"for object {thor_object.name}"
            )
save_defaults 
save_defaults()

Save default geom and material parameter values from the current model state. Optimized for large scenes by only processing visual geoms (contype == 0).

Source code in molmo_spaces/env/arena/randomization/texture.py 
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
def save_defaults(self):
    """
    Save default geom and material parameter values from the current model state.
    Optimized for large scenes by only processing visual geoms (contype == 0).
    """
    # Build name-to-geom_id mapping once, only for visual geoms (O(n) instead of O(n*m))
    name_to_geom_id = {}
    for i in range(self.model.ngeom):
        # Only process visual geoms (contype == 0 means no collision, visual only)
        if self.model.geom_contype[i] == 0:
            name_adr = self.model.name_geomadr[i]
            if name_adr >= 0:
                name_bytes = self.model.names[name_adr:]
                name = name_bytes.split(b"\x00")[0].decode("utf-8")
                if name:
                    name_to_geom_id[name] = i

    self._defaults = {}
    self._geom_id_to_defaults: dict[int, dict] = {}  # Fast lookup by geom_id
    # Only process geoms that are in our list
    for name in self.geom_names:
        geom_id = name_to_geom_id.get(name, -1)
        if geom_id < 0:
            continue

        defaults = {}
        defaults["geom_rgba"] = np.array(self.model.geom_rgba[geom_id])

        # Check if geom has a material with texture
        mat_id = int(self.model.geom_matid[geom_id])
        if mat_id >= 0:
            defaults["mat_rgba"] = np.array(self.model.mat_rgba[mat_id])
            defaults["mat_specular"] = float(self.model.mat_specular[mat_id])
            defaults["mat_shininess"] = float(self.model.mat_shininess[mat_id])
            defaults["mat_id"] = mat_id

            # Save texture ID instead of bitmap (much faster, less memory)
            if self.randomize_texture:
                tex_id = self._get_texture_id_for_geom(geom_id)
                defaults["texture_id"] = tex_id  # Store ID, not bitmap
            else:
                defaults["texture_id"] = -1
        else:
            defaults["mat_rgba"] = None
            defaults["mat_specular"] = None
            defaults["mat_shininess"] = None
            defaults["mat_id"] = -1
            defaults["texture_id"] = -1

        self._defaults[name] = defaults
        self._geom_id_to_defaults[geom_id] = defaults  # Fast lookup by geom_id
assign_texture_to_material 
assign_texture_to_material(material, tex_name: str, mat_name: str, log) -> None

Assign a texture to a material's RGB role.

Source code in molmo_spaces/env/arena/randomization/texture.py 
51
52
53
54
55
56
57
58
def assign_texture_to_material(material, tex_name: str, mat_name: str, log) -> None:
    """Assign a texture to a material's RGB role."""
    try:
        rgb_role = mujoco.mjtTextureRole.mjTEXROLE_RGB.value
        material.textures[rgb_role] = tex_name
        log.debug(f"Assigned texture {tex_name} to material {mat_name}")
    except Exception as e:
        log.warning(f"Failed to assign texture {tex_name} to material {mat_name}: {e}")
create_empty_texture 
create_empty_texture(spec: MjSpec, tex_name: str, placeholder_file: str, log) -> bool

Create an empty texture in the spec using the placeholder file.

Source code in molmo_spaces/env/arena/randomization/texture.py 
41
42
43
44
45
46
47
48
def create_empty_texture(spec: mujoco.MjSpec, tex_name: str, placeholder_file: str, log) -> bool:
    """Create an empty texture in the spec using the placeholder file."""
    try:
        spec.add_texture(name=tex_name, type=mujoco.mjtTexture.mjTEXTURE_2D, file=placeholder_file)
        return True
    except Exception as e:
        log.warning(f"Failed to create texture {tex_name}: {e}")
        return False
create_placeholder_texture_file 
create_placeholder_texture_file(texture_size: int, log) -> str | None

Create a temporary placeholder texture file for MuJoCo 2D textures.

Source code in molmo_spaces/env/arena/randomization/texture.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def create_placeholder_texture_file(texture_size: int, log) -> str | None:
    """Create a temporary placeholder texture file for MuJoCo 2D textures."""
    import os
    import tempfile

    from PIL import Image

    try:
        temp_dir = tempfile.mkdtemp(prefix="mujoco_texture_placeholder_")
        placeholder_path = os.path.join(temp_dir, "__TEXTURE_RANDOMIZER_PLACEHOLDER__.png")

        img = Image.new("RGB", (texture_size, texture_size), color=(255, 255, 255))
        img.save(placeholder_path, "PNG")

        if os.path.exists(placeholder_path) and os.path.getsize(placeholder_path) > 0:
            log.debug(f"Created placeholder texture file: {placeholder_path}")
            return placeholder_path
        else:
            log.error(f"Placeholder texture file was not created or is empty: {placeholder_path}")
            return None
    except Exception as e:
        log.error(f"Failed to create placeholder texture file: {e}")
        return None
setup_empty_materials 
setup_empty_materials(spec: MjSpec | None = None, num_materials: int = 200) -> None

Create a pool of empty materials and textures in the MjSpec that can be assigned to geoms at runtime. This allows texture randomization to modify materials and textures without affecting other geoms.

Parameters:

Name Type Description Default
spec MjSpec | None

MjSpec to modify

 None
num_materials int

Maximum number of empty materials/textures to create. Actual number is based on visual geom count with a safety buffer.

 200
Source code in molmo_spaces/env/arena/randomization/texture.py 
 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
def setup_empty_materials(spec: mujoco.MjSpec | None = None, num_materials: int = 200) -> None:
    """
    Create a pool of empty materials and textures in the MjSpec that can be assigned to geoms at runtime.
    This allows texture randomization to modify materials and textures without affecting other geoms.

    Args:
        spec: MjSpec to modify
        num_materials: Maximum number of empty materials/textures to create. Actual number is based on
                      visual geom count with a safety buffer.
    """
    import logging

    log = logging.getLogger(__name__)

    if spec is None:
        raise ValueError("spec cannot be None")

    # Check if empty materials already exist to avoid duplicates
    if spec.material("__TEXTURE_RANDOMIZER_MAT_0__") is not None:
        log.debug("Empty materials already exist in spec, skipping creation")
        return

    # Count visual geoms to estimate how many materials we might need
    visual_geom_count = sum(
        1
        for geom in spec.geoms
        if (
            hasattr(geom, "contype")
            and hasattr(geom, "conaffinity")
            and geom.contype == 0
            and geom.conaffinity == 0
        )
        or (
            hasattr(geom, "classname")
            and ("__VISUAL_MJT__" in str(geom.classname) or "visual" in str(geom.classname).lower())
        )
    )

    # Calculate number to create: use visual_geom_count with safety buffer, cap at num_materials
    safety_multiplier = 1.0
    num_to_create = min(int(visual_geom_count * safety_multiplier), num_materials)
    num_to_create = max(num_to_create, min(50, num_materials))

    # Use 512x512 as default texture size (can be overridden by caller if needed)
    texture_size = 512
    memory_per_texture_mb = texture_size * texture_size * 3 / (1024**2)
    total_memory_mb = num_to_create * memory_per_texture_mb

    log.info(
        f"Creating {num_to_create} empty materials and textures "
        f"(visual geoms: {visual_geom_count}, texture size: {texture_size}x{texture_size}, "
        f"estimated memory: ~{total_memory_mb:.1f} MB per model)"
    )

    # Create placeholder texture file (MuJoCo requires a file path for 2D textures)
    placeholder_file = create_placeholder_texture_file(texture_size, log)
    if placeholder_file is None:
        log.warning(
            "Failed to create placeholder texture file. Materials will be created without textures."
        )

    # Create empty materials and textures
    created_count = 0
    for i in range(num_to_create):
        mat_name = f"__TEXTURE_RANDOMIZER_MAT_{i}__"
        tex_name = f"__TEXTURE_RANDOMIZER_TEX_{i}__"

        # Create texture if placeholder file is available
        texture_exists = False
        if placeholder_file is not None and spec.texture(tex_name) is None:
            texture_exists = create_empty_texture(spec, tex_name, placeholder_file, log)

        # Create material
        empty_mat = spec.add_material(name=mat_name)
        empty_mat.rgba = np.array([1.0, 1.0, 1.0, 1.0], dtype=np.float64)
        empty_mat.specular = 0.5
        empty_mat.shininess = 0.5

        # Assign texture to material if it was created
        if texture_exists or (placeholder_file is not None and spec.texture(tex_name) is not None):
            assign_texture_to_material(empty_mat, tex_name, mat_name, log)
        else:
            log.debug(f"Material {mat_name} created without texture")

        created_count += 1

    log.info(f"Created {created_count} empty materials and textures for texture randomization")

scene_tweaks

Classes:

Name Description
Context
ContextFlags

Functions:

Name Description
does_body_aabb_intersect_box_site
is_body_com_within_box_site
is_body_within_any_site
is_body_within_site_in_freespace
key_callback
run_check_body_com_all_sites

Attributes:

Name Type Description
DISTANCE_RAYCAST_THRESHOLD
OFFSET_UP_RAYCAST
SITE_SIZE_TOLERANCE
args
body_handle
body_name
context
data
house_path
model
parser
result
t_start
DISTANCE_RAYCAST_THRESHOLD module-attribute 
DISTANCE_RAYCAST_THRESHOLD = 0.5
OFFSET_UP_RAYCAST module-attribute 
OFFSET_UP_RAYCAST = 0.05
SITE_SIZE_TOLERANCE module-attribute 
SITE_SIZE_TOLERANCE = array([0.05, 0.05, 0.05], dtype=float64)
args module-attribute 
args = parse_args()
body_handle module-attribute 
body_handle = body(body)
body_name module-attribute 
body_name = name
context module-attribute 
context = Context()
data module-attribute 
data = MjData(model)
house_path module-attribute 
house_path = Path(house)
model module-attribute 
model = from_xml_path(as_posix())
parser module-attribute 
parser = ArgumentParser()
result module-attribute 
result = is_body_com_within_box_site(site_id, body_id, model, data)
t_start module-attribute 
t_start = time
Context 
Context()

Attributes:

Name Type Description
body_id
flags
site_id
Source code in molmo_spaces/env/arena/scene_tweaks.py 
141
142
143
144
def __init__(self) -> None:
    self.body_id = -1
    self.site_id = -1
    self.flags = ContextFlags()
body_id instance-attribute 
body_id = -1
flags instance-attribute 
flags = ContextFlags()
site_id instance-attribute 
site_id = -1
ContextFlags dataclass 
ContextFlags()

Attributes:

Name Type Description
dirty_pause
dirty_reset
dirty_test_aabb_intersect
dirty_test_com_inside
dirty_test_in_free_space
dirty_test_with_all_sites
dirty_pause class-attribute instance-attribute 
dirty_pause = False
dirty_reset class-attribute instance-attribute 
dirty_reset = False
dirty_test_aabb_intersect class-attribute instance-attribute 
dirty_test_aabb_intersect = False
dirty_test_com_inside class-attribute instance-attribute 
dirty_test_com_inside = False
dirty_test_in_free_space class-attribute instance-attribute 
dirty_test_in_free_space = False
dirty_test_with_all_sites class-attribute instance-attribute 
dirty_test_with_all_sites = False
does_body_aabb_intersect_box_site 
does_body_aabb_intersect_box_site(site_id: int, body_id: int, model: MjModel, data: MjData) -> bool
Source code in molmo_spaces/env/arena/scene_tweaks.py 
55
56
57
58
59
def does_body_aabb_intersect_box_site(
    site_id: int, body_id: int, model: mj.MjModel, data: mj.MjData
) -> bool:
    # TODO(wilbert): implement this helper function, as it might work better for the tests
    pass
is_body_com_within_box_site 
is_body_com_within_box_site(site_id: int, body_id: int, model: MjModel, data: MjData) -> bool
Source code in molmo_spaces/env/arena/scene_tweaks.py 
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
def is_body_com_within_box_site(
    site_id: int, body_id: int, model: mj.MjModel, data: mj.MjData
) -> bool:
    site_type = model.site_type[site_id].item()
    if site_type != mj.mjtGeom.mjGEOM_BOX:
        return False

    # Make sure we're not considering site-body comparison of sites that belong to a body already
    site_bodyid = model.site_bodyid[site_id].item()
    site_rootid = model.body_rootid[site_bodyid].item()
    tgt_rootid = model.body_rootid[body_id].item()
    if site_rootid in (tgt_rootid, body_id):
        return False

    # Transform the body's position into the site's local frame, and do an AABB check in there
    body_com_pos = data.xpos[body_id]
    site_xpos = data.site_xpos[site_id]
    site_xmat = data.site_xmat[site_id].reshape(3, 3)
    site_size = model.site_size[site_id]

    plocal = body_com_pos - site_xpos
    plocal = site_xmat.T @ plocal

    # Converted is in the local frame of the site, so just check if it's within its AABB
    result = bool(np.all(np.abs(plocal) < (site_size + SITE_SIZE_TOLERANCE)))

    return result
is_body_within_any_site 
is_body_within_any_site(model: MjModel, data: MjData, body_id: int) -> tuple[bool, int]
Source code in molmo_spaces/env/arena/scene_tweaks.py 
43
44
45
46
47
48
49
50
51
52
def is_body_within_any_site(model: mj.MjModel, data: mj.MjData, body_id: int) -> tuple[bool, int]:
    is_within_site = False
    site_id_within = -1
    for site_id in range(model.nsite):
        if is_body_com_within_box_site(site_id, body_id, model, data):
            is_within_site = True
            site_id_within = site_id
            break

    return is_within_site, site_id_within
is_body_within_site_in_freespace 
is_body_within_site_in_freespace(site_id: int, body_id: int, model: MjModel, data: MjData) -> tuple[bool, float, str]
Source code in molmo_spaces/env/arena/scene_tweaks.py 
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
def is_body_within_site_in_freespace(
    site_id: int, body_id: int, model: mj.MjModel, data: mj.MjData
) -> tuple[bool, float, str]:
    site_size = model.site_size[site_id]
    geomid = np.zeros(1, dtype=np.int32)
    world_up = np.array([0, 0, 1], dtype=np.float64)
    pnt = data.xpos[body_id] + OFFSET_UP_RAYCAST * world_up

    distance_up = mj.mj_ray(model, data, pnt, world_up, None, 1, body_id, geomid)

    if distance_up == -1.0:  # No intersection, so there's space above
        return True, -1.0, ""

    is_there_space_above = distance_up > min(DISTANCE_RAYCAST_THRESHOLD, 2.0 * site_size[2])

    geom_name = ""
    if not is_there_space_above:  # just in case, try excluding the first body we touched
        geom_bodyid = model.geom_bodyid[geomid.item()].item()
        geom_rootid = model.body_rootid[geom_bodyid].item()
        if geom_rootid == body_id:
            distance_up = mj.mj_ray(model, data, pnt, world_up, None, 1, geom_bodyid, geomid)
        geom_name = mj.mj_id2name(model, mj.mjtObj.mjOBJ_GEOM, geomid.item())

    if distance_up == -1.0:  # No intersection, so there's space above
        return True, -1.0, ""

    is_there_space_above = distance_up > min(DISTANCE_RAYCAST_THRESHOLD, 2.0 * site_size[2])

    return is_there_space_above, distance_up, geom_name
key_callback 
key_callback(keycode: int) -> None
Source code in molmo_spaces/env/arena/scene_tweaks.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def key_callback(keycode: int) -> None:
    global model, data, context

    if keycode == 265:  # up arrow key
        context.body_id = (context.body_id + 1) % model.nbody
        print(f"Using body: {model.body(context.body_id).name}")
    elif keycode == 264:  # down arrow key
        context.body_id = (context.body_id - 1) % model.nbody
        print(f"Using body: {model.body(context.body_id).name}")
    elif keycode == 262:  # right arrow key
        context.site_id = (context.site_id + 1) % model.nsite
        print(f"Using site: {model.site(context.site_id).name}")
    elif keycode == 263:  # left arrow key
        context.site_id = (context.site_id - 1) % model.nsite
        print(f"Using site: {model.site(context.site_id).name}")
    elif keycode == 259:  # backspace key
        context.flags.dirty_reset = True
    elif keycode == 32:  # space key
        context.flags.dirty_pause = not context.flags.dirty_pause
        print("Paused" if context.flags.dirty_pause else "Running")
    elif keycode == 81:  # Q key
        context.flags.dirty_test_com_inside = True
    elif keycode == 80:  # P key
        context.flags.dirty_test_aabb_intersect = True
    elif keycode == 79:  # O key
        context.flags.dirty_test_with_all_sites = not context.flags.dirty_test_with_all_sites
        msg = "all sites" if context.flags.dirty_test_with_all_sites else "single site"
        print(f"Testing with {msg}")
run_check_body_com_all_sites 
run_check_body_com_all_sites(body_id: int) -> tuple[bool, str]
Source code in molmo_spaces/env/arena/scene_tweaks.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
def run_check_body_com_all_sites(body_id: int) -> tuple[bool, str]:
    found_within_site = False
    found_site_name = ""
    for site_id in range(model.nsite):
        if is_body_com_within_box_site(site_id, body_id, model, data):
            in_free_space, _, _ = is_body_within_site_in_freespace(
                site_id, body_id, model, data
            )
            if not in_free_space:
                found_within_site = True
                found_site_name = model.site(site_id).name
                break
    return found_within_site, found_site_name

camera_manager

Camera management for MolmoSpaces environments.

This module handles all camera-related functionality including camera registration, pose updates, rendering, and setup from camera configurations.

Classes:

Name Description
Camera

Base camera class with position and orientation.
CameraManager

Manages all camera-related operations for an environment.
CameraRegistry

Registry for camera objects with auto-updating support.
RobotMountedCamera

Generic robot-mounted camera that can attach to any body/joint with configurable offsets.

Attributes:

Name Type Description
log

log module-attribute 

log = getLogger(__name__)

Camera 

Camera(name: str, pos: NDArray[float32] | None = None, forward: NDArray[float32] | None = None, up: NDArray[float32] | None = None, fov: float = 45.0)

Base camera class with position and orientation.

Methods:

Name Description
get_pose

return 4x4 pose
update_pose

Update camera pose. Returns True if pose changed, False otherwise.

Attributes:

Name Type Description
forward NDArray[float32]
fov float
name str
pos NDArray[float32]
up NDArray[float32]
Source code in molmo_spaces/env/camera_manager.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def __init__(
    self,
    name: str,
    pos: NDArray[np.float32] | None = None,
    forward: NDArray[np.float32] | None = None,
    up: NDArray[np.float32] | None = None,
    fov: float = 45.0,
) -> None:
    self.name: str = name
    self.pos: NDArray[np.float32] = (
        pos if pos is not None else np.array([0.0, 0.0, 1.0], dtype=np.float32)
    )
    self.forward: NDArray[np.float32] = (
        forward if forward is not None else np.array([1.0, 0.0, 0.0], dtype=np.float32)
    )
    self.up: NDArray[np.float32] = (
        up if up is not None else np.array([0.0, 0.0, 1.0], dtype=np.float32)
    )
    self.fov: float = fov
forward instance-attribute 
forward: NDArray[float32] = forward if forward is not None else array([1.0, 0.0, 0.0], dtype=float32)
fov instance-attribute 
fov: float = fov
name instance-attribute 
name: str = name
pos instance-attribute 
pos: NDArray[float32] = pos if pos is not None else array([0.0, 0.0, 1.0], dtype=float32)
up instance-attribute 
up: NDArray[float32] = up if up is not None else array([0.0, 0.0, 1.0], dtype=float32)
get_pose 
get_pose() -> NDArray[float32]

return 4x4 pose

Source code in molmo_spaces/env/camera_manager.py 
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
def get_pose(self) -> NDArray[np.float32]:
    """
    return 4x4 pose
    """
    # Validate and normalize camera vectors
    forward_norm = np.linalg.norm(self.forward)
    up_norm = np.linalg.norm(self.up)

    if forward_norm < 1e-6 or up_norm < 1e-6:
        print(
            f"Warning: Camera '{self.camera_name}' has degenerate vectors (forward_norm={forward_norm}, up_norm={up_norm})"
        )
        return np.eye(4, 4, dtype=np.float32)

    forward = self.forward / forward_norm
    up = self.up / up_norm
    right = np.cross(forward, up)

    right_norm = np.linalg.norm(right)
    if right_norm < 1e-6:
        print(f"Warning: Camera '{self.self}' has collinear forward/up vectors")
        return np.eye(4, 4, dtype=np.float32)

    right = right / right_norm

    # Recompute orthogonal up to ensure proper orthogonal basis
    up = np.cross(right, forward)

    # Create cam2world matrix (standard camera convention)
    world2cam = np.eye(4)
    world2cam[:3, 0] = right  # X-axis (right)
    world2cam[:3, 1] = -up  # Y-axis (up)
    world2cam[:3, 2] = forward  # Z-axis - camera looks down negative Z
    world2cam[:3, 3] = self.pos  # Translation
    return world2cam
update_pose 
update_pose(env: CPUMujocoEnv) -> bool

Update camera pose. Returns True if pose changed, False otherwise.

Source code in molmo_spaces/env/camera_manager.py 
54
55
56
def update_pose(self, env: CPUMujocoEnv) -> bool:
    """Update camera pose. Returns True if pose changed, False otherwise."""
    return False  # by default cameras don't update

CameraManager 

CameraManager()

Manages all camera-related operations for an environment.

This class encapsulates camera setup, registration, and rendering operations, keeping camera logic separate from core environment concerns.

Note: This class does not store a reference to the environment to avoid circular references and enable pickling for multiprocessing. Instead, the environment is passed as a parameter to methods that need it.

Initialize the camera manager with an empty registry.

Methods:

Name Description
add_camera

Adds or updates a static camera in the registry with its world pose.
add_robot_mounted_camera

Add a robot-mounted camera that follows a specific body/joint.
add_robot_mounted_camera_with_quaternion

Add a robot-mounted camera that follows a specific body/joint using quaternion orientation.
apply_mjcf_camera_noise

Apply position, orientation, and FOV noise to MJCF camera parameters.
create_lookat_pose

Create camera position and orientation vectors using enhanced lookat approach.
create_lookat_pose_world

Create camera position and orientation vectors using enhanced lookat approach.
create_quaternion_camera_pose

Create camera pose using quaternion-based orientation relative to reference body.
create_robot_mounted_camera_pose

Create generic robot-mounted camera pose using the existing create_lookat_pose method.
setup_cameras

Set up all cameras from a CameraSystemConfig.

Attributes:

Name Type Description
registry CameraRegistry
Source code in molmo_spaces/env/camera_manager.py 
266
267
268
def __init__(self) -> None:
    """Initialize the camera manager with an empty registry."""
    self.registry: CameraRegistry = CameraRegistry()
registry instance-attribute 
registry: CameraRegistry = CameraRegistry()
add_camera 
add_camera(camera_name: str, pos: NDArray[float32], forward: NDArray[float32], up: NDArray[float32], fov: float = 45.0) -> None

Adds or updates a static camera in the registry with its world pose.

Source code in molmo_spaces/env/camera_manager.py 
719
720
721
722
723
724
725
726
727
728
def add_camera(
    self,
    camera_name: str,
    pos: NDArray[np.float32],
    forward: NDArray[np.float32],
    up: NDArray[np.float32],
    fov: float = 45.0,
) -> None:
    """Adds or updates a static camera in the registry with its world pose."""
    self.registry.add_camera(Camera(camera_name, pos, forward, up, fov))
add_robot_mounted_camera 
add_robot_mounted_camera(env, camera_name: str, reference_body_names: str | list[str], camera_offset: NDArray[float32] | None = None, lookat_offset: NDArray[float32] | None = None, up_axis: str = 'z') -> None

Add a robot-mounted camera that follows a specific body/joint.

Parameters:

Name Type Description Default
env

The environment instance (CPUMujocoEnv)

 required
camera_name str

Name for the camera

 required
reference_body_names str | list[str]

Body name(s) to attach camera to. If list, tries each until one works.

 required
camera_offset NDArray[float32] | None

Camera position relative to reference body frame

 None
lookat_offset NDArray[float32] | None

Offset from reference body to look at

 None
up_axis str

Which local axis of reference frame is "up" ("x", "y", or "z")

 'z'
Source code in molmo_spaces/env/camera_manager.py 
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
def add_robot_mounted_camera(
    self,
    env,
    camera_name: str,
    reference_body_names: str | list[str],
    camera_offset: NDArray[np.float32] | None = None,
    lookat_offset: NDArray[np.float32] | None = None,
    up_axis: str = "z",
) -> None:
    """
    Add a robot-mounted camera that follows a specific body/joint.

    Args:
        env: The environment instance (CPUMujocoEnv)
        camera_name: Name for the camera
        reference_body_names: Body name(s) to attach camera to. If list, tries each until one works.
        camera_offset: Camera position relative to reference body frame
        lookat_offset: Offset from reference body to look at
        up_axis: Which local axis of reference frame is "up" ("x", "y", or "z")
    """
    robot_camera = RobotMountedCamera(
        camera_name, reference_body_names, camera_offset, lookat_offset, up_axis
    )
    # Initialize pose
    robot_camera.update_pose(env)
    self.registry.add_camera(robot_camera)
add_robot_mounted_camera_with_quaternion 
add_robot_mounted_camera_with_quaternion(env, camera_name: str, reference_body_names: str | list[str], camera_offset: NDArray[float32] | None = None, camera_quaternion: NDArray[float32] | None = None, camera_fov: float = 45) -> None

Add a robot-mounted camera that follows a specific body/joint using quaternion orientation.

Parameters:

Name Type Description Default
env

The environment instance (CPUMujocoEnv)

 required
camera_name str

Name for the camera

 required
reference_body_names str | list[str]

Body name(s) to attach camera to. If list, tries each until one works.

 required
camera_offset NDArray[float32] | None

Camera position relative to reference body frame

 None
camera_quaternion NDArray[float32] | None

Quaternion [w, x, y, z] relative to reference body frame

 None
Source code in molmo_spaces/env/camera_manager.py 
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
def add_robot_mounted_camera_with_quaternion(
    self,
    env,
    camera_name: str,
    reference_body_names: str | list[str],
    camera_offset: NDArray[np.float32] | None = None,
    camera_quaternion: NDArray[np.float32] | None = None,
    camera_fov: float = 45,
) -> None:
    """
    Add a robot-mounted camera that follows a specific body/joint using quaternion orientation.

    Args:
        env: The environment instance (CPUMujocoEnv)
        camera_name: Name for the camera
        reference_body_names: Body name(s) to attach camera to. If list, tries each until one works.
        camera_offset: Camera position relative to reference body frame
        camera_quaternion: Quaternion [w, x, y, z] relative to reference body frame
    """
    robot_camera = RobotMountedCamera(
        camera_name,
        reference_body_names,
        camera_offset=camera_offset,
        camera_quaternion=camera_quaternion,
        camera_fov=camera_fov,
    )
    # Initialize pose
    robot_camera.update_pose(env)
    self.registry.add_camera(robot_camera)
apply_mjcf_camera_noise staticmethod 
apply_mjcf_camera_noise(camera_pos: ndarray, camera_quat: ndarray, camera_fov: float, camera_config: MjcfCameraConfig, rng=random) -> tuple[ndarray, ndarray, float]

Apply position, orientation, and FOV noise to MJCF camera parameters.

Parameters:

Name Type Description Default
camera_pos ndarray

Camera position (body-frame offset), modified in place.

 required
camera_quat ndarray

Camera quaternion [w,x,y,z] (body-frame), modified in place.

 required
camera_fov float

Base FOV in degrees.

 required
camera_config MjcfCameraConfig

Config carrying noise ranges.

 required
rng

Random generator (np.random module or np.random.RandomState).

 random

Returns:

Type Description
tuple[ndarray, ndarray, float]

Tuple of (noised_pos, noised_quat, noised_fov).
Source code in molmo_spaces/env/camera_manager.py 
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
@staticmethod
def apply_mjcf_camera_noise(
    camera_pos: np.ndarray,
    camera_quat: np.ndarray,
    camera_fov: float,
    camera_config: MjcfCameraConfig,
    rng=np.random,
) -> tuple[np.ndarray, np.ndarray, float]:
    """Apply position, orientation, and FOV noise to MJCF camera parameters.

    Args:
        camera_pos: Camera position (body-frame offset), modified in place.
        camera_quat: Camera quaternion [w,x,y,z] (body-frame), modified in place.
        camera_fov: Base FOV in degrees.
        camera_config: Config carrying noise ranges.
        rng: Random generator (np.random module or np.random.RandomState).

    Returns:
        Tuple of (noised_pos, noised_quat, noised_fov).
    """
    if camera_config.fov_noise_degrees is not None:
        noise = rng.uniform(
            camera_config.fov_noise_degrees[0], camera_config.fov_noise_degrees[1]
        )
        camera_fov += noise
        log.debug(
            f"[CAMERA SETUP] Applied FOV noise to '{camera_config.name}': {noise} degrees"
        )

    if camera_config.pos_noise_range is not None:
        noise = rng.uniform(
            camera_config.pos_noise_range[0], camera_config.pos_noise_range[1], size=3
        )
        original_rot = R.from_quat(camera_quat, scalar_first=True)
        camera_pos = camera_pos + original_rot.apply(noise)
        log.debug(f"[CAMERA SETUP] Applied position noise to '{camera_config.name}': {noise}")

    if camera_config.orientation_noise_degrees is not None:
        noise_euler = rng.uniform(
            -np.array(camera_config.orientation_noise_degrees),
            np.array(camera_config.orientation_noise_degrees),
            size=3,
        )
        noise_rotation = R.from_euler("xyz", noise_euler, degrees=True)
        original_rotation = R.from_quat(camera_quat, scalar_first=True)
        noisy_rotation = original_rotation * noise_rotation
        camera_quat = noisy_rotation.as_quat(scalar_first=True)
        log.debug(
            f"[CAMERA SETUP] Applied orientation noise to '{camera_config.name}': {noise_euler} degrees"
        )

    return camera_pos, camera_quat, camera_fov
create_lookat_pose 
create_lookat_pose(env, camera_relative_pos: ndarray, rpy: ndarray, reference_body_name: str, lookat_target: ndarray = None, lookat_body_name: str = None, camera_up: ndarray = None) -> tuple[ndarray, ndarray, ndarray]

Create camera position and orientation vectors using enhanced lookat approach.

Parameters:

Name Type Description Default
env

The environment instance (CPUMujocoEnv)

 required
camera_relative_pos ndarray

Camera position relative to reference body

 required
rpy ndarray

Roll, pitch, yaw (currently unused but kept for compatibility)

 required
reference_body_name str

Body name for camera positioning reference

 required
lookat_target ndarray

Optional 3D point (np.array) to look at in world coordinates

 None
lookat_body_name str

Optional body name to look at (used if lookat_target is None)

 None
camera_up ndarray

Optional desired up direction for camera (np.array in world coordinates) If None, uses world Z-up [0, 0, 1]

 None
Source code in molmo_spaces/env/camera_manager.py 
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
def create_lookat_pose(
    self,
    env,
    camera_relative_pos: np.ndarray,
    rpy: np.ndarray,
    reference_body_name: str,
    lookat_target: np.ndarray = None,
    lookat_body_name: str = None,
    camera_up: np.ndarray = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Create camera position and orientation vectors using enhanced lookat approach.

    Args:
        env: The environment instance (CPUMujocoEnv)
        camera_relative_pos: Camera position relative to reference body
        rpy: Roll, pitch, yaw (currently unused but kept for compatibility)
        reference_body_name: Body name for camera positioning reference
        lookat_target: Optional 3D point (np.array) to look at in world coordinates
        lookat_body_name: Optional body name to look at (used if lookat_target is None)
        camera_up: Optional desired up direction for camera (np.array in world coordinates)
                  If None, uses world Z-up [0, 0, 1]
    """
    from molmo_spaces.env.data_views import create_mlspaces_body

    data = env.current_data

    coordinate_reference_body = create_mlspaces_body(data, reference_body_name)
    camera_pos_world = (
        coordinate_reference_body.pose[:3, :3] @ camera_relative_pos
        + coordinate_reference_body.pose[:3, 3]
    )

    # Determine what to look at
    if lookat_target is not None:
        # Look at the provided 3D point
        lookat_pos_world = lookat_target
    elif lookat_body_name is not None:
        # Look at the specified body
        lookat_reference_body = create_mlspaces_body(data, lookat_body_name)
        lookat_pos_world = lookat_reference_body.position
    else:
        raise ValueError("Either lookat_target or lookat_body_name must be provided")

    # Calculate forward direction (from camera to target)
    forward = lookat_pos_world - camera_pos_world
    forward = forward / np.linalg.norm(forward)

    # Calculate right and up vectors using the desired up direction
    if camera_up is None:
        camera_up = np.array([0.0, 0.0, 1.0])  # Default world Z-up

    # Calculate right vector (perpendicular to both forward and desired up)
    right = np.cross(forward, camera_up)
    right_norm = np.linalg.norm(right)

    # Handle case where forward is parallel to desired up
    if right_norm < 1e-6:
        # Use a fallback reference direction
        fallback_ref = np.array([1.0, 0.0, 0.0])
        if np.abs(np.dot(forward, fallback_ref)) > 0.9:
            fallback_ref = np.array([0.0, 1.0, 0.0])
        right = np.cross(forward, fallback_ref)
        right = right / np.linalg.norm(right)
    else:
        right = right / right_norm

    # Calculate actual up vector (perpendicular to forward and right)
    up = np.cross(right, forward)

    return camera_pos_world, forward, up
create_lookat_pose_world 
create_lookat_pose_world(env, camera_pos_world: ndarray, rpy: ndarray, lookat_target: ndarray | None = None, lookat_body_name: str | None = None, camera_up: ndarray | None = None) -> tuple[ndarray, ndarray, ndarray]

Create camera position and orientation vectors using enhanced lookat approach. Functionally equivalent to create_lookat_pose, but without use of a reference body.

Parameters:

Name Type Description Default
env

The environment instance (CPUMujocoEnv)

 required
camera_pos_world ndarray

Camera position in world coordinates

 required
rpy ndarray

Roll, pitch, yaw (currently unused but kept for compatibility)

 required
lookat_target ndarray | None

Optional 3D point (np.array) to look at in world coordinates

 None
lookat_body_name str | None

Optional body name to look at (used if lookat_target is None)

 None
camera_up ndarray | None

Optional desired up direction for camera (np.array in world coordinates) If None, uses world Z-up [0, 0, 1]

 None

Returns: Tuple of (camera_pos_world, forward_vector, up_vector), each of shape (3,)

Source code in molmo_spaces/env/camera_manager.py 
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
def create_lookat_pose_world(
    self,
    env,
    camera_pos_world: np.ndarray,
    rpy: np.ndarray,
    lookat_target: np.ndarray | None = None,
    lookat_body_name: str | None = None,
    camera_up: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Create camera position and orientation vectors using enhanced lookat approach.
    Functionally equivalent to create_lookat_pose, but without use of a reference body.

    Args:
        env: The environment instance (CPUMujocoEnv)
        camera_pos_world: Camera position in world coordinates
        rpy: Roll, pitch, yaw (currently unused but kept for compatibility)
        lookat_target: Optional 3D point (np.array) to look at in world coordinates
        lookat_body_name: Optional body name to look at (used if lookat_target is None)
        camera_up: Optional desired up direction for camera (np.array in world coordinates)
                  If None, uses world Z-up [0, 0, 1]
    Returns:
        Tuple of (camera_pos_world, forward_vector, up_vector), each of shape (3,)
    """
    from molmo_spaces.utils.pose import compute_lookat_forward_up

    if lookat_target is None:
        if lookat_body_name is not None:
            from molmo_spaces.env.data_views import create_mlspaces_body

            lookat_target = create_mlspaces_body(env.current_data, lookat_body_name).position
        else:
            raise ValueError("Either lookat_target or lookat_body_name must be provided")

    forward, up = compute_lookat_forward_up(camera_pos_world, lookat_target, camera_up)
    return camera_pos_world, forward, up
create_quaternion_camera_pose 
create_quaternion_camera_pose(env, reference_body_name: str, camera_offset: ndarray, camera_quaternion: ndarray) -> tuple[ndarray, ndarray, ndarray]

Create camera pose using quaternion-based orientation relative to reference body.

Parameters:

Name Type Description Default
env

The environment instance (CPUMujocoEnv)

 required
reference_body_name str

Name of the body to attach camera to

 required
camera_offset ndarray

Camera position relative to reference body frame

 required
camera_quaternion ndarray

Quaternion [w, x, y, z] relative to reference body frame

 required

Returns:

Type Description
tuple[ndarray, ndarray, ndarray]

Tuple of (camera_pos_world, forward_vector, up_vector)
Source code in molmo_spaces/env/camera_manager.py 
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
def create_quaternion_camera_pose(
    self,
    env,
    reference_body_name: str,
    camera_offset: np.ndarray,
    camera_quaternion: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Create camera pose using quaternion-based orientation relative to reference body.

    Args:
        env: The environment instance (CPUMujocoEnv)
        reference_body_name: Name of the body to attach camera to
        camera_offset: Camera position relative to reference body frame
        camera_quaternion: Quaternion [w, x, y, z] relative to reference body frame

    Returns:
        Tuple of (camera_pos_world, forward_vector, up_vector)
    """
    from scipy.spatial.transform import Rotation as R

    from molmo_spaces.env.data_views import create_mlspaces_body

    # Get reference body for transformations
    reference_body = create_mlspaces_body(env.current_data, reference_body_name)

    # Calculate camera position in world coordinates
    camera_pos_world = reference_body.pose[:3, :3] @ camera_offset + reference_body.pose[:3, 3]

    # Calculate camera orientation in world coordinates
    norm_quat = np.linalg.norm(camera_quaternion)
    if norm_quat > 0:
        camera_quaternion = camera_quaternion / norm_quat

    camera_rotation_matrix = R.from_quat(camera_quaternion, scalar_first=True).as_matrix()
    world_rotation = reference_body.pose[:3, :3] @ camera_rotation_matrix

    # Extract forward and up vectors from the rotation matrix
    # Camera convention: forward is negative Z-axis, up is Y-axis
    forward = -world_rotation[:, 2]  # Negative Z-axis
    up = world_rotation[:, 1]  # Y-axis

    return camera_pos_world, forward, up
create_robot_mounted_camera_pose 
create_robot_mounted_camera_pose(env, reference_body_name: str, camera_offset: ndarray, lookat_offset: ndarray = None, up_axis: str = 'z') -> tuple[ndarray, ndarray, ndarray]

Create generic robot-mounted camera pose using the existing create_lookat_pose method. Camera positioned relative to reference body, looking at a specific point offset from the body, with camera up aligned with the specified local axis.

Parameters:

Name Type Description Default
env

The environment instance (CPUMujocoEnv)

 required
reference_body_name str

Name of the body to attach camera to

 required
camera_offset ndarray

Camera position relative to reference body frame

 required
lookat_offset ndarray

Offset from reference body to look at (default: 8cm forward along local Z axis)

 None
up_axis str

Which local axis of reference frame is "up" ("x", "y", or "z")

 'z'
Source code in molmo_spaces/env/camera_manager.py 
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
def create_robot_mounted_camera_pose(
    self,
    env,
    reference_body_name: str,
    camera_offset: np.ndarray,
    lookat_offset: np.ndarray = None,
    up_axis: str = "z",
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Create generic robot-mounted camera pose using the existing create_lookat_pose method.
    Camera positioned relative to reference body, looking at a specific point offset from the body,
    with camera up aligned with the specified local axis.

    Args:
        env: The environment instance (CPUMujocoEnv)
        reference_body_name: Name of the body to attach camera to
        camera_offset: Camera position relative to reference body frame
        lookat_offset: Offset from reference body to look at (default: 8cm forward along local Z axis)
        up_axis: Which local axis of reference frame is "up" ("x", "y", or "z")
    """
    from molmo_spaces.env.data_views import create_mlspaces_body

    # Default lookat offset: look forward along reference body Z axis
    if lookat_offset is None:
        lookat_offset = np.array([0.0, 0.0, 0.08])  # 8cm forward along local Z axis

    # Get reference body for transformations
    reference_body = create_mlspaces_body(env.current_data, reference_body_name)

    # Calculate the target point in world coordinates
    lookat_target_world = (
        reference_body.pose[:3, :3] @ lookat_offset + reference_body.pose[:3, 3]
    )

    # Get the reference body's up direction in world coordinates
    local_up_map = {
        "x": np.array([1.0, 0.0, 0.0]),
        "y": np.array([0.0, 1.0, 0.0]),
        "z": np.array([0.0, 0.0, 1.0]),
    }
    if up_axis not in local_up_map:
        raise ValueError(f"up_axis must be 'x', 'y', or 'z', got {up_axis}")

    local_up = local_up_map[up_axis]
    reference_up_world = reference_body.pose[:3, :3] @ local_up

    # Use create_lookat_pose with the 3D target point and reference body's up direction
    return self.create_lookat_pose(
        env,
        camera_offset,
        np.zeros(3),
        reference_body_name,
        lookat_target=lookat_target_world,
        camera_up=reference_up_world,
    )
setup_cameras 
setup_cameras(env, camera_system_config: CameraSystemConfig, workspace_center=None, visibility_resolver: Callable[[str], list[str]] | None = None, deterministic_only: bool = False) -> None

Set up all cameras from a CameraSystemConfig.

This is the main entry point for camera setup. It processes each camera spec and delegates to the appropriate setup method based on camera type.

Parameters:

Name Type Description Default
env

The environment instance (CPUMujocoEnv)

 required
camera_system_config CameraSystemConfig

CameraSystemConfig instance with all camera specs

 required
workspace_center

Optional workspace center position (np.ndarray) for camera placement

 None
visibility_resolver Callable[[str], list[str]] | None

Optional callable(key: str) -> str that resolves special visibility keys

 None
Source code in molmo_spaces/env/camera_manager.py 
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
def setup_cameras(
    self,
    env,
    camera_system_config: CameraSystemConfig,
    workspace_center=None,
    visibility_resolver: Callable[[str], list[str]] | None = None,
    deterministic_only: bool = False,
) -> None:
    """Set up all cameras from a CameraSystemConfig.

    This is the main entry point for camera setup. It processes each camera spec
    and delegates to the appropriate setup method based on camera type.

    Args:
        env: The environment instance (CPUMujocoEnv)
        camera_system_config: CameraSystemConfig instance with all camera specs
        workspace_center: Optional workspace center position (np.ndarray) for camera placement
        visibility_resolver: Optional callable(key: str) -> str that resolves special visibility keys
    """
    from molmo_spaces.configs.camera_configs import (
        FixedExocentricCameraConfig,
        MjcfCameraConfig,
        RandomizedExocentricCameraConfig,
        RobotMountedCameraConfig,
    )

    # Store workspace center and visibility resolver for use in camera setup methods
    self._workspace_center = workspace_center
    self._visibility_resolver = visibility_resolver

    log.info(f"[CAMERA SETUP] Setting up {len(camera_system_config.cameras)} cameras")

    # allow runtime errors to propagate here instead of catching them and logging
    for camera_spec in camera_system_config.cameras:
        if isinstance(camera_spec, MjcfCameraConfig):
            self._setup_mjcf_camera(env, camera_spec)
        elif isinstance(camera_spec, RobotMountedCameraConfig):
            self._setup_robot_mounted_camera(env, camera_spec)
        elif isinstance(camera_spec, FixedExocentricCameraConfig):
            self._setup_fixed_exocentric_camera(env, camera_spec)
        elif isinstance(camera_spec, RandomizedExocentricCameraConfig):
            if deterministic_only:
                log.info(
                    f"[CAMERA SETUP] skipping randomized exocentric camera '{camera_spec.name}' for pre-workspace setup"
                )
            else:
                self._setup_randomized_exocentric_camera(env, camera_spec)
        else:
            log.warning(
                f"[CAMERA SETUP] Unknown camera spec type: {type(camera_spec).__name__}"
            )

    # Clean up
    self._workspace_center = None
    self._visibility_resolver = None

    log.info(f"[CAMERA SETUP] Successfully set up {len(self.registry.cameras)} cameras")

CameraRegistry 

CameraRegistry()

Registry for camera objects with auto-updating support.

Methods:

Name Description
__contains__
__getitem__
__iter__
__setitem__
add_camera

Adds a camera object to the registry.
add_static_camera

Adds a static camera to the registry.
keys
update_all_cameras

Update all cameras that need updating. Returns list of cameras that changed.

Attributes:

Name Type Description
cameras dict[str, Camera]
Source code in molmo_spaces/env/camera_manager.py 
210
211
def __init__(self) -> None:
    self.cameras: dict[str, Camera] = {}
cameras instance-attribute 
cameras: dict[str, Camera] = {}
__contains__ 
__contains__(key: str) -> bool
Source code in molmo_spaces/env/camera_manager.py 
245
246
def __contains__(self, key: str) -> bool:
    return key in self.cameras
__getitem__ 
__getitem__(key: str) -> Camera
Source code in molmo_spaces/env/camera_manager.py 
235
236
237
238
239
240
def __getitem__(self, key: str) -> Camera:
    try:
        return self.cameras[key]
    except KeyError as e:
        log.warning(f"For KeyError, camera options are options are: {self.cameras.keys()}")
        raise e
__iter__ 
__iter__()
Source code in molmo_spaces/env/camera_manager.py 
251
252
def __iter__(self):
    return iter(self.cameras.values())
__setitem__ 
__setitem__(key: str, value: Camera) -> None
Source code in molmo_spaces/env/camera_manager.py 
242
243
def __setitem__(self, key: str, value: Camera) -> None:
    self.cameras[key] = value
add_camera 
add_camera(camera: Camera) -> None

Adds a camera object to the registry.

Source code in molmo_spaces/env/camera_manager.py 
213
214
215
def add_camera(self, camera: Camera) -> None:
    """Adds a camera object to the registry."""
    self.cameras[camera.name] = camera
add_static_camera 
add_static_camera(name: str, pos: NDArray[float32], forward: NDArray[float32], up: NDArray[float32]) -> None

Adds a static camera to the registry.

Source code in molmo_spaces/env/camera_manager.py 
217
218
219
220
221
222
223
224
225
def add_static_camera(
    self,
    name: str,
    pos: NDArray[np.float32],
    forward: NDArray[np.float32],
    up: NDArray[np.float32],
) -> None:
    """Adds a static camera to the registry."""
    self.cameras[name] = Camera(name=name, pos=pos, forward=forward, up=up)
keys 
keys()
Source code in molmo_spaces/env/camera_manager.py 
248
249
def keys(self):
    return self.cameras.keys()
update_all_cameras 
update_all_cameras(env: CPUMujocoEnv) -> list[str]

Update all cameras that need updating. Returns list of cameras that changed.

Source code in molmo_spaces/env/camera_manager.py 
227
228
229
230
231
232
233
def update_all_cameras(self, env: CPUMujocoEnv) -> list[str]:
    """Update all cameras that need updating. Returns list of cameras that changed."""
    updated_cameras: list[str] = []
    for camera in self.cameras.values():
        if camera.update_pose(env):
            updated_cameras.append(camera.name)
    return updated_cameras

RobotMountedCamera 

RobotMountedCamera(name: str, reference_body_names: str | list[str], camera_offset: NDArray[float32] | list[float] | None = None, lookat_offset: NDArray[float32] | list[float] | None = None, up_axis: str = 'z', camera_quaternion: NDArray[float32] | list[float] | None = None, camera_fov: float = 45)

Bases: Camera

Generic robot-mounted camera that can attach to any body/joint with configurable offsets.

Parameters:

Name Type Description Default
name str

Camera name

 required
reference_body_names str | list[str]

Body name(s) to attach camera to. If list, tries each until one works.

 required
camera_offset NDArray[float32] | list[float] | None

Camera position relative to reference body frame

 None
lookat_offset NDArray[float32] | list[float] | None

Offset from reference body to look at (used when camera_quaternion is None)

 None
up_axis str

Which local axis of reference frame is "up" ("x", "y", or "z") (used when camera_quaternion is None)

 'z'
camera_quaternion NDArray[float32] | list[float] | None

Quaternion [w, x, y, z] relative to reference body frame. If provided, overrides lookat_offset and up_axis

 None

Methods:

Name Description
get_pose

return 4x4 pose
update_pose

Update camera pose based on current reference body state. Returns True if pose changed.

Attributes:

Name Type Description
camera_offset NDArray[float32]
camera_quaternion NDArray[float32] | None
forward NDArray[float32]
fov float
lookat_offset NDArray[float32]
name str
pos NDArray[float32]
reference_body_names list[str]
up NDArray[float32]
up_axis str
Source code in molmo_spaces/env/camera_manager.py 
 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
def __init__(
    self,
    name: str,
    reference_body_names: str | list[str],
    camera_offset: NDArray[np.float32] | list[float] | None = None,
    lookat_offset: NDArray[np.float32] | list[float] | None = None,
    up_axis: str = "z",
    camera_quaternion: NDArray[np.float32] | list[float] | None = None,
    camera_fov: float = 45,
) -> None:
    """
    Args:
        name: Camera name
        reference_body_names: Body name(s) to attach camera to. If list, tries each until one works.
        camera_offset: Camera position relative to reference body frame
        lookat_offset: Offset from reference body to look at (used when camera_quaternion is None)
        up_axis: Which local axis of reference frame is "up" ("x", "y", or "z") (used when camera_quaternion is None)
        camera_quaternion: Quaternion [w, x, y, z] relative to reference body frame. If provided, overrides lookat_offset and up_axis
    """
    # Store configuration for dynamic updates
    self.reference_body_names: list[str] = (
        [reference_body_names]
        if isinstance(reference_body_names, str)
        else reference_body_names
    )
    self.camera_offset: NDArray[np.float32] = (
        np.array(camera_offset, dtype=np.float32)
        if camera_offset is not None
        else np.array([0.10, 0.0, -0.15], dtype=np.float32)
    )
    self.lookat_offset: NDArray[np.float32] = (
        np.array(lookat_offset, dtype=np.float32)
        if lookat_offset is not None
        else np.array([0.0, 0.0, 0.08], dtype=np.float32)
    )
    self.up_axis: str = up_axis
    self.camera_quaternion: NDArray[np.float32] | None = (
        np.array(camera_quaternion, dtype=np.float32) if camera_quaternion is not None else None
    )

    # Initialize with default pose (will be updated on first update_pose call)
    super().__init__(name, fov=camera_fov)

    # Cache for avoiding unnecessary recalculations
    self._last_reference_pose: NDArray[np.float32] | None = None
    self._active_reference_body_name: str | None = None
camera_offset instance-attribute 
camera_offset: NDArray[float32] = array(camera_offset, dtype=float32) if camera_offset is not None else array([0.1, 0.0, -0.15], dtype=float32)
camera_quaternion instance-attribute 
camera_quaternion: NDArray[float32] | None = array(camera_quaternion, dtype=float32) if camera_quaternion is not None else None
forward instance-attribute 
forward: NDArray[float32] = forward if forward is not None else array([1.0, 0.0, 0.0], dtype=float32)
fov instance-attribute 
fov: float = fov
lookat_offset instance-attribute 
lookat_offset: NDArray[float32] = array(lookat_offset, dtype=float32) if lookat_offset is not None else array([0.0, 0.0, 0.08], dtype=float32)
name instance-attribute 
name: str = name
pos instance-attribute 
pos: NDArray[float32] = pos if pos is not None else array([0.0, 0.0, 1.0], dtype=float32)
reference_body_names instance-attribute 
reference_body_names: list[str] = [reference_body_names] if isinstance(reference_body_names, str) else reference_body_names
up instance-attribute 
up: NDArray[float32] = up if up is not None else array([0.0, 0.0, 1.0], dtype=float32)
up_axis instance-attribute 
up_axis: str = up_axis
get_pose 
get_pose() -> NDArray[float32]

return 4x4 pose

Source code in molmo_spaces/env/camera_manager.py 
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
def get_pose(self) -> NDArray[np.float32]:
    """
    return 4x4 pose
    """
    # Validate and normalize camera vectors
    forward_norm = np.linalg.norm(self.forward)
    up_norm = np.linalg.norm(self.up)

    if forward_norm < 1e-6 or up_norm < 1e-6:
        print(
            f"Warning: Camera '{self.camera_name}' has degenerate vectors (forward_norm={forward_norm}, up_norm={up_norm})"
        )
        return np.eye(4, 4, dtype=np.float32)

    forward = self.forward / forward_norm
    up = self.up / up_norm
    right = np.cross(forward, up)

    right_norm = np.linalg.norm(right)
    if right_norm < 1e-6:
        print(f"Warning: Camera '{self.self}' has collinear forward/up vectors")
        return np.eye(4, 4, dtype=np.float32)

    right = right / right_norm

    # Recompute orthogonal up to ensure proper orthogonal basis
    up = np.cross(right, forward)

    # Create cam2world matrix (standard camera convention)
    world2cam = np.eye(4)
    world2cam[:3, 0] = right  # X-axis (right)
    world2cam[:3, 1] = -up  # Y-axis (up)
    world2cam[:3, 2] = forward  # Z-axis - camera looks down negative Z
    world2cam[:3, 3] = self.pos  # Translation
    return world2cam
update_pose 
update_pose(env: CPUMujocoEnv) -> bool

Update camera pose based on current reference body state. Returns True if pose changed.

Source code in molmo_spaces/env/camera_manager.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
def update_pose(self, env: CPUMujocoEnv) -> bool:
    """Update camera pose based on current reference body state. Returns True if pose changed."""

    # Find reference body
    reference_body, body_name = self._find_reference_body(env)
    if reference_body is None:
        log.warning(
            f"No valid reference body found for camera {self.name} {self.reference_body_names}"
        )
        return False  # Cannot find any reference body

    # Update active reference if it changed
    if self._active_reference_body_name != body_name:
        self._active_reference_body_name = body_name
        self._last_reference_pose = None  # Force update

    current_reference_pose = reference_body.pose.copy()

    # Check if reference pose has changed (with small tolerance for numerical precision)
    if self._last_reference_pose is not None:
        pose_diff = np.linalg.norm(current_reference_pose - self._last_reference_pose)
        if pose_diff < 1e-6:  # No significant change
            return False

    # Reference pose has changed, recalculate camera pose
    if self.camera_quaternion is not None:
        pos, forward, up = env.camera_manager.create_quaternion_camera_pose(
            env, self._active_reference_body_name, self.camera_offset, self.camera_quaternion
        )
    else:
        pos, forward, up = env.camera_manager.create_robot_mounted_camera_pose(
            env,
            self._active_reference_body_name,
            self.camera_offset,
            self.lookat_offset,
            self.up_axis,
        )

    # Update our pose
    self.pos = pos
    self.forward = forward
    self.up = up

    # Cache the reference pose
    self._last_reference_pose = current_reference_pose

    return True

data_views

Classes:

Name Description
Door
MlSpacesArticulationObject
MlSpacesBody
MlSpacesCamera
MlSpacesFreeJointBody
MlSpacesImmovableBody
MlSpacesMocapBody
MlSpacesObject

Class to represent scene object data
MlSpacesObjectAbstract

Functions:

Name Description
create_mlspaces_body

Door 

Door(door_name: str, data: MjData)

Bases: MlSpacesArticulationObject

Methods:

Name Description
body_name2id
body_parent_id
build_children_lists
find_top_object_body_id
get_ancestors
get_descendants
get_direct_children
get_friction
get_geom_infos

Get geom information for this object.
get_geom_type_name
get_handle_bboxes_array

Get handle bounding boxes as an array.
get_handle_joint_index

Get the index of the handle joint (the joint that has 'handle' in its name).
get_handle_pose

Get handle pose (position + quaternion) as a single array.
get_hinge_joint_index

Get the index of the door hinge joint (the joint that does not have 'handle' in its name).
get_joint_anchor_position
get_joint_armature
get_joint_axis
get_joint_body_orientation
get_joint_damping
get_joint_frictionloss
get_joint_leaf_body_position
get_joint_position
get_joint_qpos_adr
get_joint_range
get_joint_stiffness
get_joint_type
get_mass
get_swing_arc_circle
get_top_level_bodies

Return bodies whose parent is the world body.
handle_name
is_child_name_of
is_object_in_gripper
is_object_picked_up
is_point_in_swing_arc
set_joint_position

Attributes:

Name Type Description
aabb_center ndarray
aabb_size ndarray
body_id int
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).
bvh_root
center_of_mass
door_name str
geom_ids

Get all geom IDs belonging to this object (lazy, cached).
joint_id2name
joint_id2qpos_adr
joint_ids
joint_names
mj_data
mj_model
name str
njoints int
num_handles int
object_id int
object_root_id
pose ndarray
position
quat
Source code in molmo_spaces/env/data_views.py 
606
607
608
609
610
611
def __init__(
    self,
    door_name: str,
    data: mujoco.MjData,
) -> None:
    super().__init__(data=data, object_name=door_name)
aabb_center property 
aabb_center: ndarray
aabb_size property 
aabb_size: ndarray
body_id cached property 
body_id: int
body_ids cached property 
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).

bvh_root property 
bvh_root
center_of_mass property 
center_of_mass
door_name property 
door_name: str
geom_ids cached property 
geom_ids

Get all geom IDs belonging to this object (lazy, cached).

joint_id2name instance-attribute 
joint_id2name = {}
joint_id2qpos_adr instance-attribute 
joint_id2qpos_adr = {}
joint_ids instance-attribute 
joint_ids = []
joint_names instance-attribute 
joint_names = []
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
njoints property 
njoints: int
num_handles property 
num_handles: int
object_id property 
object_id: int
object_root_id cached property 
object_root_id
pose property writable 
pose: ndarray
position property 
position
quat property 
quat
body_name2id staticmethod 
body_name2id(model: MjModel) -> dict[str, int]
Source code in molmo_spaces/env/data_views.py 
467
468
469
@staticmethod
def body_name2id(model: mujoco.MjModel) -> dict[str, int]:
    return {model.body(i).name: i for i in range(model.nbody)}
body_parent_id staticmethod 
body_parent_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
387
388
389
390
391
392
393
@staticmethod
def body_parent_id(model: mujoco.MjModel, body_id: int) -> int:
    try:
        return int(model.body_parentid[body_id])
    except Exception as e:
        print(f"Error getting body parent id: {e}")
        return int(model.body(body_id).parentid)
build_children_lists staticmethod 
build_children_lists(model: MjModel) -> list[list[int]]
Source code in molmo_spaces/env/data_views.py 
395
396
397
398
399
400
401
402
@staticmethod
def build_children_lists(model: mujoco.MjModel) -> list[list[int]]:
    children: list[list[int]] = [[] for _ in range(model.nbody)]
    for child_id in range(model.nbody):
        pid = MlSpacesObject.body_parent_id(model, child_id)
        if pid >= 0:
            children[pid].append(child_id)
    return children
find_top_object_body_id staticmethod 
find_top_object_body_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
@staticmethod
def find_top_object_body_id(model: mujoco.MjModel, body_id: int) -> int:
    cur = body_id
    visited = set()
    while True:
        if cur in visited:
            break
        visited.add(cur)
        pid = MlSpacesObject.body_parent_id(model, cur)
        if pid < 0:
            break
        pname = model.body(pid).name
        cname = model.body(cur).name
        if MlSpacesObject.is_child_name_of(pname, cname):
            cur = pid
        else:
            break
    return int(cur)
get_ancestors staticmethod 
get_ancestors(model: MjModel, body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
404
405
406
407
408
409
410
411
412
413
@staticmethod
def get_ancestors(model: mujoco.MjModel, body_id: int) -> list[int]:
    ancestors: list[int] = []
    visited = {body_id}
    pid = MlSpacesObject.body_parent_id(model, body_id)
    while pid >= 0 and pid not in visited:
        ancestors.append(pid)
        visited.add(pid)
        pid = MlSpacesObject.body_parent_id(model, pid)
    return ancestors
get_descendants staticmethod 
get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
415
416
417
418
419
420
421
422
423
424
@staticmethod
def get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]:
    stack = [body_id]
    desc: list[int] = []
    while stack:
        cur = stack.pop()
        for c in children_lists[cur]:
            desc.append(int(c))
            stack.append(c)
    return desc
get_direct_children staticmethod 
get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
426
427
428
@staticmethod
def get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]:
    return list(children_lists[body_id])
get_friction 
get_friction()
Source code in molmo_spaces/env/data_views.py 
381
382
def get_friction(self):
    return [self.mj_model.geom(geom_id).friction for geom_id in self.geom_ids]
get_geom_infos 
get_geom_infos(include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]

Get geom information for this object.

Parameters:

Name Type Description Default
include_descendants bool

If True, includes geoms from all descendant bodies. If False, only includes geoms directly attached to this object's body.

 True
max_geoms int | None

Maximum number of geoms to return. If None, returns all.

 2048

Returns:

Type Description
list[dict[str, object]]

List of dicts with geom info: id, name, position, size, type, type_name
Source code in molmo_spaces/env/data_views.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def get_geom_infos(
    self, include_descendants: bool = True, max_geoms: int | None = 2048
) -> list[dict[str, object]]:
    """Get geom information for this object.

    Args:
        include_descendants: If True, includes geoms from all descendant bodies.
            If False, only includes geoms directly attached to this object's body.
        max_geoms: Maximum number of geoms to return. If None, returns all.

    Returns:
        List of dicts with geom info: id, name, position, size, type, type_name
    """
    # Build set of body IDs to include
    if include_descendants:
        body_ids = set(self.body_ids)
    else:
        body_ids = {self.object_id}

    # Find geoms belonging to these bodies
    geoms: list[dict[str, object]] = []
    count = 0
    for geom_id in range(self.mj_model.ngeom):
        if max_geoms is not None and count >= max_geoms:
            break
        if int(self.mj_model.geom_bodyid[geom_id]) in body_ids:
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id))
            pos = self.mj_data.geom_xpos[geom_id].copy()
            size = self.mj_model.geom_size[geom_id].copy()
            gtype = int(self.mj_model.geom_type[geom_id])
            geoms.append(
                {
                    "id": int(geom_id),
                    "name": name,
                    "position": pos,
                    "size": size,
                    "type": gtype,
                    "type_name": MlSpacesObject.get_geom_type_name(gtype),
                }
            )
            count += 1

    return geoms
get_geom_type_name staticmethod 
get_geom_type_name(geom_type: int) -> str
Source code in molmo_spaces/env/data_views.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
@staticmethod
def get_geom_type_name(geom_type: int) -> str:
    names = {
        0: "PLANE",
        1: "HFIELD",
        2: "SPHERE",
        3: "CAPSULE",
        4: "ELLIPSOID",
        5: "CYLINDER",
        6: "BOX",
        7: "MESH",
    }
    return names.get(int(geom_type), f"UNKNOWN({geom_type})")
get_handle_bboxes_array 
get_handle_bboxes_array() -> ndarray

Get handle bounding boxes as an array. Finds leaf bodies in the door hierarchy (assumed to be handles) and returns their visual geom AABBs. Returns: np.ndarray: Array of AABBs (center, size) for handle visual geoms

Source code in molmo_spaces/env/data_views.py 
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
def get_handle_bboxes_array(self) -> np.ndarray:
    """Get handle bounding boxes as an array.
    Finds leaf bodies in the door hierarchy (assumed to be handles) and
    returns their visual geom AABBs.
    Returns:
        np.ndarray: Array of AABBs (center, size) for handle visual geoms
    """
    model = self.mj_model
    leaf_body_ids = self._get_handle_leaf_body_ids()
    handle_bboxes_array = []
    # Find visual geoms for leaf bodies
    for leaf_body_id in leaf_body_ids:
        leaf_ginds = np.where(model.geom_bodyid == leaf_body_id)[0]
        for leaf_gind in leaf_ginds:
            if model.geom(leaf_gind).contype == 0 and model.geom(leaf_gind).conaffinity == 0:
                # visual geom (this is usually the target geom of the handle)
                aabb = model.geom_aabb[leaf_gind]  # axis aligned bounding box (center, size)
                handle_bboxes_array.append(aabb)

    return np.array(handle_bboxes_array)
get_handle_joint_index 
get_handle_joint_index() -> int

Get the index of the handle joint (the joint that has 'handle' in its name). Returns: int: Index of the handle joint in self.joint_ids list

Source code in molmo_spaces/env/data_views.py 
663
664
665
666
667
668
669
670
671
672
673
674
675
def get_handle_joint_index(self) -> int:
    """Get the index of the handle joint (the joint that has 'handle' in its name).
    Returns:
        int: Index of the handle joint in self.joint_ids list
    """
    for i, joint_name in enumerate(self.joint_names):
        if "handle" in joint_name.lower():
            return i

    raise ValueError(
        f"No joint found on door '{self.name}' that contains 'handle' in its name. "
        f"Available joints: {self.joint_names}"
    )
get_handle_pose 
get_handle_pose() -> ndarray

Get handle pose (position + quaternion) as a single array. Finds the first handle visual geom and returns its pose. If multiple handles exist, returns the first one. Returns: np.ndarray: Handle pose [x, y, z, qw, qx, qy, qz]

Source code in molmo_spaces/env/data_views.py 
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def get_handle_pose(self) -> np.ndarray:
    """Get handle pose (position + quaternion) as a single array.
    Finds the first handle visual geom and returns its pose.
    If multiple handles exist, returns the first one.
    Returns:
        np.ndarray: Handle pose [x, y, z, qw, qx, qy, qz]
    """
    model = self.mj_model
    leaf_body_ids = self._get_handle_leaf_body_ids()

    # Find first visual geom for leaf bodies
    for leaf_body_id in leaf_body_ids:
        leaf_ginds = np.where(model.geom_bodyid == leaf_body_id)[0]
        for leaf_gind in leaf_ginds:
            if model.geom(leaf_gind).contype == 0 and model.geom(leaf_gind).conaffinity == 0:
                # visual geom - get its pose
                geom_pos = self.mj_data.geom_xpos[leaf_gind]
                geom_rot_mat = self.mj_data.geom_xmat[leaf_gind].reshape(3, 3)
                geom_quat = R.from_matrix(geom_rot_mat).as_quat(scalar_first=True)
                geom_pose = np.concatenate([geom_pos, geom_quat])
                return geom_pose

    # No handle found, return zero pose
    warnings.warn(
        f"No handle visual geom found for door '{self.name}'", UserWarning, stacklevel=2
    )
    return np.zeros(7)  # 3 for position + 4 for quaternion
get_hinge_joint_index 
get_hinge_joint_index() -> int

Get the index of the door hinge joint (the joint that does not have 'handle' in its name). Returns: int: Index of the hinge joint in self.joint_ids list (not the joint ID itself) Warns: UserWarning: If no joint without 'handle' in its name is found.

Source code in molmo_spaces/env/data_views.py 
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
def get_hinge_joint_index(self) -> int:
    """Get the index of the door hinge joint (the joint that does not have 'handle' in its name).
    Returns:
        int: Index of the hinge joint in self.joint_ids list (not the joint ID itself)
    Warns:
        UserWarning: If no joint without 'handle' in its name is found.
    """
    for i, joint_name in enumerate(self.joint_names):
        if "handle" not in joint_name.lower():
            return i

    raise ValueError(
        f"No joint found on door '{self.name}' that does not contain 'handle' in its name. "
        f"Available joints: {self.joint_names}"
    )
get_joint_anchor_position 
get_joint_anchor_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
550
551
552
553
554
555
556
557
558
559
def get_joint_anchor_position(self, i: int) -> np.ndarray:
    # Get the position in world frame of the joint anchor point
    joint_id = self.joint_ids[i]
    body_id = self.mj_model.jnt_bodyid[joint_id]
    local_anchor = self.mj_model.jnt_pos[joint_id]
    body_pos = self.mj_data.xpos[body_id]
    body_quat = self.mj_data.xquat[body_id]
    body_rot = R.from_quat(body_quat, scalar_first=True).as_matrix()
    world_anchor = body_pos + body_rot @ local_anchor
    return world_anchor
get_joint_armature 
get_joint_armature()
Source code in molmo_spaces/env/data_views.py 
570
571
def get_joint_armature(self):
    return [self.mj_model.joint(joint_id).armature for joint_id in self.joint_ids]
get_joint_axis 
get_joint_axis(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
526
527
def get_joint_axis(self, i: int) -> np.ndarray:
    return self.mj_model.joint(self.joint_ids[i]).axis
get_joint_body_orientation 
get_joint_body_orientation(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
545
546
547
548
def get_joint_body_orientation(self, i: int) -> np.ndarray:
    # orientation of the joint body not the root orientation of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    return self.mj_data.xmat[body_id].copy().reshape(3, 3)
get_joint_damping 
get_joint_damping()
Source code in molmo_spaces/env/data_views.py 
567
568
def get_joint_damping(self):
    return [self.mj_model.joint(joint_id).damping for joint_id in self.joint_ids]
get_joint_frictionloss 
get_joint_frictionloss()
Source code in molmo_spaces/env/data_views.py 
561
562
def get_joint_frictionloss(self):
    return [self.mj_model.joint(joint_id).frictionloss for joint_id in self.joint_ids]
get_joint_leaf_body_position 
get_joint_leaf_body_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
535
536
537
538
539
540
541
542
543
def get_joint_leaf_body_position(self, i: int) -> np.ndarray:
    # position of the joint body not the root center of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    # find child body of this joint body
    for child_body_id in range(self.mj_model.nbody):
        if self.mj_model.body(child_body_id).parentid[0] == body_id:
            body_id = child_body_id
            break
    return self.mj_data.xpos[body_id].copy()
get_joint_position 
get_joint_position(i: int) -> float
Source code in molmo_spaces/env/data_views.py 
520
521
def get_joint_position(self, i: int) -> float:
    return self.mj_data.qpos[self.get_joint_qpos_adr(i)].copy()
get_joint_qpos_adr 
get_joint_qpos_adr(i: int) -> int
Source code in molmo_spaces/env/data_views.py 
532
533
def get_joint_qpos_adr(self, i: int) -> int:
    return int(self.mj_model.joint(self.joint_ids[i]).qposadr[0])
get_joint_range 
get_joint_range(i: int) -> tuple[float, float]
Source code in molmo_spaces/env/data_views.py 
523
524
def get_joint_range(self, i: int) -> tuple[float, float]:
    return self.mj_model.joint(self.joint_ids[i]).range
get_joint_stiffness 
get_joint_stiffness()
Source code in molmo_spaces/env/data_views.py 
564
565
def get_joint_stiffness(self):
    return [self.mj_model.joint(joint_id).stiffness for joint_id in self.joint_ids]
get_joint_type 
get_joint_type(i: int) -> mjtJoint
Source code in molmo_spaces/env/data_views.py 
529
530
def get_joint_type(self, i: int) -> mujoco.mjtJoint:
    return self.mj_model.joint(self.joint_ids[i]).type
get_mass 
get_mass()
Source code in molmo_spaces/env/data_views.py 
384
385
def get_mass(self):
    return [self.mj_model.body_mass[body_id] for body_id in self.body_ids]
get_swing_arc_circle 
get_swing_arc_circle() -> dict[str, ndarray | float]
Source code in molmo_spaces/env/data_views.py 
743
744
745
746
747
748
749
750
751
752
753
754
def get_swing_arc_circle(self) -> dict[str, np.ndarray | float]:
    hinge_joint_idx = self.get_hinge_joint_index()
    center = self.get_joint_anchor_position(hinge_joint_idx)

    # Calculate radius as door extent on Z axis (height)
    # Use the door's AABB size in Z direction
    radius = self.aabb_size[2]

    return {
        "center": center,
        "radius": float(radius),
    }
get_top_level_bodies staticmethod 
get_top_level_bodies(model: MjModel) -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/data_views.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
@staticmethod
def get_top_level_bodies(model: mujoco.MjModel) -> list[int]:
    """Return bodies whose parent is the world body."""
    # Identify world body id: the one with parent < 0
    try:
        world_id = next(
            b for b in range(model.nbody) if MlSpacesObject.body_parent_id(model, b) < 0
        )
    except StopIteration:
        world_id = 0
    tops: list[int] = []
    for b in range(model.nbody):
        try:
            pid = MlSpacesObject.body_parent_id(model, b)
            if pid == world_id:
                tops.append(b)
        except Exception as e:
            print(f"Error getting top-level bodies: {e}")
            continue
    return tops
handle_name 
handle_name(handle_id=0) -> str
Source code in molmo_spaces/env/data_views.py 
621
622
def handle_name(self, handle_id=0) -> str:
    return self.mj_model.body(self._get_handle_leaf_body_ids()[handle_id]).name
is_child_name_of staticmethod 
is_child_name_of(parent_name: str, child_name: str) -> bool
Source code in molmo_spaces/env/data_views.py 
430
431
432
@staticmethod
def is_child_name_of(parent_name: str, child_name: str) -> bool:
    return child_name.startswith(parent_name + "_")
is_object_in_gripper 
is_object_in_gripper() -> None
Source code in molmo_spaces/env/data_views.py 
277
278
def is_object_in_gripper(self) -> None:
    pass
is_object_picked_up 
is_object_picked_up() -> None
Source code in molmo_spaces/env/data_views.py 
274
275
def is_object_picked_up(self) -> None:
    pass
is_point_in_swing_arc 
is_point_in_swing_arc(point: ndarray, safety_margin: float = 0.1) -> bool
Source code in molmo_spaces/env/data_views.py 
756
757
758
759
760
761
762
763
764
765
766
def is_point_in_swing_arc(self, point: np.ndarray, safety_margin: float = 0.1) -> bool:
    circle_info = self.get_swing_arc_circle()
    center = circle_info["center"]
    radius = circle_info["radius"] + safety_margin

    # Extract 2D position
    point_2d = point[:2] if len(point) >= 2 else point

    # Check if point is within circle
    dist_from_center = np.linalg.norm(point_2d - center[:2])
    return (dist_from_center <= radius).item()
set_joint_position 
set_joint_position(i: int, position: float) -> None
Source code in molmo_spaces/env/data_views.py 
624
625
626
def set_joint_position(self, i: int, position: float) -> None:
    super().set_joint_position(i, position)
    mujoco.mj_forward(self.mj_model, self.mj_data)

MlSpacesArticulationObject 

MlSpacesArticulationObject(object_name: str, data: MjData)

Bases: MlSpacesObject

Methods:

Name Description
body_name2id
body_parent_id
build_children_lists
find_top_object_body_id
get_ancestors
get_descendants
get_direct_children
get_friction
get_geom_infos

Get geom information for this object.
get_geom_type_name
get_joint_anchor_position
get_joint_armature
get_joint_axis
get_joint_body_orientation
get_joint_damping
get_joint_frictionloss
get_joint_leaf_body_position
get_joint_position
get_joint_qpos_adr
get_joint_range
get_joint_stiffness
get_joint_type
get_mass
get_top_level_bodies

Return bodies whose parent is the world body.
is_child_name_of
is_object_in_gripper
is_object_picked_up
set_joint_position

Attributes:

Name Type Description
aabb_center ndarray
aabb_size ndarray
body_id int
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).
bvh_root
center_of_mass
geom_ids

Get all geom IDs belonging to this object (lazy, cached).
joint_id2name
joint_id2qpos_adr
joint_ids
joint_names
mj_data
mj_model
name str
njoints int
object_id int
object_root_id
pose ndarray
position
quat
Source code in molmo_spaces/env/data_views.py 
494
495
496
497
498
499
500
501
502
def __init__(self, object_name: str, data: mujoco.MjData) -> None:
    super().__init__(object_name, data)

    # gather all joint information of the articulation object
    self.joint_ids = []
    self.joint_names = []
    self.joint_id2name = {}
    self.joint_id2qpos_adr = {}
    self._get_joint_info()
aabb_center property 
aabb_center: ndarray
aabb_size property 
aabb_size: ndarray
body_id cached property 
body_id: int
body_ids cached property 
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).

bvh_root property 
bvh_root
center_of_mass property 
center_of_mass
geom_ids cached property 
geom_ids

Get all geom IDs belonging to this object (lazy, cached).

joint_id2name instance-attribute 
joint_id2name = {}
joint_id2qpos_adr instance-attribute 
joint_id2qpos_adr = {}
joint_ids instance-attribute 
joint_ids = []
joint_names instance-attribute 
joint_names = []
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
njoints property 
njoints: int
object_id property 
object_id: int
object_root_id cached property 
object_root_id
pose property writable 
pose: ndarray
position property 
position
quat property 
quat
body_name2id staticmethod 
body_name2id(model: MjModel) -> dict[str, int]
Source code in molmo_spaces/env/data_views.py 
467
468
469
@staticmethod
def body_name2id(model: mujoco.MjModel) -> dict[str, int]:
    return {model.body(i).name: i for i in range(model.nbody)}
body_parent_id staticmethod 
body_parent_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
387
388
389
390
391
392
393
@staticmethod
def body_parent_id(model: mujoco.MjModel, body_id: int) -> int:
    try:
        return int(model.body_parentid[body_id])
    except Exception as e:
        print(f"Error getting body parent id: {e}")
        return int(model.body(body_id).parentid)
build_children_lists staticmethod 
build_children_lists(model: MjModel) -> list[list[int]]
Source code in molmo_spaces/env/data_views.py 
395
396
397
398
399
400
401
402
@staticmethod
def build_children_lists(model: mujoco.MjModel) -> list[list[int]]:
    children: list[list[int]] = [[] for _ in range(model.nbody)]
    for child_id in range(model.nbody):
        pid = MlSpacesObject.body_parent_id(model, child_id)
        if pid >= 0:
            children[pid].append(child_id)
    return children
find_top_object_body_id staticmethod 
find_top_object_body_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
@staticmethod
def find_top_object_body_id(model: mujoco.MjModel, body_id: int) -> int:
    cur = body_id
    visited = set()
    while True:
        if cur in visited:
            break
        visited.add(cur)
        pid = MlSpacesObject.body_parent_id(model, cur)
        if pid < 0:
            break
        pname = model.body(pid).name
        cname = model.body(cur).name
        if MlSpacesObject.is_child_name_of(pname, cname):
            cur = pid
        else:
            break
    return int(cur)
get_ancestors staticmethod 
get_ancestors(model: MjModel, body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
404
405
406
407
408
409
410
411
412
413
@staticmethod
def get_ancestors(model: mujoco.MjModel, body_id: int) -> list[int]:
    ancestors: list[int] = []
    visited = {body_id}
    pid = MlSpacesObject.body_parent_id(model, body_id)
    while pid >= 0 and pid not in visited:
        ancestors.append(pid)
        visited.add(pid)
        pid = MlSpacesObject.body_parent_id(model, pid)
    return ancestors
get_descendants staticmethod 
get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
415
416
417
418
419
420
421
422
423
424
@staticmethod
def get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]:
    stack = [body_id]
    desc: list[int] = []
    while stack:
        cur = stack.pop()
        for c in children_lists[cur]:
            desc.append(int(c))
            stack.append(c)
    return desc
get_direct_children staticmethod 
get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
426
427
428
@staticmethod
def get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]:
    return list(children_lists[body_id])
get_friction 
get_friction()
Source code in molmo_spaces/env/data_views.py 
381
382
def get_friction(self):
    return [self.mj_model.geom(geom_id).friction for geom_id in self.geom_ids]
get_geom_infos 
get_geom_infos(include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]

Get geom information for this object.

Parameters:

Name Type Description Default
include_descendants bool

If True, includes geoms from all descendant bodies. If False, only includes geoms directly attached to this object's body.

 True
max_geoms int | None

Maximum number of geoms to return. If None, returns all.

 2048

Returns:

Type Description
list[dict[str, object]]

List of dicts with geom info: id, name, position, size, type, type_name
Source code in molmo_spaces/env/data_views.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def get_geom_infos(
    self, include_descendants: bool = True, max_geoms: int | None = 2048
) -> list[dict[str, object]]:
    """Get geom information for this object.

    Args:
        include_descendants: If True, includes geoms from all descendant bodies.
            If False, only includes geoms directly attached to this object's body.
        max_geoms: Maximum number of geoms to return. If None, returns all.

    Returns:
        List of dicts with geom info: id, name, position, size, type, type_name
    """
    # Build set of body IDs to include
    if include_descendants:
        body_ids = set(self.body_ids)
    else:
        body_ids = {self.object_id}

    # Find geoms belonging to these bodies
    geoms: list[dict[str, object]] = []
    count = 0
    for geom_id in range(self.mj_model.ngeom):
        if max_geoms is not None and count >= max_geoms:
            break
        if int(self.mj_model.geom_bodyid[geom_id]) in body_ids:
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id))
            pos = self.mj_data.geom_xpos[geom_id].copy()
            size = self.mj_model.geom_size[geom_id].copy()
            gtype = int(self.mj_model.geom_type[geom_id])
            geoms.append(
                {
                    "id": int(geom_id),
                    "name": name,
                    "position": pos,
                    "size": size,
                    "type": gtype,
                    "type_name": MlSpacesObject.get_geom_type_name(gtype),
                }
            )
            count += 1

    return geoms
get_geom_type_name staticmethod 
get_geom_type_name(geom_type: int) -> str
Source code in molmo_spaces/env/data_views.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
@staticmethod
def get_geom_type_name(geom_type: int) -> str:
    names = {
        0: "PLANE",
        1: "HFIELD",
        2: "SPHERE",
        3: "CAPSULE",
        4: "ELLIPSOID",
        5: "CYLINDER",
        6: "BOX",
        7: "MESH",
    }
    return names.get(int(geom_type), f"UNKNOWN({geom_type})")
get_joint_anchor_position 
get_joint_anchor_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
550
551
552
553
554
555
556
557
558
559
def get_joint_anchor_position(self, i: int) -> np.ndarray:
    # Get the position in world frame of the joint anchor point
    joint_id = self.joint_ids[i]
    body_id = self.mj_model.jnt_bodyid[joint_id]
    local_anchor = self.mj_model.jnt_pos[joint_id]
    body_pos = self.mj_data.xpos[body_id]
    body_quat = self.mj_data.xquat[body_id]
    body_rot = R.from_quat(body_quat, scalar_first=True).as_matrix()
    world_anchor = body_pos + body_rot @ local_anchor
    return world_anchor
get_joint_armature 
get_joint_armature()
Source code in molmo_spaces/env/data_views.py 
570
571
def get_joint_armature(self):
    return [self.mj_model.joint(joint_id).armature for joint_id in self.joint_ids]
get_joint_axis 
get_joint_axis(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
526
527
def get_joint_axis(self, i: int) -> np.ndarray:
    return self.mj_model.joint(self.joint_ids[i]).axis
get_joint_body_orientation 
get_joint_body_orientation(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
545
546
547
548
def get_joint_body_orientation(self, i: int) -> np.ndarray:
    # orientation of the joint body not the root orientation of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    return self.mj_data.xmat[body_id].copy().reshape(3, 3)
get_joint_damping 
get_joint_damping()
Source code in molmo_spaces/env/data_views.py 
567
568
def get_joint_damping(self):
    return [self.mj_model.joint(joint_id).damping for joint_id in self.joint_ids]
get_joint_frictionloss 
get_joint_frictionloss()
Source code in molmo_spaces/env/data_views.py 
561
562
def get_joint_frictionloss(self):
    return [self.mj_model.joint(joint_id).frictionloss for joint_id in self.joint_ids]
get_joint_leaf_body_position 
get_joint_leaf_body_position(i: int) -> ndarray
Source code in molmo_spaces/env/data_views.py 
535
536
537
538
539
540
541
542
543
def get_joint_leaf_body_position(self, i: int) -> np.ndarray:
    # position of the joint body not the root center of the object
    body_id = self.mj_model.joint(self.joint_ids[i]).bodyid[0]
    # find child body of this joint body
    for child_body_id in range(self.mj_model.nbody):
        if self.mj_model.body(child_body_id).parentid[0] == body_id:
            body_id = child_body_id
            break
    return self.mj_data.xpos[body_id].copy()
get_joint_position 
get_joint_position(i: int) -> float
Source code in molmo_spaces/env/data_views.py 
520
521
def get_joint_position(self, i: int) -> float:
    return self.mj_data.qpos[self.get_joint_qpos_adr(i)].copy()
get_joint_qpos_adr 
get_joint_qpos_adr(i: int) -> int
Source code in molmo_spaces/env/data_views.py 
532
533
def get_joint_qpos_adr(self, i: int) -> int:
    return int(self.mj_model.joint(self.joint_ids[i]).qposadr[0])
get_joint_range 
get_joint_range(i: int) -> tuple[float, float]
Source code in molmo_spaces/env/data_views.py 
523
524
def get_joint_range(self, i: int) -> tuple[float, float]:
    return self.mj_model.joint(self.joint_ids[i]).range
get_joint_stiffness 
get_joint_stiffness()
Source code in molmo_spaces/env/data_views.py 
564
565
def get_joint_stiffness(self):
    return [self.mj_model.joint(joint_id).stiffness for joint_id in self.joint_ids]
get_joint_type 
get_joint_type(i: int) -> mjtJoint
Source code in molmo_spaces/env/data_views.py 
529
530
def get_joint_type(self, i: int) -> mujoco.mjtJoint:
    return self.mj_model.joint(self.joint_ids[i]).type
get_mass 
get_mass()
Source code in molmo_spaces/env/data_views.py 
384
385
def get_mass(self):
    return [self.mj_model.body_mass[body_id] for body_id in self.body_ids]
get_top_level_bodies staticmethod 
get_top_level_bodies(model: MjModel) -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/data_views.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
@staticmethod
def get_top_level_bodies(model: mujoco.MjModel) -> list[int]:
    """Return bodies whose parent is the world body."""
    # Identify world body id: the one with parent < 0
    try:
        world_id = next(
            b for b in range(model.nbody) if MlSpacesObject.body_parent_id(model, b) < 0
        )
    except StopIteration:
        world_id = 0
    tops: list[int] = []
    for b in range(model.nbody):
        try:
            pid = MlSpacesObject.body_parent_id(model, b)
            if pid == world_id:
                tops.append(b)
        except Exception as e:
            print(f"Error getting top-level bodies: {e}")
            continue
    return tops
is_child_name_of staticmethod 
is_child_name_of(parent_name: str, child_name: str) -> bool
Source code in molmo_spaces/env/data_views.py 
430
431
432
@staticmethod
def is_child_name_of(parent_name: str, child_name: str) -> bool:
    return child_name.startswith(parent_name + "_")
is_object_in_gripper 
is_object_in_gripper() -> None
Source code in molmo_spaces/env/data_views.py 
277
278
def is_object_in_gripper(self) -> None:
    pass
is_object_picked_up 
is_object_picked_up() -> None
Source code in molmo_spaces/env/data_views.py 
274
275
def is_object_picked_up(self) -> None:
    pass
set_joint_position 
set_joint_position(i: int, position: float) -> None
Source code in molmo_spaces/env/data_views.py 
573
574
575
576
def set_joint_position(self, i: int, position: float) -> None:
    if isinstance(position, np.ndarray):
        position = position.item()
    self.mj_data.qpos[self.get_joint_qpos_adr(i)] = position

MlSpacesBody 

MlSpacesBody(data: MjData, name: str)

Bases: MlSpacesObjectAbstract

Attributes:

Name Type Description
body_id int
mj_data
mj_model
name str
pose ndarray
position ndarray
quat ndarray
Source code in molmo_spaces/env/data_views.py 
56
57
def __init__(self, data: mujoco.MjData, name: str) -> None:
    super().__init__(data, name)
body_id cached property 
body_id: int
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
pose property writable 
pose: ndarray
position abstractmethod property writable 
position: ndarray
quat abstractmethod property writable 
quat: ndarray

MlSpacesCamera 

MlSpacesCamera(data: MjData, camera_name: str)

Bases: MlSpacesObjectAbstract

Attributes:

Name Type Description
camera_id int
fovy float
mj_data
mj_model
name str
pose ndarray
position ndarray
quat ndarray
Source code in molmo_spaces/env/data_views.py 
142
143
def __init__(self, data: mujoco.MjData, camera_name: str) -> None:
    super().__init__(data, camera_name)
camera_id cached property 
camera_id: int
fovy property writable 
fovy: float
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
pose property 
pose: ndarray
position property writable 
position: ndarray
quat property writable 
quat: ndarray

MlSpacesFreeJointBody 

MlSpacesFreeJointBody(data: MjData, body_name: str)

Bases: MlSpacesBody

Attributes:

Name Type Description
body_id int
mj_data
mj_model
name str
pose ndarray
position ndarray
quat ndarray
Source code in molmo_spaces/env/data_views.py 
65
66
67
68
69
70
71
72
73
def __init__(self, data: mujoco.MjData, body_name: str) -> None:
    super().__init__(data, body_name)
    self._jnt_id = self.mj_model.body_jntadr[self.body_id]
    self._qposadr = self.mj_model.jnt_qposadr[self._jnt_id]
    # TODO(all): this should be a ValueError
    assert (
        self._jnt_id != -1
        and self.mj_model.jnt_type[self._jnt_id] == mujoco.mjtJoint.mjJNT_FREE
    ), f"Body {body_name} does not have a free joint! {self._jnt_id}"
body_id cached property 
body_id: int
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
pose property writable 
pose: ndarray
position property writable 
position: ndarray
quat property writable 
quat: ndarray

MlSpacesImmovableBody 

MlSpacesImmovableBody(data: MjData, body_name: str)

Bases: MlSpacesBody

Attributes:

Name Type Description
body_id int
mj_data
mj_model
name str
pose ndarray
position ndarray
quat ndarray
Source code in molmo_spaces/env/data_views.py 
121
122
def __init__(self, data: mujoco.MjData, body_name: str) -> None:
    super().__init__(data, body_name)
body_id cached property 
body_id: int
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
pose property writable 
pose: ndarray
position property writable 
position: ndarray
quat property writable 
quat: ndarray

MlSpacesMocapBody 

MlSpacesMocapBody(data: MjData, body_name: str)

Bases: MlSpacesBody

Attributes:

Name Type Description
body_id int
mj_data
mj_model
mocap_id int
name str
pose ndarray
position ndarray
quat ndarray
Source code in molmo_spaces/env/data_views.py 
94
95
96
97
def __init__(self, data: mujoco.MjData, body_name: str) -> None:
    super().__init__(data, body_name)
    self._mocap_id = int(self.mj_model.body_mocapid[self.body_id])
    assert self._mocap_id != -1, f"Body {body_name} is not a mocap body!"
body_id cached property 
body_id: int
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
mocap_id property 
mocap_id: int
name property 
name: str
pose property writable 
pose: ndarray
position property writable 
position: ndarray
quat property writable 
quat: ndarray

MlSpacesObject 

MlSpacesObject(object_name: str, data: MjData)

Bases: MlSpacesBody

Class to represent scene object data

Methods:

Name Description
body_name2id
body_parent_id
build_children_lists
find_top_object_body_id
get_ancestors
get_descendants
get_direct_children
get_friction
get_geom_infos

Get geom information for this object.
get_geom_type_name
get_mass
get_top_level_bodies

Return bodies whose parent is the world body.
is_child_name_of
is_object_in_gripper
is_object_picked_up

Attributes:

Name Type Description
aabb_center ndarray
aabb_size ndarray
body_id int
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).
bvh_root
center_of_mass
geom_ids

Get all geom IDs belonging to this object (lazy, cached).
mj_data
mj_model
name str
object_id int
object_root_id
pose ndarray
position
quat
Source code in molmo_spaces/env/data_views.py 
209
210
211
212
213
214
215
216
217
218
219
220
221
222
def __init__(self, object_name: str, data: mujoco.MjData) -> None:
    super().__init__(data, object_name)

    # Note: _object_root_id, _geom_ids and _body_ids are no longer needed as backing stores
    # since @cached_property handles caching automatically

    # Local position of object center of mass in world coordinate
    self._center_of_mass_ref = self.mj_model.body_ipos[
        self.object_id
    ].copy()  # inertial frame position

    # Lazy axes-aligned bounding box (initial pose from model)
    self._bvhadr = None
    self._aabb_size: np.ndarray | None = None
aabb_center property 
aabb_center: ndarray
aabb_size property 
aabb_size: ndarray
body_id cached property 
body_id: int
body_ids cached property 
body_ids

Get all body IDs belonging to this object including descendants (lazy, cached).

bvh_root property 
bvh_root
center_of_mass property 
center_of_mass
geom_ids cached property 
geom_ids

Get all geom IDs belonging to this object (lazy, cached).

mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
object_id property 
object_id: int
object_root_id cached property 
object_root_id
pose property writable 
pose: ndarray
position property 
position
quat property 
quat
body_name2id staticmethod 
body_name2id(model: MjModel) -> dict[str, int]
Source code in molmo_spaces/env/data_views.py 
467
468
469
@staticmethod
def body_name2id(model: mujoco.MjModel) -> dict[str, int]:
    return {model.body(i).name: i for i in range(model.nbody)}
body_parent_id staticmethod 
body_parent_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
387
388
389
390
391
392
393
@staticmethod
def body_parent_id(model: mujoco.MjModel, body_id: int) -> int:
    try:
        return int(model.body_parentid[body_id])
    except Exception as e:
        print(f"Error getting body parent id: {e}")
        return int(model.body(body_id).parentid)
build_children_lists staticmethod 
build_children_lists(model: MjModel) -> list[list[int]]
Source code in molmo_spaces/env/data_views.py 
395
396
397
398
399
400
401
402
@staticmethod
def build_children_lists(model: mujoco.MjModel) -> list[list[int]]:
    children: list[list[int]] = [[] for _ in range(model.nbody)]
    for child_id in range(model.nbody):
        pid = MlSpacesObject.body_parent_id(model, child_id)
        if pid >= 0:
            children[pid].append(child_id)
    return children
find_top_object_body_id staticmethod 
find_top_object_body_id(model: MjModel, body_id: int) -> int
Source code in molmo_spaces/env/data_views.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
@staticmethod
def find_top_object_body_id(model: mujoco.MjModel, body_id: int) -> int:
    cur = body_id
    visited = set()
    while True:
        if cur in visited:
            break
        visited.add(cur)
        pid = MlSpacesObject.body_parent_id(model, cur)
        if pid < 0:
            break
        pname = model.body(pid).name
        cname = model.body(cur).name
        if MlSpacesObject.is_child_name_of(pname, cname):
            cur = pid
        else:
            break
    return int(cur)
get_ancestors staticmethod 
get_ancestors(model: MjModel, body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
404
405
406
407
408
409
410
411
412
413
@staticmethod
def get_ancestors(model: mujoco.MjModel, body_id: int) -> list[int]:
    ancestors: list[int] = []
    visited = {body_id}
    pid = MlSpacesObject.body_parent_id(model, body_id)
    while pid >= 0 and pid not in visited:
        ancestors.append(pid)
        visited.add(pid)
        pid = MlSpacesObject.body_parent_id(model, pid)
    return ancestors
get_descendants staticmethod 
get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
415
416
417
418
419
420
421
422
423
424
@staticmethod
def get_descendants(children_lists: list[list[int]], body_id: int) -> list[int]:
    stack = [body_id]
    desc: list[int] = []
    while stack:
        cur = stack.pop()
        for c in children_lists[cur]:
            desc.append(int(c))
            stack.append(c)
    return desc
get_direct_children staticmethod 
get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]
Source code in molmo_spaces/env/data_views.py 
426
427
428
@staticmethod
def get_direct_children(children_lists: list[list[int]], body_id: int) -> list[int]:
    return list(children_lists[body_id])
get_friction 
get_friction()
Source code in molmo_spaces/env/data_views.py 
381
382
def get_friction(self):
    return [self.mj_model.geom(geom_id).friction for geom_id in self.geom_ids]
get_geom_infos 
get_geom_infos(include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]

Get geom information for this object.

Parameters:

Name Type Description Default
include_descendants bool

If True, includes geoms from all descendant bodies. If False, only includes geoms directly attached to this object's body.

 True
max_geoms int | None

Maximum number of geoms to return. If None, returns all.

 2048

Returns:

Type Description
list[dict[str, object]]

List of dicts with geom info: id, name, position, size, type, type_name
Source code in molmo_spaces/env/data_views.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def get_geom_infos(
    self, include_descendants: bool = True, max_geoms: int | None = 2048
) -> list[dict[str, object]]:
    """Get geom information for this object.

    Args:
        include_descendants: If True, includes geoms from all descendant bodies.
            If False, only includes geoms directly attached to this object's body.
        max_geoms: Maximum number of geoms to return. If None, returns all.

    Returns:
        List of dicts with geom info: id, name, position, size, type, type_name
    """
    # Build set of body IDs to include
    if include_descendants:
        body_ids = set(self.body_ids)
    else:
        body_ids = {self.object_id}

    # Find geoms belonging to these bodies
    geoms: list[dict[str, object]] = []
    count = 0
    for geom_id in range(self.mj_model.ngeom):
        if max_geoms is not None and count >= max_geoms:
            break
        if int(self.mj_model.geom_bodyid[geom_id]) in body_ids:
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id))
            pos = self.mj_data.geom_xpos[geom_id].copy()
            size = self.mj_model.geom_size[geom_id].copy()
            gtype = int(self.mj_model.geom_type[geom_id])
            geoms.append(
                {
                    "id": int(geom_id),
                    "name": name,
                    "position": pos,
                    "size": size,
                    "type": gtype,
                    "type_name": MlSpacesObject.get_geom_type_name(gtype),
                }
            )
            count += 1

    return geoms
get_geom_type_name staticmethod 
get_geom_type_name(geom_type: int) -> str
Source code in molmo_spaces/env/data_views.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
@staticmethod
def get_geom_type_name(geom_type: int) -> str:
    names = {
        0: "PLANE",
        1: "HFIELD",
        2: "SPHERE",
        3: "CAPSULE",
        4: "ELLIPSOID",
        5: "CYLINDER",
        6: "BOX",
        7: "MESH",
    }
    return names.get(int(geom_type), f"UNKNOWN({geom_type})")
get_mass 
get_mass()
Source code in molmo_spaces/env/data_views.py 
384
385
def get_mass(self):
    return [self.mj_model.body_mass[body_id] for body_id in self.body_ids]
get_top_level_bodies staticmethod 
get_top_level_bodies(model: MjModel) -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/data_views.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
@staticmethod
def get_top_level_bodies(model: mujoco.MjModel) -> list[int]:
    """Return bodies whose parent is the world body."""
    # Identify world body id: the one with parent < 0
    try:
        world_id = next(
            b for b in range(model.nbody) if MlSpacesObject.body_parent_id(model, b) < 0
        )
    except StopIteration:
        world_id = 0
    tops: list[int] = []
    for b in range(model.nbody):
        try:
            pid = MlSpacesObject.body_parent_id(model, b)
            if pid == world_id:
                tops.append(b)
        except Exception as e:
            print(f"Error getting top-level bodies: {e}")
            continue
    return tops
is_child_name_of staticmethod 
is_child_name_of(parent_name: str, child_name: str) -> bool
Source code in molmo_spaces/env/data_views.py 
430
431
432
@staticmethod
def is_child_name_of(parent_name: str, child_name: str) -> bool:
    return child_name.startswith(parent_name + "_")
is_object_in_gripper 
is_object_in_gripper() -> None
Source code in molmo_spaces/env/data_views.py 
277
278
def is_object_in_gripper(self) -> None:
    pass
is_object_picked_up 
is_object_picked_up() -> None
Source code in molmo_spaces/env/data_views.py 
274
275
def is_object_picked_up(self) -> None:
    pass

MlSpacesObjectAbstract 

MlSpacesObjectAbstract(data: MjData, name: str)

Bases: ABC

Attributes:

Name Type Description
mj_data
mj_model
name str
pose ndarray
position ndarray
quat ndarray
Source code in molmo_spaces/env/data_views.py 
14
15
16
def __init__(self, data: mujoco.MjData, name: str) -> None:
    self.mj_data = data
    self._name = name
mj_data instance-attribute 
mj_data = data
mj_model cached property 
mj_model
name property 
name: str
pose property writable 
pose: ndarray
position abstractmethod property writable 
position: ndarray
quat abstractmethod property writable 
quat: ndarray

create_mlspaces_body 

create_mlspaces_body(data: MjData, body_name_or_id: str | int) -> MlSpacesBody
Source code in molmo_spaces/env/data_views.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
def create_mlspaces_body(data: mujoco.MjData, body_name_or_id: str | int) -> MlSpacesBody:
    if isinstance(body_name_or_id, int):
        body_id = body_name_or_id
        body_name = data.model.body(body_id).name
    else:
        body_name = body_name_or_id
        body_id = data.body(body_name).id

    if data.model.body_mocapid[body_id] != -1:
        return MlSpacesMocapBody(data, body_name)

    elif (
        data.model.body_jntadr[body_id] != -1
        and data.model.jnt_type[data.model.body_jntadr[body_id]] == mujoco.mjtJoint.mjJNT_FREE
    ):
        return MlSpacesFreeJointBody(data, body_name)

    else:
        return MlSpacesImmovableBody(data, body_name)

env

Classes:

Name Description
BaseMujocoEnv
CPUMujocoEnv

Attributes:

Name Type Description
HAS_FILAMENT bool
log

HAS_FILAMENT module-attribute 

HAS_FILAMENT: bool = getattr(mujoco, 'mjRENDERER', 'classic') == 'filament'

log module-attribute 

log = getLogger(__name__)

BaseMujocoEnv 

BaseMujocoEnv(exp_config: MlSpacesExpConfig, mj_model: MjModel = None)

Bases: ABC

Methods:

Name Description
is_loaded

Check if a scene is currently loaded.
reset
step

Attributes:

Name Type Description
config
current_batch_index int

Current batch index for accessing data and robots.
current_data MjData

Current MjData instance based on current_batch_index.
current_model MjModel

Current MjModel instance (always the same across batches).
current_model_path str

Current string xml path instance (always the same across batches).
current_robot Robot

Current robot instance based on current_batch_index.
current_scene_metadata dict

Current scene metadata instance (always the same across batches).
depth_frame ndarray

Get the latest depth frame.
mj_datas Sequence[MjData]
mj_model MjModel
n_batch int
object_managers list[ObjectManager]
rgb_frame ndarray

Get the latest RGB frame.
robots Sequence[Robot]
segmentation_frame ndarray

Get the latest segmentation frame.
Source code in molmo_spaces/env/env.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
def __init__(self, exp_config: "MlSpacesExpConfig", mj_model: MjModel = None) -> None:
    self._mj_model = mj_model
    self._current_batch_index = 0

    self.config = exp_config

    # Rendering state
    # TODO(anyone): can we remove these? too simple, we're on camera manager now.
    self._renderer = None
    self._rgb_frame = None
    self._depth_frame = None
    self._segmentation_frame = None
    self._camera_name = "camera"
config instance-attribute 
config = exp_config
current_batch_index property writable 
current_batch_index: int

Current batch index for accessing data and robots.

current_data property 
current_data: MjData

Current MjData instance based on current_batch_index.

current_model property 
current_model: MjModel

Current MjModel instance (always the same across batches).

current_model_path property 
current_model_path: str

Current string xml path instance (always the same across batches).

current_robot property 
current_robot: Robot

Current robot instance based on current_batch_index.

current_scene_metadata property 
current_scene_metadata: dict

Current scene metadata instance (always the same across batches).

depth_frame property 
depth_frame: ndarray

Get the latest depth frame.

mj_datas abstractmethod property 
mj_datas: Sequence[MjData]
mj_model property 
mj_model: MjModel
n_batch abstractmethod property 
n_batch: int
object_managers instance-attribute 
object_managers: list[ObjectManager]
rgb_frame property 
rgb_frame: ndarray

Get the latest RGB frame.

robots abstractmethod property 
robots: Sequence[Robot]
segmentation_frame property 
segmentation_frame: ndarray

Get the latest segmentation frame.

is_loaded 
is_loaded() -> bool

Check if a scene is currently loaded.

Source code in molmo_spaces/env/env.py 
55
56
57
def is_loaded(self) -> bool:
    """Check if a scene is currently loaded."""
    return self._mj_model is not None
reset abstractmethod 
reset(idxs: Collection[int] | None = None) -> None
Source code in molmo_spaces/env/env.py 
136
137
138
@abstractmethod
def reset(self, idxs: Collection[int] | None = None) -> None:
    raise NotImplementedError
step abstractmethod 
step(n_steps: int = 1) -> None
Source code in molmo_spaces/env/env.py 
140
141
142
@abstractmethod
def step(self, n_steps: int = 1) -> None:
    raise NotImplementedError

CPUMujocoEnv 

CPUMujocoEnv(exp_config: MlSpacesExpConfig, robot_factory: Callable[[MjData], Robot], mj_model: MjModel, mj_base_scene_path: str, parallelize: bool = True)

Bases: BaseMujocoEnv

Methods:

Name Description
__del__
check_camera_visibility_constraints

Check if all cameras with visibility constraints can see their target objects.
check_if_robot_collision_at_base_pose

Check if robot placement would result in collision with environment.
check_robot_collision_in_current_pose
check_visibility

Check visibility of one or more target objects from a specific camera.
cleanup_rendering

Clean up rendering resources.
close
get_camera_parameters

Get camera parameters for a specific camera.
get_robot_gripper_positions

Get current world positions of all robot grippers/end-effectors.
get_segmentation_mask_of_object

Get binary segmentation mask for a specific object from a camera view.
get_thormap

Get or create a cached occupancy map for robot placement.
get_visible_objects

Get list of visible objects for a specific camera (placeholder).
is_loaded

Check if a scene is currently loaded.
place_robot_near

Place robot near a target point or object with collision checking.
render_depth_frame

Renders a depth frame from the perspective of the specified camera.
render_rgb_frame

Renders an RGB frame from the perspective of the specified camera.
render_segmentation_frame

Renders a segmentation frame from the perspective of the specified camera.
reset
segmentation_fraction

Calculate visibility of a body in the segmentation image.
step

Attributes:

Name Type Description
camera_manager
config
current_batch_index int

Current batch index for accessing data and robots.
current_data MjData

Current MjData instance based on current_batch_index.
current_model MjModel

Current MjModel instance (always the same across batches).
current_model_path str

Current string xml path instance (always the same across batches).
current_robot Robot

Current robot instance based on current_batch_index.
current_scene_metadata dict

Current scene metadata instance (always the same across batches).
depth_frame ndarray

Get the latest depth frame.
mj_datas Sequence[MjData]
mj_model MjModel
n_batch int
object_managers
rgb_frame ndarray

Get the latest RGB frame.
robots Sequence[Robot]
segmentation_frame ndarray

Get the latest segmentation frame.
Source code in molmo_spaces/env/env.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
def __init__(
    self,
    exp_config: "MlSpacesExpConfig",
    robot_factory: Callable[[MjData], Robot],
    mj_model: MjModel,
    mj_base_scene_path: str,
    parallelize: bool = True,
) -> None:
    super().__init__(exp_config, mj_model)

    # Store configuration for scene loading
    self._robot_factory = robot_factory
    self._n_batch = exp_config.task_sampler_config.task_batch_size
    self._parallelize = parallelize

    # Initialize empty - will be populated when scene is loaded
    self._mj_datas = None
    self._robots = None
    self._executor = None
    self._mj_base_scene_path = None
    self._scene_metadata = None

    self.camera_manager = CameraManager()
    self._renderer: MjAbstractRenderer | None = None

    self.object_managers = []

    # Cached occupancy map for robot placement (expensive to create)
    self._cached_thormap = None
    self._cached_thormap_key = None  # (model_path, agent_radius, px_per_m)

    self._initialize_with_model(mj_model, mj_base_scene_path)
camera_manager instance-attribute 
camera_manager = CameraManager()
config instance-attribute 
config = exp_config
current_batch_index property writable 
current_batch_index: int

Current batch index for accessing data and robots.

current_data property 
current_data: MjData

Current MjData instance based on current_batch_index.

current_model property 
current_model: MjModel

Current MjModel instance (always the same across batches).

current_model_path property 
current_model_path: str

Current string xml path instance (always the same across batches).

current_robot property 
current_robot: Robot

Current robot instance based on current_batch_index.

current_scene_metadata property 
current_scene_metadata: dict

Current scene metadata instance (always the same across batches).

depth_frame property 
depth_frame: ndarray

Get the latest depth frame.

mj_datas property 
mj_datas: Sequence[MjData]
mj_model property 
mj_model: MjModel
n_batch property 
n_batch: int
object_managers instance-attribute 
object_managers = []
rgb_frame property 
rgb_frame: ndarray

Get the latest RGB frame.

robots property 
robots: Sequence[Robot]
segmentation_frame property 
segmentation_frame: ndarray

Get the latest segmentation frame.

__del__ 
__del__() -> None
Source code in molmo_spaces/env/env.py 
1055
1056
def __del__(self) -> None:
    self.close()
check_camera_visibility_constraints 
check_camera_visibility_constraints(visibility_resolver: Callable[[str], list[str]] | None = None, save_frames_dir: Path | str | None = None) -> tuple[bool, dict[str, dict[str, float]]]

Check if all cameras with visibility constraints can see their target objects.

This is used during robot placement to ensure that cameras (fixed exo, robot-mounted base, etc.) have good views of important objects (e.g., gripper, target objects).

Parameters:

Name Type Description Default
visibility_resolver Callable[[str], list[str]] | None

Optional callable that resolves special visibility keys like "gripper" to (possibly multiple) actual body names

 None
save_frames_dir Path | str | None

Optional directory to save RGB frames from each camera for debugging

 None

Returns:

Type Description
bool

Tuple of (all_satisfied, detailed_results)
dict[str, dict[str, float]]
  • all_satisfied: bool indicating if ALL constraints on ALL cameras are met
tuple[bool, dict[str, dict[str, float]]]
  • detailed_results: dict mapping camera_name -> {object_name: visibility_fraction}
Source code in molmo_spaces/env/env.py 
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
def check_camera_visibility_constraints(
    self,
    visibility_resolver: Callable[[str], list[str]] | None = None,
    save_frames_dir: Path | str | None = None,
) -> tuple[bool, dict[str, dict[str, float]]]:
    """
    Check if all cameras with visibility constraints can see their target objects.

    This is used during robot placement to ensure that cameras (fixed exo, robot-mounted base, etc.)
    have good views of important objects (e.g., gripper, target objects).

    Args:
        visibility_resolver: Optional callable that resolves special visibility keys
                           like "__gripper__" to (possibly multiple) actual body names
        save_frames_dir: Optional directory to save RGB frames from each camera for debugging

    Returns:
        Tuple of (all_satisfied, detailed_results)
        - all_satisfied: bool indicating if ALL constraints on ALL cameras are met
        - detailed_results: dict mapping camera_name -> {object_name: visibility_fraction}
    """
    all_satisfied = True
    detailed_results = {}

    # Prepare save directory if specified
    if save_frames_dir is not None:
        save_frames_dir = Path(save_frames_dir)
        visibility_frames_dir = save_frames_dir / "visibility_check_frames"
        visibility_frames_dir.mkdir(parents=True, exist_ok=True)

    # Iterate through all cameras in the registry
    for camera in self.camera_manager.registry:
        # Check if this camera has visibility constraints
        if (
            not hasattr(camera, "visibility_constraints")
            or camera.visibility_constraints is None
        ):
            continue

        camera_name = camera.name
        constraints = camera.visibility_constraints

        # Resolve special visibility keys
        resolved_constraints = {}
        for key, threshold in constraints.items():
            if key.startswith("__") and key.endswith("__"):
                # Special key - resolve via callback
                if visibility_resolver is not None:
                    resolved_keys = visibility_resolver(key)
                    if resolved_keys:
                        for resolved_key in resolved_keys:
                            resolved_constraints[resolved_key] = threshold
                    else:
                        log.warning(
                            f"[VISIBILITY CHECK] Could not resolve visibility key '{key}' for camera '{camera_name}'"
                        )
                else:
                    log.warning(
                        f"[VISIBILITY CHECK] No visibility resolver provided for key '{key}' in camera '{camera_name}'"
                    )
            else:
                # Regular body name
                resolved_constraints[key] = threshold

        if not resolved_constraints:
            # No valid constraints for this camera
            continue

        # Check visibility for all objects
        try:
            visibility_results = self.check_visibility(
                camera_name, *resolved_constraints.keys()
            )

            # Convert to dict if single object
            if not isinstance(visibility_results, dict):
                visibility_results = {list(resolved_constraints.keys())[0]: visibility_results}

            detailed_results[camera_name] = visibility_results

            # Save frame if directory specified
            if save_frames_dir is not None:
                try:
                    import time

                    from PIL import Image

                    rgb_frame = self.render_rgb_frame(camera_name)
                    # Convert from float [0,1] to uint8 if needed
                    if rgb_frame.dtype == np.float32 or rgb_frame.dtype == np.float64:
                        rgb_frame = (rgb_frame * 255).astype(np.uint8)
                    img = Image.fromarray(rgb_frame)
                    # Add timestamp to filename to avoid overwriting
                    timestamp = int(time.time() * 1000)
                    frame_path = visibility_frames_dir / f"{camera_name}_{timestamp}.png"
                    img.save(frame_path)
                    log.debug(f"[VISIBILITY CHECK] Saved frame to {frame_path}")
                except Exception as e:
                    log.warning(
                        f"[VISIBILITY CHECK] Failed to save frame for camera '{camera_name}': {e}"
                    )

            # Check if all constraints are satisfied for this camera
            for obj_name, threshold in resolved_constraints.items():
                actual_visibility = visibility_results.get(obj_name, 0.0)
                if actual_visibility < threshold:
                    all_satisfied = False
                    log.debug(
                        f"[VISIBILITY CHECK] Camera '{camera_name}': object '{obj_name}' "
                        f"visibility {actual_visibility:.5f} < threshold {threshold:.5f}"
                    )

        except Exception as e:
            log.warning(
                f"[VISIBILITY CHECK] Failed to check visibility for camera '{camera_name}': {e}"
            )
            all_satisfied = False

    return all_satisfied, detailed_results
check_if_robot_collision_at_base_pose 
check_if_robot_collision_at_base_pose(robot_view, robot_pose: ndarray, robot_namespace: str = 'robot_0/') -> bool

Check if robot placement would result in collision with environment.

Source code in molmo_spaces/env/env.py 
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
def check_if_robot_collision_at_base_pose(
    self, robot_view, robot_pose: np.ndarray, robot_namespace: str = "robot_0/"
) -> bool:
    """Check if robot placement would result in collision with environment."""
    # TODO: there has to be a better way to do this
    model = self.current_model
    data = self.current_data

    # Store current robot pose
    original_pose = robot_view.base.pose.copy()

    try:
        # Temporarily place robot at candidate position
        robot_view.base.pose = robot_pose
        mujoco.mj_forward(model, data)
        return self.check_robot_collision_in_current_pose(robot_namespace)

    finally:
        # Restore original robot pose
        robot_view.base.pose = original_pose
        mujoco.mj_forward(model, data)
check_robot_collision_in_current_pose 
check_robot_collision_in_current_pose(robot_namespace: str = 'robot_0/') -> bool
Source code in molmo_spaces/env/env.py 
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
def check_robot_collision_in_current_pose(self, robot_namespace: str = "robot_0/") -> bool:
    model = self.current_model
    data = self.current_data

    # Check for contacts
    contacts = data.contact
    collision_found = False

    for i in range(data.ncon):
        contact = contacts[i]
        if contact.dist > 0:  # Only consider actual contacts
            continue

        # Get body IDs for the contacting geometries
        body1_id = model.geom_bodyid[contact.geom1]
        body2_id = model.geom_bodyid[contact.geom2]

        # Get root body IDs to identify which system each contact belongs to
        root1_id = model.body_rootid[body1_id]
        root2_id = model.body_rootid[body2_id]

        # Check if this involves the robot by looking at body names
        body1_name = model.body(root1_id).name
        body2_name = model.body(root2_id).name

        # Check if either body belongs to the robot (using namespace)
        robot_involved = body1_name.startswith(robot_namespace) or body2_name.startswith(
            robot_namespace
        )

        if robot_involved:
            # Get the non-robot body name
            other_body_name = (
                body2_name if body1_name.startswith(robot_namespace) else body1_name
            )

            actual_body1_name = model.body(body1_id).name
            actual_body2_name = model.body(body2_id).name
            log.debug(
                f"Collision is between {actual_body1_name} and {actual_body2_name} (roots: {body1_name}, {body2_name})"
            )

            # Allow floor contacts but reject walls/obstacles
            if "floor" not in other_body_name.lower() and body1_name != body2_name:
                collision_found = True
                break

    return collision_found
check_visibility 
check_visibility(camera_name: str, *target_objects) -> float | dict[str, float]

Check visibility of one or more target objects from a specific camera.

Parameters:

Name Type Description Default
camera_name str

Name of camera in the registry

 required
*target_objects

Variable number of object/body names to check visibility for

 ()

Returns:

Name Type Description
float float | dict[str, float]

If single object provided, returns visibility fraction (0.0 to 1.0)
dict float | dict[str, float]

If multiple objects provided, returns mapping of object names to visibility fractions
Source code in molmo_spaces/env/env.py 
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
def check_visibility(self, camera_name: str, *target_objects) -> float | dict[str, float]:
    """
    Check visibility of one or more target objects from a specific camera.

    Args:
        camera_name: Name of camera in the registry
        *target_objects: Variable number of object/body names to check visibility for

    Returns:
        float: If single object provided, returns visibility fraction (0.0 to 1.0)
        dict: If multiple objects provided, returns mapping of object names to visibility fractions
    """
    if camera_name not in self.camera_manager.registry:
        raise KeyError(f"Camera '{camera_name}' not found in registry.")

    if len(target_objects) == 0:
        raise ValueError("At least one target object must be provided")

    try:
        # Render segmentation frame once for all objects
        seg_frame = self.render_segmentation_frame(camera_name)

        results = {}
        # Check visibility for each target object
        for obj_name in target_objects:
            try:
                visibility = self.segmentation_fraction(seg_frame, obj_name)
                results[obj_name] = visibility
            except ValueError:
                results[obj_name] = 0.0

        # Return single float if only one object, dict if multiple
        if len(target_objects) == 1:
            return results[target_objects[0]]
        else:
            return results

    except Exception as e:
        # Return 0 visibility for all objects if camera render fails
        log.warning(
            f"[VISIBILITY CHECK] Failed to check visibility for camera {camera_name}: {e}"
        )
        if len(target_objects) == 1:
            return 0.0
        else:
            return {obj_name: 0.0 for obj_name in target_objects}
cleanup_rendering 
cleanup_rendering() -> None

Clean up rendering resources.

Source code in molmo_spaces/env/env.py 
365
366
367
368
369
def cleanup_rendering(self) -> None:
    """Clean up rendering resources."""
    if self._renderer:
        self._renderer.close()
        self._renderer = None
close 
close() -> None
Source code in molmo_spaces/env/env.py 
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
def close(self) -> None:
    if self._executor is not None:
        self._executor.shutdown(wait=True)

    # clean up object managers
    for om in self.object_managers:
        om.clear()
    self.object_managers.clear()

    # Clean up camera renderers
    self.cleanup_rendering()

    # add garbage collection to ensure all resources are released
    gc.collect()
get_camera_parameters 
get_camera_parameters(camera_name: str) -> dict

Get camera parameters for a specific camera.

Source code in molmo_spaces/env/env.py 
327
328
329
330
331
332
def get_camera_parameters(self, camera_name: str) -> dict:
    """Get camera parameters for a specific camera."""
    if camera_name not in self.camera_manager.registry:
        raise KeyError(f"Camera '{camera_name}' not found.")
    camera = self.camera_manager.registry[camera_name]
    return {"pos": camera.pos, "forward": camera.forward, "up": camera.up}
get_robot_gripper_positions 
get_robot_gripper_positions(robot_index: int = 0) -> dict[str, ndarray]

Get current world positions of all robot grippers/end-effectors.

It's particularly useful for camera positioning to ensure good gripper visibility.

Parameters:

Name Type Description Default
robot_index int

Index of robot in batch (default 0)

 0

Returns:

Type Description
dict[str, ndarray]

Dict mapping gripper names to their world positions as numpy arrays.
dict[str, ndarray]

Returns empty dict if the robot does not have any grippers.
Source code in molmo_spaces/env/env.py 
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
def get_robot_gripper_positions(self, robot_index: int = 0) -> dict[str, np.ndarray]:
    """Get current world positions of all robot grippers/end-effectors.

    It's particularly useful for camera positioning to ensure good gripper visibility.

    Args:
        robot_index: Index of robot in batch (default 0)

    Returns:
        Dict mapping gripper names to their world positions as numpy arrays.
        Returns empty dict if the robot does not have any grippers.
    """
    assert 0 <= robot_index < self.n_batch, (
        f"Robot index {robot_index} out of range [0, {self.n_batch})"
    )

    robot = self.robots[robot_index]
    robot_view = robot.robot_view

    gripper_positions = {}
    for gripper_name in robot_view.get_gripper_movegroup_ids():
        mg = robot_view.get_move_group(gripper_name)
        gripper_positions[gripper_name] = mg.leaf_frame_to_world[:3, 3]

    return gripper_positions
get_segmentation_mask_of_object 
get_segmentation_mask_of_object(object_name: str, camera_name: str, batch_index: int = 0) -> ndarray | None

Get binary segmentation mask for a specific object from a camera view.

Parameters:

Name Type Description Default
object_name str

Name of the object body to segment

 required
camera_name str

Name of the camera to render from

 required
batch_index int

Batch index for the environment

 0

Returns:

Type Description
ndarray | None

Binary mask (HxW) where True indicates object pixels, or None if object not found
Source code in molmo_spaces/env/env.py 
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
def get_segmentation_mask_of_object(
    self, object_name: str, camera_name: str, batch_index: int = 0
) -> np.ndarray | None:
    """Get binary segmentation mask for a specific object from a camera view.

    Args:
        object_name: Name of the object body to segment
        camera_name: Name of the camera to render from
        batch_index: Batch index for the environment

    Returns:
        Binary mask (HxW) where True indicates object pixels, or None if object not found
    """
    # Set current batch index for rendering
    prev_batch_index = self.current_batch_index
    try:
        self.current_batch_index = batch_index

        # Get the body ID for the object
        model = self.current_model
        try:
            body_id = model.body(object_name).id
        except KeyError:
            log.warning(f"Object '{object_name}' not found in model")
            return None

        # Render segmentation frame for the camera
        seg_frame = self.render_segmentation_frame(camera_name)

        # Extract binary mask for this specific object using get_geom_seg_mask
        object_mask = get_geom_seg_mask(model, seg_frame[..., :2], body_id)

        return object_mask.astype(bool)

    finally:
        # Restore previous batch index
        self.current_batch_index = prev_batch_index
get_thormap 
get_thormap(agent_radius: float = 0.35, px_per_m: int = 200) -> ProcTHORMap | iTHORMap

Get or create a cached occupancy map for robot placement.

The map is cached per scene and parameters to avoid expensive recreation. Creating the map involves re-parsing XML, re-compiling MuJoCo model, and rendering multiple depth views - typically taking 10-30 seconds.

Parameters:

Name Type Description Default
agent_radius float

Safety radius for occupancy dilation (default 0.35m)

 0.35
px_per_m int

Pixels per meter resolution (default 200)

 200

Returns:

Type Description
ProcTHORMap | iTHORMap

ProcTHORMap or iTHORMap depending on scene type
Source code in molmo_spaces/env/env.py 
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
def get_thormap(
    self, agent_radius: float = 0.35, px_per_m: int = 200
) -> "ProcTHORMap | iTHORMap":
    """
    Get or create a cached occupancy map for robot placement.

    The map is cached per scene and parameters to avoid expensive recreation.
    Creating the map involves re-parsing XML, re-compiling MuJoCo model, and
    rendering multiple depth views - typically taking 10-30 seconds.

    Args:
        agent_radius: Safety radius for occupancy dilation (default 0.35m)
        px_per_m: Pixels per meter resolution (default 200)

    Returns:
        ProcTHORMap or iTHORMap depending on scene type
    """
    cache_key = (self.current_model_path, agent_radius, px_per_m)

    # Return cached map if available and matches parameters
    if self._cached_thormap is not None and self._cached_thormap_key == cache_key:
        log.debug("[THORMAP] Using cached occupancy map")
        return self._cached_thormap

    # Create new map
    log.info(
        f"[THORMAP] Creating occupancy map (agent_radius={agent_radius}, px_per_m={px_per_m})"
    )

    if "ithor" in self.current_model_path:
        model_path = Path(self.current_model_path.replace("_ceiling", ""))
        precomputed_map = model_path.parent / f"{model_path.stem}_map.png"
        if precomputed_map.is_file():
            thormap = iTHORMap.load(
                path=precomputed_map.as_posix(),
                agent_radius=agent_radius,
            )
        else:
            thormap = iTHORMap.from_mj_model_path(
                model_path=self.current_model_path,
                agent_radius=agent_radius,
                px_per_m=px_per_m,
                device_id=None,
                use_filament=HAS_FILAMENT,
            )
    elif "procthor" in self.current_model_path or "holodeck" in self.current_model_path:
        model_path = Path(self.current_model_path.replace("_ceiling", ""))
        precomputed_map = model_path.parent / f"{model_path.stem}_map.png"
        if precomputed_map.is_file():
            thormap = ProcTHORMap.load(
                path=precomputed_map.as_posix(),
                agent_radius=agent_radius,
            )
        else:
            thormap = ProcTHORMap.from_mj_model_path(
                model_path=self.current_model_path,
                px_per_m=px_per_m,
                agent_radius=agent_radius,
                device_id=None,
                use_filament=HAS_FILAMENT,
            )
    else:
        raise ValueError(f"Unknown scene type: {self.current_model_path}")

    # Cache the map
    self._cached_thormap = thormap
    self._cached_thormap_key = cache_key
    log.info("[THORMAP] Occupancy map created and cached")

    return thormap
get_visible_objects 
get_visible_objects(camera_name: str) -> list

Get list of visible objects for a specific camera (placeholder).

Source code in molmo_spaces/env/env.py 
334
335
336
337
def get_visible_objects(self, camera_name: str) -> list:
    """Get list of visible objects for a specific camera (placeholder)."""
    # This would require proper visibility calculation using segmentation - TODO
    return []
is_loaded 
is_loaded() -> bool

Check if a scene is currently loaded.

Source code in molmo_spaces/env/env.py 
55
56
57
def is_loaded(self) -> bool:
    """Check if a scene is currently loaded."""
    return self._mj_model is not None
place_robot_near 
place_robot_near(robot_view, target, max_tries: int = 10, sampling_radius_range: tuple[float, float] = (0.0, 1.0), robot_safety_radius: float = 0.35, preserve_z: float = None, face_target: bool = True, check_camera_visibility: bool = False, visibility_resolver=None, excluded_positions: list[ndarray] | None = None, exclusion_threshold: float | None = None, save_visibility_frames_dir: Path | str | None = None) -> bool

Place robot near a target point or object with collision checking.

Parameters:

Name Type Description Default
robot_view

Robot view object to position

 required
target

Either a 3D point (np.ndarray) or an Object instance or object name (str)

 required
max_tries int

Maximum number of placement attempts

 10
sampling_radius_range tuple[float, float]

Radius range around target to sample robot positions. For picking up objects, the min radius is 0.0.

 (0.0, 1.0)
robot_safety_radius float

Safety radius for occupancy map collision checking

 0.35
preserve_z float

Z height to preserve for robot (if None, uses current robot Z)

 None
face_target bool

Whether to orient robot to face the target

 True
check_camera_visibility bool

Whether to check fixed camera visibility constraints after placing robot

 False
visibility_resolver

Optional callable(key: str) -> str for resolving special visibility keys

 None
excluded_positions list[ndarray] | None

List of positions to avoid (e.g. previously used positions)

 None
exclusion_threshold float | None

Minimum distance from any excluded position

 None
save_visibility_frames_dir Path | str | None

Optional directory to save camera frames during visibility check

 None

Returns: bool: True if placement was successful, False otherwise

Source code in molmo_spaces/env/env.py 
 761
 762
 763
 764
 765
 766
 767
 768
 769
 770
 771
 772
 773
 774
 775
 776
 777
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
def place_robot_near(
    self,
    robot_view,
    target,
    max_tries: int = 10,
    sampling_radius_range: tuple[float, float] = (0.0, 1.0),
    robot_safety_radius: float = 0.35,
    preserve_z: float = None,
    face_target: bool = True,
    check_camera_visibility: bool = False,
    visibility_resolver=None,
    excluded_positions: list[np.ndarray] | None = None,
    exclusion_threshold: float | None = None,
    save_visibility_frames_dir: Path | str | None = None,
) -> bool:
    """
    Place robot near a target point or object with collision checking.

    Args:
        robot_view: Robot view object to position
        target: Either a 3D point (np.ndarray) or an Object instance or object name (str)
        max_tries: Maximum number of placement attempts
        sampling_radius_range: Radius range around target to sample robot positions. For picking up objects, the min radius is 0.0.
        robot_safety_radius: Safety radius for occupancy map collision checking
        preserve_z: Z height to preserve for robot (if None, uses current robot Z)
        face_target: Whether to orient robot to face the target
        check_camera_visibility: Whether to check fixed camera visibility constraints after placing robot
        visibility_resolver: Optional callable(key: str) -> str for resolving special visibility keys
        excluded_positions: List of positions to avoid (e.g. previously used positions)
        exclusion_threshold: Minimum distance from any excluded position
        save_visibility_frames_dir: Optional directory to save camera frames during visibility check
    Returns:
        bool: True if placement was successful, False otherwise
    """
    if exclusion_threshold is None:
        exclusion_threshold = (
            self.config.task_sampler_config.robot_placement_exclusion_threshold
        )

    log.debug(
        f"[PLACE_ROBOT_NEAR] Starting robot placement near target (max_tries={max_tries})"
    )

    # Extract target position
    if isinstance(target, np.ndarray):
        if target.shape != (3,):
            raise ValueError("Target point must be a 3D numpy array")
        target_pos = target
        log.debug(
            f"[PLACE_ROBOT_NEAR] Target point: ({target_pos[0]:.3f}, {target_pos[1]:.3f}, {target_pos[2]:.3f})"
        )
    elif isinstance(target, MlSpacesArticulationObject):
        target_pos = target.get_joint_leaf_body_position(0)

    elif isinstance(target, MlSpacesObject):
        target_pos = target.position
        log.debug(
            f"[PLACE_ROBOT_NEAR] Target object '{target.name}': ({target_pos[0]:.3f}, {target_pos[1]:.3f}, {target_pos[2]:.3f})"
        )
    elif isinstance(target, str):
        # Assume it's an object name
        target_obj = create_mlspaces_body(self.current_data, target)
        target_pos = target_obj.position
        log.debug(
            f"[PLACE_ROBOT_NEAR] Target object '{target}': ({target_pos[0]:.3f}, {target_pos[1]:.3f}, {target_pos[2]:.3f})"
        )
    else:
        raise ValueError(
            "Target must be a 3D point (np.ndarray), Object instance, or object name (str)"
        )

    # Get robot Z height to preserve
    initial_robot_z = preserve_z if preserve_z is not None else robot_view.base.pose[2, 3]
    log.debug(f"[PLACE_ROBOT_NEAR] Robot Z height to preserve: {initial_robot_z:.6f}m")
    log.debug(
        f"[PLACE_ROBOT_NEAR] Sampling radius range: {sampling_radius_range[0]:.3f}m - {sampling_radius_range[1]:.3f}m"
    )

    # Track timing breakdown for profiling
    import time

    place_start_time = time.perf_counter()
    map_time_ms = 0.0  # Will be set after map creation
    attempts_made = 0

    # Try occupancy map approach first (using cached map for performance)
    try:
        map_start_time = time.perf_counter()
        thormap = self.get_thormap(agent_radius=robot_safety_radius, px_per_m=200)
        map_time_ms = (time.perf_counter() - map_start_time) * 1000

        # Get free points within sampling radius of target
        free_points = thormap.get_free_points()
        target_dist = np.linalg.norm(free_points[:, :2] - target_pos[:2], axis=1)
        # Filter points within the radius range [min, max]
        valid_mask = (target_dist > sampling_radius_range[0]) & (
            target_dist < sampling_radius_range[1]
        )
        valid_points = free_points[valid_mask]

        if len(valid_points) > 0:
            log.debug(
                f"[PLACE_ROBOT_NEAR] Found {len(valid_points)} free points within sampling radius"
            )

            # PHASE 1: Find all collision-free candidate poses (no visibility checks yet)
            # This avoids expensive rendering on every attempt
            collision_free_poses = []
            for attempt in range(max_tries):
                attempts_made = attempt + 1
                # Sample a random point from valid points using the radius range
                sampled_point = sample_around_point(
                    thormap, target_pos[:2], sampling_radius_range
                )

                robot_base_pos = np.array([sampled_point[0], sampled_point[1], initial_robot_z])

                # Check against excluded positions
                if excluded_positions:
                    is_excluded = np.any(
                        np.linalg.norm(
                            np.stack(excluded_positions)[:, :2] - robot_base_pos[None, :2],
                            axis=-1,
                        )
                        < exclusion_threshold
                    )
                    if is_excluded:
                        if attempt < 5:
                            log.debug(
                                f"[PLACE_ROBOT_NEAR] Attempt {attempt + 1}: Position excluded, trying another..."
                            )
                        continue  # Skip this point

                # Reject pose if target is a Door and point is inside swing arc
                if isinstance(target, Door):
                    if target.is_point_in_swing_arc(
                        robot_base_pos, safety_margin=robot_safety_radius
                    ):
                        if attempt < 5:
                            log.debug(
                                f"[PLACE_ROBOT_NEAR] Attempt {attempt + 1}: Point in door swing arc, rejecting..."
                            )
                        continue  # Skip this point

                # Calculate robot orientation
                if face_target:
                    # Orient robot to face the target
                    xy_vec_robot_to_target = target_pos[:2] - robot_base_pos[:2]
                    if np.linalg.norm(xy_vec_robot_to_target) > 1e-6:  # Avoid division by zero
                        robot_base_yaw = np.arctan2(
                            xy_vec_robot_to_target[1], xy_vec_robot_to_target[0]
                        )
                    else:
                        robot_base_yaw = (
                            0.0  # Default orientation if target is at same XY position
                        )
                    # NOTE(yejin): is this robot-specific logic okay?
                    if "rum" in self.config.robot_config.robot_cls.__name__.lower():
                        # Orient robot pitch to face the target
                        robot_base_pitch = -np.arctan2(
                            target_pos[2] - robot_base_pos[2],
                            np.linalg.norm(xy_vec_robot_to_target),
                        )
                    else:
                        robot_base_pitch = 0.0  # Default pitch
                else:
                    robot_base_yaw = 0.0  # Default orientation
                    robot_base_pitch = 0.0  # Default pitch

                # Apply randomization to yaw
                randomization_range = (
                    self.config.task_sampler_config.robot_placement_rotation_range_rad
                )
                if randomization_range > 0:
                    robot_base_yaw += np.random.uniform(
                        -randomization_range, randomization_range
                    )

                robot_base_quat = R.from_euler(
                    "xyz", [0, robot_base_pitch, robot_base_yaw], degrees=False
                ).as_quat(scalar_first=True)

                # Create robot pose matrix
                robot_pose = np.eye(4)
                robot_pose[:3, 3] = robot_base_pos
                robot_pose[:3, :3] = R.from_quat(robot_base_quat, scalar_first=True).as_matrix()

                log.debug(
                    f"[PLACE_ROBOT_NEAR] Attempt {attempt + 1}: pos={robot_base_pos} yaw={np.degrees(robot_base_yaw):.1f}deg"
                )

                # Check for collisions
                if not self.check_if_robot_collision_at_base_pose(
                    robot_view, robot_pose, "robot_0/"
                ):
                    # Valid collision-free placement found - add to candidates
                    collision_free_poses.append((robot_pose, robot_base_pos, robot_base_yaw))
                    log.debug(
                        f"[PLACE_ROBOT_NEAR] Found collision-free pose #{len(collision_free_poses)}"
                    )
                    # If not checking visibility, we can return immediately with first valid pose
                    if not check_camera_visibility:
                        break
                    # Otherwise, collect a few candidates for visibility checking
                    # Stop early once we have enough candidates to avoid unnecessary collision checks
                    if len(collision_free_poses) >= self.config.collision_free_pose_limit:
                        break
                elif attempt < 5:
                    log.debug(
                        "[PLACE_ROBOT_NEAR]   Collision detected, trying another point..."
                    )

            # PHASE 2: Check visibility only on collision-free candidates (much fewer renders)
            for pose_idx, (robot_pose, robot_base_pos, robot_base_yaw) in enumerate(
                collision_free_poses
            ):
                robot_view.base.pose = robot_pose
                mujoco.mj_forward(self.current_model, self.current_data)

                # Check visibility constraints if requested
                if check_camera_visibility:
                    # Update camera poses (important for robot-mounted cameras)
                    self.camera_manager.registry.update_all_cameras(self)

                    visibility_satisfied, visibility_results = (
                        self.check_camera_visibility_constraints(
                            visibility_resolver=visibility_resolver
                        )
                    )

                    if not visibility_satisfied:
                        log.debug(
                            f"[PLACE_ROBOT_NEAR] Candidate {pose_idx + 1}/{len(collision_free_poses)}: Visibility constraints not met"
                        )
                        continue  # Try next collision-free candidate
                    else:
                        log.debug(
                            f"[PLACE_ROBOT_NEAR] Visibility constraints satisfied: {visibility_results}"
                        )

                # Success!
                total_time_ms = (time.perf_counter() - place_start_time) * 1000
                retry_time_ms = total_time_ms - map_time_ms
                log.info(
                    f"[PLACE_ROBOT_NEAR] Success after {attempts_made} samples, {len(collision_free_poses)} collision-free, {pose_idx + 1} visibility checks | "
                    f"map={map_time_ms:.1f}ms, retries={retry_time_ms:.1f}ms, total={total_time_ms:.1f}ms"
                )
                log.debug(
                    f"[PLACE_ROBOT_NEAR] Final position: ({robot_base_pos[0]:.3f}, {robot_base_pos[1]:.3f}, {robot_base_pos[2]:.3f})"
                )
                log.debug(
                    f"[PLACE_ROBOT_NEAR] Final orientation: {np.degrees(robot_base_yaw):.1f} deg"
                )
                distance_to_target = np.linalg.norm(robot_base_pos[:2] - target_pos[:2])
                log.debug(f"[PLACE_ROBOT_NEAR] Distance to target: {distance_to_target:.3f}m")
                return True

            if len(collision_free_poses) > 0:
                log.debug(
                    f"[PLACE_ROBOT_NEAR] Found {len(collision_free_poses)} collision-free poses but none satisfied visibility"
                )
        else:
            log.debug(
                "[PLACE_ROBOT_NEAR] No free points found within sampling radius, trying fallback..."
            )

    except Exception as e:
        log.exception(e)
        log.debug(f"[PLACE_ROBOT_NEAR] ❌ Occupancy map failed: {e}, failed to place robot...")

    # Log timing even on failure
    total_time_ms = (time.perf_counter() - place_start_time) * 1000
    retry_time_ms = total_time_ms - map_time_ms
    log.info(
        f"[PLACE_ROBOT_NEAR] ❌ Failed after {attempts_made} attempts | "
        f"map={map_time_ms:.1f}ms, retries={retry_time_ms:.1f}ms, total={total_time_ms:.1f}ms"
    )
    return False
render_depth_frame 
render_depth_frame(camera_name: str) -> ndarray

Renders a depth frame from the perspective of the specified camera.

Returns raw metric depth values in meters as float32 array. Depth encoding to RGB happens at save time for video storage.

Returns:

Type Description
ndarray

np.ndarray: (H, W) float32 array of depth values in meters
Source code in molmo_spaces/env/env.py 
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
def render_depth_frame(self, camera_name: str) -> np.ndarray:
    """Renders a depth frame from the perspective of the specified camera.

    Returns raw metric depth values in meters as float32 array.
    Depth encoding to RGB happens at save time for video storage.

    Returns:
        np.ndarray: (H, W) float32 array of depth values in meters
    """
    if camera_name not in self.camera_manager.registry:
        raise KeyError(f"Camera '{camera_name}' not found in registry.")

    camera = self.camera_manager.registry[camera_name]
    depth_frame = self._render_frame(
        camera.pos, camera.forward, camera.up, camera.fov, depth=True
    )

    # Return raw depth in meters (encoding to RGB happens at save time)
    return depth_frame.astype(np.float32)
render_rgb_frame 
render_rgb_frame(camera_name: str) -> ndarray

Renders an RGB frame from the perspective of the specified camera.

Source code in molmo_spaces/env/env.py 
287
288
289
290
291
292
293
294
295
def render_rgb_frame(self, camera_name: str) -> np.ndarray:
    """Renders an RGB frame from the perspective of the specified camera."""
    if camera_name not in self.camera_manager.registry:
        raise KeyError(f"Camera '{camera_name}' not found in registry.")

    camera = self.camera_manager.registry[camera_name]
    return self._render_frame(
        camera.pos, camera.forward, camera.up, camera.fov, segmentation=False
    )
render_segmentation_frame 
render_segmentation_frame(camera_name: str) -> ndarray

Renders a segmentation frame from the perspective of the specified camera.

Source code in molmo_spaces/env/env.py 
317
318
319
320
321
322
323
324
325
def render_segmentation_frame(self, camera_name: str) -> np.ndarray:
    """Renders a segmentation frame from the perspective of the specified camera."""
    if camera_name not in self.camera_manager.registry:
        raise KeyError(f"Camera '{camera_name}' not found in registry.")

    camera = self.camera_manager.registry[camera_name]
    return self._render_frame(
        camera.pos, camera.forward, camera.up, camera.fov, segmentation=True
    )
reset 
reset(idxs: Collection[int] | None = None) -> None
Source code in molmo_spaces/env/env.py 
375
376
377
378
379
380
381
382
383
384
def reset(self, idxs: Collection[int] | None = None) -> None:
    if idxs is None:
        idxs = range(self.n_batch)
    if self._executor is not None and len(idxs) > 1:
        futures = [self._executor.submit(self._reset_single, idx) for idx in idxs]
        for future in as_completed(futures):
            future.result()
    else:
        for idx in idxs:
            self._reset_single(idx)
segmentation_fraction 
segmentation_fraction(seg: ndarray, body_name_or_id: str | int) -> float

Calculate visibility of a body in the segmentation image.

Source code in molmo_spaces/env/env.py 
404
405
406
407
408
409
410
411
412
def segmentation_fraction(self, seg: np.ndarray, body_name_or_id: str | int) -> float:
    """Calculate visibility of a body in the segmentation image."""
    model = self.current_model
    if isinstance(body_name_or_id, str):
        body_id = model.body(body_name_or_id).id
    else:
        body_id = body_name_or_id

    return np.mean(get_geom_seg_mask(model, seg[..., :2], body_id)).item()
step 
step(n_steps: int = 1) -> None
Source code in molmo_spaces/env/env.py 
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
def step(self, n_steps: int = 1) -> None:
    if self._executor is not None:
        futures = [
            self._executor.submit(mujoco.mj_step, self._mj_model, mj_data, n_steps)
            for mj_data in self._mj_datas
        ]
        for future in as_completed(futures):
            future.result()
    else:
        for mj_data in self._mj_datas:
            mujoco.mj_step(self._mj_model, mj_data, n_steps)

    # We got new scene state, so anything depending on data must be refreshed
    for om in self.object_managers:
        om.invalidate_data_cache()

    self.camera_manager.registry.update_all_cameras(self)

mj_extensions

Classes:

Name Description
MjModelBindings

MjModelBindings 

MjModelBindings(model: MjModel, xml_path: str | None = None)

Methods:

Name Description
__getattr__
extract_mj_names
from_xml_path

Attributes:

Name Type Description
actuator_id2name dict[int, str]
actuator_name2id dict[str, int]
actuator_names list[str]
body_id2name dict[int, str]
body_name2id dict[str, int]
body_names list[str]
camera_id2name dict[int, str]
camera_name2id dict[str, int]
camera_names list[str]
equality_id2name dict[int, str]
equality_name2id dict[str, int]
equality_names list[str]
geom_id2name dict[int, str]
geom_name2id dict[str, int]
geom_names list[str]
id int
joint_id2name dict[int, str]
joint_name2id dict[str, int]
joint_names list[str]
light_id2name dict[int, str]
light_name2id dict[str, int]
light_names list[str]
mesh_id2name dict[int, str]
mesh_name2id dict[str, int]
mesh_names list[str]
model MjModel
sensor_id2name dict[int, str]
sensor_name2id dict[str, int]
sensor_names list[str]
site_id2name dict[int, str]
site_name2id dict[str, int]
site_names list[str]
tendon_id2name dict[int, str]
tendon_name2id dict[str, int]
tendon_names list[str]
xml Element
xml_path
Source code in molmo_spaces/env/mj_extensions.py 
52
53
54
55
def __init__(self, model: MjModel, xml_path: str | None = None) -> None:
    self.model = model
    self.xml_path = xml_path
    self.id = id(model)
actuator_id2name instance-attribute 
actuator_id2name: dict[int, str]
actuator_name2id instance-attribute 
actuator_name2id: dict[str, int]
actuator_names instance-attribute 
actuator_names: list[str]
body_id2name instance-attribute 
body_id2name: dict[int, str]
body_name2id instance-attribute 
body_name2id: dict[str, int]
body_names instance-attribute 
body_names: list[str]
camera_id2name instance-attribute 
camera_id2name: dict[int, str]
camera_name2id instance-attribute 
camera_name2id: dict[str, int]
camera_names instance-attribute 
camera_names: list[str]
equality_id2name instance-attribute 
equality_id2name: dict[int, str]
equality_name2id instance-attribute 
equality_name2id: dict[str, int]
equality_names instance-attribute 
equality_names: list[str]
geom_id2name instance-attribute 
geom_id2name: dict[int, str]
geom_name2id instance-attribute 
geom_name2id: dict[str, int]
geom_names instance-attribute 
geom_names: list[str]
id instance-attribute 
id: int = id(model)
joint_id2name instance-attribute 
joint_id2name: dict[int, str]
joint_name2id instance-attribute 
joint_name2id: dict[str, int]
joint_names instance-attribute 
joint_names: list[str]
light_id2name instance-attribute 
light_id2name: dict[int, str]
light_name2id instance-attribute 
light_name2id: dict[str, int]
light_names instance-attribute 
light_names: list[str]
mesh_id2name instance-attribute 
mesh_id2name: dict[int, str]
mesh_name2id instance-attribute 
mesh_name2id: dict[str, int]
mesh_names instance-attribute 
mesh_names: list[str]
model instance-attribute 
model: MjModel = model
sensor_id2name instance-attribute 
sensor_id2name: dict[int, str]
sensor_name2id instance-attribute 
sensor_name2id: dict[str, int]
sensor_names instance-attribute 
sensor_names: list[str]
site_id2name instance-attribute 
site_id2name: dict[int, str]
site_name2id instance-attribute 
site_name2id: dict[str, int]
site_names instance-attribute 
site_names: list[str]
tendon_id2name instance-attribute 
tendon_id2name: dict[int, str]
tendon_name2id instance-attribute 
tendon_name2id: dict[str, int]
tendon_names instance-attribute 
tendon_names: list[str]
xml cached property 
xml: Element
xml_path instance-attribute 
xml_path = xml_path
__getattr__ 
__getattr__(item)
Source code in molmo_spaces/env/mj_extensions.py 
 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
def __getattr__(self, item):
    if item.endswith("_names") or item.endswith("_name2id") or item.endswith("_id2name"):
        # Lazy-load the names, name2id, and id2name mappings

        k = item.split("_")[0]

        assert k in [
            "body",
            "joint",
            "geom",
            "site",
            "light",
            "camera",
            "actuator",
            "sensor",
            "tendon",
            "mesh",
            "equality",
        ]

        def key_to_nkey(key: str):
            return {
                "actuator": "nu",
                "joint": "njnt",
                "camera": "ncam",
                "equality": "neq",
            }.get(key, f"n{key}")

        def key_to_short_key(key: str):
            return {
                "joint": "jnt",
                "camera": "cam",
                "equality": "eq",
            }.get(key, f"{key}")

        names, name2id, id2name = self.extract_mj_names(
            getattr(self.model, f"name_{key_to_short_key(k)}adr"),
            getattr(self.model, key_to_nkey(k)),
            getattr(mujoco.mjtObj, f"mjOBJ_{k.upper()}"),
        )
        setattr(self, f"{k}_names", names)
        setattr(self, f"{k}_name2id", name2id)
        setattr(self, f"{k}_id2name", id2name)
        return getattr(self, item)
    else:
        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'")
extract_mj_names 
extract_mj_names(name_adr: ndarray | None, num_obj: int, obj_type: mjtObj)
Source code in molmo_spaces/env/mj_extensions.py 
68
69
def extract_mj_names(self, name_adr: np.ndarray | None, num_obj: int, obj_type: mjtObj):
    return extract_mj_names(self.model, name_adr, num_obj, obj_type)
from_xml_path classmethod 
from_xml_path(model: MjModel, filename: str) -> MjModelBindings
Source code in molmo_spaces/env/mj_extensions.py 
64
65
66
@classmethod
def from_xml_path(cls, model: MjModel, filename: str) -> "MjModelBindings":
    return cls(model, xml_path=filename)

object_manager

Classes:

Name Description
Context
ObjectManager

Attributes:

Name Type Description
ObjectOrNameOrIdType
log

ObjectOrNameOrIdType module-attribute 

ObjectOrNameOrIdType = MlSpacesObject | str | int

log module-attribute 

log = getLogger(__name__)

Context

Bases: Enum

Attributes:

Name Type Description
BENCH
OBJECT
SCENE
BENCH class-attribute instance-attribute 
BENCH = 'bench'
OBJECT class-attribute instance-attribute 
OBJECT = 'object'
SCENE class-attribute instance-attribute 
SCENE = 'scene'

ObjectManager 

ObjectManager(env: BaseMujocoEnv, batch_idx: int, caching_enabled: bool = True, name_caching_enabled: bool = True)

Methods:

Name Description
ancestors
approximate_supporting_geoms
category_from_name
children_lists
clear
compute_placement_region
default_object_context_synsets
descendants
expression_probs

Compute probabilities for each expression in the priority list.
fallback_expression
find_door_names

Find all valid door body names in the scene.
get_annotation_category
get_annotation_synset
get_body_to_geoms
get_cache_key
get_collision_obj_files

Return collision mesh file paths and the world pose of the body holding them.
get_context_objects
get_context_synsets
get_direct_children_names
get_door_bboxes_array

Get door collision geometry bounding boxes as an array.
get_free_objects

Return list of all bodies with free joints
get_geom_infos
get_mobile_objects

Return of list of all task relevant bodies i.e. not robots/policy objects
get_natural_object_names
get_object
get_object_body_id
get_object_by_name

Return the top-level object with the specified name, or None if not found.
get_object_hypernyms
get_object_name
get_objects_of_type

Return top-level scene objects for which at least one valid identifier matches any in object_types.
get_objects_that_are_on_top_of_object
get_parent_chain_names
get_pickup_candidates

Return top-level candidate small objects for pickup.
get_possible_object_types
get_receptacles

Return top-level receptacle objects (tables, counters, etc.)
get_support_below

Return the name of the support surface below (receptacle or room floor).
has_free_joint
has_receptacle_site
has_some_valid_identifier

If empty valid_identifiers, equivalent to accept any.
infer_room_name
invalidate_all_caches
invalidate_data_cache
invalidate_model_cache
is_excluded
is_object_articulable

If it has at least one hinge or slide joint, return True
is_pickup_candidate

If pickup_types is None, match None
is_receptacle

If receptacle_types is None, match None
is_structural
list_top_level_objects

List all non-structural top-level objects as Object instances.
most_concrete_synset
object_bottom_z
object_metadata
object_summary_str
objects_on_bench
objects_on_receptacle
prefilter_with_clip
referral_expression_priority

Return prioritized referral expressions, each presented as a tuple with:
sample_expression

Sample a candidate expression using a softmax distribution over priority scores.
summarize_top_level_bodies
thresholded_expression_priority

Thresholding to filter out ambiguous expressions.
top_level_bodies

Return bodies whose parent is the world body.
uid_to_annotation_for_type

Return list of uids in entire object library with object_type among the possible types

Attributes:

Name Type Description
STRUCTURAL_TYPES
data
model MjModel
model_path Path
scene_metadata dict | None
Source code in molmo_spaces/env/object_manager.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def __init__(
    self,
    env: "BaseMujocoEnv",
    batch_idx: int,
    caching_enabled: bool = True,
    name_caching_enabled: bool = True,
):
    self._env = env
    self._batch_idx = batch_idx
    self.data = env.mj_datas[batch_idx]

    # Caches for natural names and possible types
    self._name_caching_enabled = name_caching_enabled
    self._object_name_and_context_to_source_to_natural_names = {}
    self._object_name_and_context_to_natural_names = {}
    self._object_name_to_possible_type_names = {}

    # More caches
    self._caching_enabled = caching_enabled
    self._model_cache: dict[str, Any] = defaultdict(dict)
    self._data_cache: dict[str, Any] = defaultdict(dict)
STRUCTURAL_TYPES class-attribute instance-attribute 
STRUCTURAL_TYPES = {'world', 'room', 'floor', 'wall', 'window', 'doorframe', 'doorway', 'ceiling'}
data instance-attribute 
data = mj_datas[batch_idx]
model property 
model: MjModel
model_path property 
model_path: Path
scene_metadata cached property 
scene_metadata: dict | None
ancestors 
ancestors(object_or_name_or_id: ObjectOrNameOrIdType)
Source code in molmo_spaces/env/object_manager.py 
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
def ancestors(self, object_or_name_or_id: ObjectOrNameOrIdType):
    cache_in_use = self._model_cache

    oname = self.get_object_name(object_or_name_or_id)

    if "ancestors" not in cache_in_use[oname]:
        ancestors = MlSpacesObject.get_ancestors(
            self.model, self.get_object_body_id(object_or_name_or_id)
        )
        if self._caching_enabled:
            cache_in_use[oname]["ancestors"] = ancestors
        else:
            cache_in_use.pop(oname, None)
            return ancestors

    return cache_in_use[oname]["ancestors"]
approximate_supporting_geoms 
approximate_supporting_geoms(object_or_name_or_id: ObjectOrNameOrIdType, body_to_geoms: dict[int, list[int]])
Source code in molmo_spaces/env/object_manager.py 
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
def approximate_supporting_geoms(
    self, object_or_name_or_id: ObjectOrNameOrIdType, body_to_geoms: dict[int, list[int]]
):
    from shapely.geometry import Point, Polygon

    from molmo_spaces.utils.mujoco_scene_utils import body_aabb, geom_aabb

    def make_poly_and_zs(c, e):
        poly = Polygon(
            [
                Point(c[0] - e[0] / 2, c[1] - e[1] / 2),
                Point(c[0] + e[0] / 2, c[1] - e[1] / 2),
                Point(c[0] + e[0] / 2, c[1] + e[1] / 2),
                Point(c[0] - e[0] / 2, c[1] + e[1] / 2),
            ]
        )
        zb = c[2] - e[2] / 2  # bottom
        zt = c[2] + e[2] / 2  # top

        return poly, zb, zt

    def get_body_box(bid):
        c, e = body_aabb(self.model, self.data, bid)
        return make_poly_and_zs(c, e)

    def get_geom_box(geom_id):
        c, e = geom_aabb(self.model, self.data, [geom_id], tight_mesh=True)
        return make_poly_and_zs(c, e)

    oid = self.get_object(object_or_name_or_id).body_id
    opoly, oz, ot = get_body_box(oid)

    half_area = opoly.area / 2

    candidates = []
    for body, geoms in body_to_geoms.items():
        if body == oid:
            continue
        bpoly, bz, bt = get_body_box(body)
        if bpoly.intersects(opoly):
            for geom in geoms:
                gpoly, gz, gt = get_geom_box(geom)
                if abs(oz - gt) < 0.02:
                    # Keep geom if it "supports" at least half of the box
                    if gpoly.intersection(opoly).area >= half_area:
                        candidates.append((oz - gt, geom, body))

    candidates = sorted(candidates, key=lambda x: x[0])

    return candidates
category_from_name 
category_from_name(object_or_name_or_id: ObjectOrNameOrIdType)
Source code in molmo_spaces/env/object_manager.py 
580
581
582
583
584
585
def category_from_name(self, object_or_name_or_id: ObjectOrNameOrIdType):
    name = self.get_object_name(object_or_name_or_id)
    try:
        return re.compile(r"^(.*?)(?=[0-9a-fA-F]{32})").match(name).group(1).strip("_").lower()
    except AttributeError:
        return re.compile(r"^([A-Za-z_]+)").match(name).group(1).strip("_").lower()
children_lists 
children_lists()
Source code in molmo_spaces/env/object_manager.py 
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
def children_lists(self):
    cache_in_use = self._model_cache

    cache_key = "__scene__children_lists__"

    if cache_key not in cache_in_use:
        children_lists = MlSpacesObject.build_children_lists(self.model)
        if self._caching_enabled:
            cache_in_use[cache_key] = children_lists
        else:
            return children_lists

    return cache_in_use[cache_key]
clear 
clear()
Source code in molmo_spaces/env/object_manager.py 
1611
1612
def clear(self):
    self.invalidate_all_caches()
compute_placement_region 
compute_placement_region(object_or_name_or_id: ObjectOrNameOrIdType, margin_xy: float = 0.05) -> dict[str, ndarray]
Source code in molmo_spaces/env/object_manager.py 
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
def compute_placement_region(
    self, object_or_name_or_id: ObjectOrNameOrIdType, margin_xy: float = 0.05
) -> dict[str, np.ndarray]:
    cache_in_use = self._data_cache

    oname = self.get_object_name(object_or_name_or_id)

    if "placement_region" not in cache_in_use[oname]:
        obj_body_id = self.get_object_body_id(object_or_name_or_id)
        body_ids = {obj_body_id, *self.descendants(obj_body_id)}
        xy_min = np.array([np.inf, np.inf])
        xy_max = np.array([-np.inf, -np.inf])
        top_z = -np.inf
        found_any = False
        for geom_id in range(self.model.ngeom):
            try:
                if int(self.model.geom_bodyid[geom_id]) in body_ids:
                    center, dims = geom_aabb(self.model, self.data, [int(geom_id)])
                    aabb_min = center - dims / 2.0
                    aabb_max = center + dims / 2.0
                    xy_min = np.minimum(xy_min, aabb_min[:2])
                    xy_max = np.maximum(xy_max, aabb_max[:2])
                    top_z = max(top_z, float(aabb_max[2]))
                    found_any = True
            except Exception as e:
                print(f"Error computing AABB for geom {geom_id}: {e}")
                continue
        if (not found_any) or (not np.isfinite(top_z)):
            # Fallback small patch around object center if no valid AABBs
            pos = self.get_object(object_or_name_or_id).position
            placement_region = {
                "xy_min": pos[:2] - 0.3,
                "xy_max": pos[:2] + 0.3,
                "top_z": float(pos[2]),
            }
        else:
            placement_region = {
                "xy_min": xy_min - margin_xy,
                "xy_max": xy_max + margin_xy,
                "top_z": top_z,
            }

        if self._caching_enabled:
            cache_in_use[oname]["placement_region"] = placement_region
        else:
            cache_in_use.pop(oname, None)
            return placement_region

    return cache_in_use[oname]["placement_region"]
default_object_context_synsets 
default_object_context_synsets(target: ObjectOrNameOrIdType) -> set[str]
Source code in molmo_spaces/env/object_manager.py 
358
359
360
361
362
363
364
365
366
367
def default_object_context_synsets(self, target: ObjectOrNameOrIdType) -> set[str]:
    all_hypernyms = generate_all_hypernyms_with_exclusions(self.get_annotation_synset(target))
    if all_hypernyms:
        object_hypernyms = cast(
            set[str],
            {hyp.name() for hyp in all_hypernyms},
        )
        return object_hypernyms

    return set()
descendants 
descendants(object_or_name_or_id: ObjectOrNameOrIdType)
Source code in molmo_spaces/env/object_manager.py 
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
def descendants(self, object_or_name_or_id: ObjectOrNameOrIdType):
    cache_in_use = self._model_cache

    oname = self.get_object_name(object_or_name_or_id)

    if "descendants" not in cache_in_use[oname]:
        descendants = MlSpacesObject.get_descendants(
            self.children_lists(), self.get_object_body_id(object_or_name_or_id)
        )
        if self._caching_enabled:
            cache_in_use[oname]["descendants"] = descendants
        else:
            cache_in_use.pop(oname, None)
            return descendants

    return cache_in_use[oname]["descendants"]
expression_probs staticmethod 
expression_probs(priority: list[tuple[float, float, str]], temperature: float = 0.02) -> ndarray

Compute probabilities for each expression in the priority list.

Parameters

priority : list[tuple[float, float, str]] A list of tuples (sim_margin, target_similarity, value): - sim_margin (float): A priority score (typically in [-1, 1]) used as the primary softmax logit. - target_similarity (float): A secondary score indicating similarity to a target representation. - value (str): The actual expression or token to be sampled. temperature : float, optional Softmax temperature. Lower values (< 0.1) make sampling more deterministic by amplifying score differences; higher values increase exploration.

Returns

np.ndarray A probability distribution over the expressions in the priority list.

Source code in molmo_spaces/env/object_manager.py 
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
@staticmethod
def expression_probs(
    priority: list[tuple[float, float, str]],
    temperature: float = 2e-2,
) -> np.ndarray:
    """
    Compute probabilities for each expression in the priority list.

    Parameters
    ----------
    priority : list[tuple[float, float, str]]
        A list of tuples `(sim_margin, target_similarity, value)`:
        - sim_margin (float):    A priority score (typically in [-1, 1]) used as the
                                 primary softmax logit.
        - target_similarity (float): A secondary score indicating similarity to a
                                     target representation.
        - value (str):           The actual expression or token to be sampled.
    temperature : float, optional
        Softmax temperature. Lower values (< 0.1) make sampling more deterministic by
        amplifying score differences; higher values increase exploration.

    Returns
    -------
    np.ndarray
        A probability distribution over the expressions in the priority list.
    """

    scores = np.asarray([p[0] for p in priority], dtype=float)
    logits = np.exp(scores / max(1e-3, temperature))
    probs = logits / logits.sum()
    return probs
fallback_expression 
fallback_expression(object_name: str) -> str
Source code in molmo_spaces/env/object_manager.py 
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
def fallback_expression(self, object_name: str) -> str:
    # skip 3 digits at end of name and a hexadecimal code, keep everything before (in lower case)
    name = " ".join(object_name.split("_")[:-4]).lower()

    # remove digits
    name = re.sub(r"\d+", "", name).strip()

    # remove "obja" prefix, if any
    if name.startswith("obja"):
        name = name[4:]  # len("obja") = 4

    # no valid identifier
    if len(name.strip()) == 0:
        metadata = self.object_metadata(object_name)
        asset_id = metadata["asset_id"]
        name = ObjectMeta.annotation(asset_id)["category"].lower()

        # remove digits
        name = re.sub(r"\d+", "", name).strip()

        # remove "obja" prefix, if any
        if name.startswith("obja"):
            name = name[4:]  # len("obja") = 4

    # collapse multiple whitespace into a single space and strip ends
    return re.sub(r"\s+", " ", name).strip()
find_door_names 
find_door_names() -> list[str]

Find all valid door body names in the scene.

Valid doors are identified by the naming convention and whether they have joints.

Returns:

Type Description
list[str]

List of door body names found in the scene.
Source code in molmo_spaces/env/object_manager.py 
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
def find_door_names(self) -> list[str]:
    """Find all valid door body names in the scene.

    Valid doors are identified by the naming convention and whether they have joints.

    Returns:
        List of door body names found in the scene.
    """
    door_body_names = []
    for key, value in self.scene_metadata["objects"].items():
        if "doorway" in key:
            name_map = value.get("name_map", {})
            bodies = name_map.get("bodies", {})
            for k, v in bodies.items():
                if "_door_" in v:
                    door_object = Door(k, self.data)
                    if door_object.njoints > 0:
                        door_body_names.append(k)
    return door_body_names
get_annotation_category 
get_annotation_category(object_or_name_or_id: ObjectOrNameOrIdType) -> str
Source code in molmo_spaces/env/object_manager.py 
587
588
589
590
591
592
def get_annotation_category(self, object_or_name_or_id: ObjectOrNameOrIdType) -> str:
    object_name = self.get_object_name(object_or_name_or_id)

    object_meta = self.object_metadata(object_name)
    # Might contain Obja/obja, be lowercase, snake case, camel case
    return object_meta.get("category", self.category_from_name(object_name).strip())
get_annotation_synset 
get_annotation_synset(object_or_name_or_id: ObjectOrNameOrIdType) -> str | None
Source code in molmo_spaces/env/object_manager.py 
594
595
596
597
def get_annotation_synset(self, object_or_name_or_id: ObjectOrNameOrIdType) -> str | None:
    object_meta = self.object_metadata(object_or_name_or_id)

    return (ObjectMeta.annotation(object_meta.get("asset_id", "DUMMY")) or {}).get("synset")
get_body_to_geoms 
get_body_to_geoms()
Source code in molmo_spaces/env/object_manager.py 
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
def get_body_to_geoms(self):
    body_to_geom_ids = defaultdict(set)
    for geom_id in range(0, self.model.ngeom):
        body_id = int(self.model.geom(geom_id).bodyid.item())
        root_id = int(self.model.body(body_id).rootid.item())
        body_to_geom_ids[root_id].add(geom_id)
    return {
        key: sorted(values)
        for key, values in body_to_geom_ids.items()
        if not self.is_excluded(key)
    }
get_cache_key 
get_cache_key(object_or_name_or_id: ObjectOrNameOrIdType, context_synsets: Collection[str] = None)
Source code in molmo_spaces/env/object_manager.py 
471
472
473
474
475
476
477
def get_cache_key(
    self, object_or_name_or_id: ObjectOrNameOrIdType, context_synsets: Collection[str] = None
):
    plain_key = "__".join(
        [self.get_object_name(object_or_name_or_id)] + sorted(context_synsets or [])
    )
    return hashlib.md5(plain_key.encode()).hexdigest()
get_collision_obj_files 
get_collision_obj_files(object_or_name_or_id: ObjectOrNameOrIdType, xml_path: Path | None = None) -> tuple[list[Path], ndarray]

Return collision mesh file paths and the world pose of the body holding them.

Parses an XML file to find all descendant collision geoms of the object's body, resolves their mesh references to file paths, and returns the list of absolute .obj file paths together with the 4x4 world pose of the child body that contains the geoms.

Parameters:

Name Type Description Default
object_or_name_or_id ObjectOrNameOrIdType

The object to look up.

 required
xml_path Path | None

Optional XML file to parse instead of the scene model path. Useful for dynamically added objects whose meshes are defined in a standalone XML rather than the scene XML.

 None

Returns:

Type Description
list[Path]

A tuple of (obj_files, pose_4x4) where obj_files is a list of
ndarray

absolute .obj file paths and pose_4x4 is the 4x4 world pose of
tuple[list[Path], ndarray]

the child body holding the collision geoms.
Source code in molmo_spaces/env/object_manager.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
def get_collision_obj_files(
    self, object_or_name_or_id: ObjectOrNameOrIdType, xml_path: Path | None = None
) -> tuple[list[Path], np.ndarray]:
    """Return collision mesh file paths and the world pose of the body holding them.

    Parses an XML file to find all descendant collision geoms of the
    object's body, resolves their mesh references to file paths, and
    returns the list of absolute .obj file paths together with the 4x4
    world pose of the child body that contains the geoms.

    Args:
        object_or_name_or_id: The object to look up.
        xml_path: Optional XML file to parse instead of the scene model path.
                  Useful for dynamically added objects whose meshes are defined
                  in a standalone XML rather than the scene XML.

    Returns:
        A tuple of (obj_files, pose_4x4) where obj_files is a list of
        absolute .obj file paths and pose_4x4 is the 4x4 world pose of
        the child body holding the collision geoms.
    """
    from molmo_spaces.molmo_spaces_constants import ASSETS_DIR

    object_name = self.get_object_name(object_or_name_or_id)
    source_path = xml_path if xml_path is not None else self.model_path
    tree = ET.parse(source_path)
    root = tree.getroot()

    identity_pose = np.eye(4)

    # Build mesh name -> file map from the source XML's <asset> section
    if xml_path is not None:
        asset_elem = root.find("asset")
        mesh_name_to_file: dict[str, str] = {}
        if asset_elem is not None:
            for mesh_elem in asset_elem.findall("mesh"):
                name = mesh_elem.get("name")
                file = mesh_elem.get("file")
                if name and file:
                    mesh_name_to_file[name] = file
    else:
        mesh_name_to_file = self._mesh_name_to_file

    # Find the body element matching the object name
    if xml_path is not None:
        base_name = object_name.split("/")[-1]
        body_elem = root.find(f".//body[@name='{base_name}']")
        if body_elem is None:
            # Fallback: standalone XMLs typically have a single top-level body
            worldbody = root.find("worldbody")
            if worldbody is not None:
                bodies = worldbody.findall("body")
                if len(bodies) == 1:
                    body_elem = bodies[0]
    else:
        body_elem = root.find(f".//body[@name='{object_name}']")

    if body_elem is None:
        return [], identity_pose

    # Collect collision geoms (descendants with "collision" in their name and type="mesh")
    obj_files: list[Path] = []
    for geom in body_elem.iter("geom"):
        if geom.get("class") == "__VISUAL_MJT__":
            continue
        if geom.get("type") != "mesh":
            continue
        mesh_name = geom.get("mesh")
        if mesh_name and mesh_name in mesh_name_to_file:
            rel_path = mesh_name_to_file[mesh_name]
            if rel_path.startswith("../../"):
                abs_path = ASSETS_DIR / rel_path[len("../../") :]
            else:
                abs_path = source_path.parent / rel_path
            obj_files.append(abs_path.resolve())

    # Find the child body that holds the collision geoms and get its world pose
    parent_body_id = self.get_object_body_id(object_or_name_or_id)
    child_ids = np.where(self.model.body_parentid == parent_body_id)[0]
    geom_body_id = parent_body_id
    for child_id in child_ids:
        if self.model.body_geomnum[child_id] > 0:
            geom_body_id = int(child_id)
            break

    pose = np.eye(4)
    pose[:3, :3] = self.data.xmat[geom_body_id].reshape(3, 3)
    pose[:3, 3] = self.data.xpos[geom_body_id]

    return obj_files, pose
get_context_objects 
get_context_objects(object_or_name_or_id: ObjectOrNameOrIdType, context_type: Context = BENCH, bench_geom_ids: list[int] = None, cameras: list[str] = None, room_ids: list[str] = None, disable_caching: bool = True, **kwargs) -> list[MlSpacesObject]
Source code in molmo_spaces/env/object_manager.py 
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
def get_context_objects(
    self,
    object_or_name_or_id: ObjectOrNameOrIdType,
    context_type: Context = Context.BENCH,
    bench_geom_ids: list[int] = None,
    cameras: list[str] = None,
    room_ids: list[str] = None,
    disable_caching: bool = True,
    **kwargs,
) -> list[MlSpacesObject]:
    # Util for task sampling, so model-scope
    scope_cache = self._model_cache

    oname = self.get_object_name(object_or_name_or_id)

    if self._caching_enabled and not disable_caching:
        plain_key = (
            f"{context_type.name}"
            f"_{sorted(bench_geom_ids or [])}"
            f"_{sorted(cameras or [])}"
            f"_{sorted(room_ids or [])}"
        )
        key = hashlib.md5(plain_key.encode()).hexdigest()
    else:
        # Make sure we can't use cache
        if not disable_caching:
            assert not scope_cache
        key = "DUMMY_KEY"
        scope_cache.pop(key, None)

    if key not in scope_cache[oname]:
        target = self.get_object(object_or_name_or_id)

        if context_type == Context.BENCH:
            assert bench_geom_ids, f"When using {context_type}, `bench_geom_ids` must be given"

            # Note: We actually take any object on top of the bench object
            # (so not only on the given geom ids surfaces)
            context_objects = self.objects_on_bench(bench_geom_ids)

        # elif context_type == Context.VISIBLE:
        #     assert cameras, f"When using {context_type}, `cameras` must be given"
        #     raise NotImplementedError

        elif context_type == Context.SCENE:
            name_to_obj = {obj.name: obj for obj in self.list_top_level_objects()}
            if target.name not in name_to_obj:
                log.warning(f"{target.name} was not in `list_top_level_objects`")
                name_to_obj[target.name] = target

            context_objects = [name_to_obj[name] for name in sorted(name_to_obj.keys())]

        # elif context_type == Context.ROOM:
        #     assert room_ids, f"When using {context_type}, `room_ids` must be given"
        #     raise NotImplementedError

        else:  # if context_type == Context.OBJECT:
            context_objects = [target]

        if self._caching_enabled and not disable_caching:
            scope_cache[oname][key] = context_objects
        else:
            scope_cache.pop(oname, None)
            return context_objects

    return scope_cache[oname][key]
get_context_synsets 
get_context_synsets(context_objects: Collection[ObjectOrNameOrIdType]) -> list[str]
Source code in molmo_spaces/env/object_manager.py 
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
def get_context_synsets(
    self,
    context_objects: Collection[ObjectOrNameOrIdType],
) -> list[str]:
    annotated_synsets = sorted(
        set(
            self.get_annotation_synset(object_or_name_or_id)
            for object_or_name_or_id in context_objects
        )
        - {None}
    )

    if not annotated_synsets:
        return []

    valid_synsets = filter_synsets_to_remove_hyponyms(annotated_synsets)

    extended_synsets = set()
    for synset in valid_synsets:
        all_hypernyms = generate_all_hypernyms_with_exclusions(synset)
        if all_hypernyms:
            extended_synsets |= {syn.name() for syn in all_hypernyms}

    return sorted(extended_synsets)
get_direct_children_names 
get_direct_children_names(object_or_name_or_id: ObjectOrNameOrIdType) -> list[str]
Source code in molmo_spaces/env/object_manager.py 
1350
1351
1352
1353
1354
1355
1356
1357
1358
def get_direct_children_names(self, object_or_name_or_id: ObjectOrNameOrIdType) -> list[str]:
    # Rare use, so no cache
    direct_children_names = [
        self.model.body(c).name
        for c in MlSpacesObject.get_direct_children(
            self.children_lists(), self.get_object_body_id(object_or_name_or_id)
        )
    ]
    return direct_children_names
get_door_bboxes_array 
get_door_bboxes_array(object_or_name_or_id: ObjectOrNameOrIdType) -> ndarray

Get door collision geometry bounding boxes as an array. Returns: np.ndarray: Array of AABBs (center, size) for door collision geoms

Source code in molmo_spaces/env/object_manager.py 
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
def get_door_bboxes_array(self, object_or_name_or_id: ObjectOrNameOrIdType) -> np.ndarray:
    """Get door collision geometry bounding boxes as an array.
    Returns:
        np.ndarray: Array of AABBs (center, size) for door collision geoms
    """
    cache_in_use = self._model_cache

    oname = self.get_object_name(object_or_name_or_id)

    if "door_bboxes_array" not in cache_in_use[oname]:
        # Get all geoms for the door object
        geom_infos = self.get_geom_infos(object_or_name_or_id, include_descendants=True)
        door_bboxes = []

        for geom_info in geom_infos:
            geom_id = geom_info["id"]
            # Check if it's a collision geom (contype != 0 or conaffinity != 0)
            if (
                self.model.geom(geom_id).contype != 0
                or self.model.geom(geom_id).conaffinity != 0
            ):
                # Get AABB from model (center, size)
                aabb = self.model.geom_aabb[geom_id]
                door_bboxes.append(aabb)

        if not door_bboxes:
            door_bboxes = np.array([]).reshape(0, 6)
        else:
            door_bboxes = np.array(door_bboxes)

        if self._caching_enabled:
            cache_in_use[oname]["door_bboxes_array"] = door_bboxes
        else:
            cache_in_use.pop(oname, None)
            return door_bboxes

    return cache_in_use[oname]["door_bboxes_array"]
get_free_objects 
get_free_objects() -> list[MlSpacesObject]

Return list of all bodies with free joints

Source code in molmo_spaces/env/object_manager.py 
1222
1223
1224
1225
1226
1227
def get_free_objects(self) -> list[MlSpacesObject]:
    """Return list of all bodies with free joints"""
    model = self.model
    freejoints = np.where(model.jnt_type == mujoco.mjtJoint.mjJNT_FREE)[0]
    body_ids = model.jnt_bodyid[freejoints]
    return [self.get_object_by_name(model.body(id).name) for id in body_ids]
get_geom_infos 
get_geom_infos(object_or_name_or_id: ObjectOrNameOrIdType, include_descendants: bool = True, max_geoms: int | None = 2048) -> list[dict[str, object]]
Source code in molmo_spaces/env/object_manager.py 
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
def get_geom_infos(
    self,
    object_or_name_or_id: ObjectOrNameOrIdType,
    include_descendants: bool = True,
    max_geoms: int | None = 2048,
) -> list[dict[str, object]]:
    cache_in_use = self._data_cache

    oname = self.get_object_name(object_or_name_or_id)

    cache_key = f"geom_infos_{include_descendants}_{max_geoms}"

    if cache_key not in cache_in_use[oname]:
        # Delegate to MlSpacesObject.get_geom_infos
        obj = self.get_object(object_or_name_or_id)
        geoms = obj.get_geom_infos(include_descendants=include_descendants, max_geoms=max_geoms)
        if self._caching_enabled:
            cache_in_use[oname][cache_key] = geoms
        else:
            cache_in_use.pop(oname, None)
            return geoms

    return cache_in_use[oname][cache_key]
get_mobile_objects 
get_mobile_objects() -> list[MlSpacesObject]

Return of list of all task relevant bodies i.e. not robots/policy objects

Source code in molmo_spaces/env/object_manager.py 
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
def get_mobile_objects(self) -> list[MlSpacesObject]:
    """Return of list of all task relevant bodies i.e. not robots/policy objects"""

    task_objects = []
    for object_name, object_dict in self.scene_metadata["objects"].items():
        if not object_dict["is_static"]:
            try:
                task_object = create_mlspaces_body(self.data, object_name)
            except KeyError:
                log.warning("Could not find object %s in scene", object_name)
                continue
            task_objects.append(task_object)

    # TODO(Abhay): do we want this?
    for object_name in self._env.config.task_config.added_objects:
        task_object = create_mlspaces_body(self.data, object_name)
        task_objects.append(task_object)

    return task_objects
get_natural_object_names 
get_natural_object_names(object_or_name_or_id: ObjectOrNameOrIdType, context_synsets: list[str] = None)
Source code in molmo_spaces/env/object_manager.py 
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
def get_natural_object_names(
    self, object_or_name_or_id: ObjectOrNameOrIdType, context_synsets: list[str] = None
):
    cache_key = self.get_cache_key(object_or_name_or_id, context_synsets=context_synsets)

    if cache_key not in self._object_name_and_context_to_natural_names:
        all_names = set(
            sum(
                self._extract_names_from_context(
                    object_or_name_or_id, context_synsets
                ).values(),
                [],
            )
        )
        res = sorted(all_names, key=lambda x: (len(x), x))
        if self._name_caching_enabled:
            self._object_name_and_context_to_natural_names[cache_key] = res
        else:
            return res

    return self._object_name_and_context_to_natural_names[cache_key]
get_object 
get_object(object_or_name_or_id: ObjectOrNameOrIdType) -> MlSpacesObject
Source code in molmo_spaces/env/object_manager.py 
248
249
250
251
252
def get_object(self, object_or_name_or_id: ObjectOrNameOrIdType) -> MlSpacesObject:
    if isinstance(object_or_name_or_id, MlSpacesObject):
        return object_or_name_or_id

    return self.get_object_by_name(self.get_object_name(object_or_name_or_id))
get_object_body_id 
get_object_body_id(object_or_name_or_id: ObjectOrNameOrIdType) -> int
Source code in molmo_spaces/env/object_manager.py 
254
255
256
257
258
259
260
261
262
263
264
def get_object_body_id(self, object_or_name_or_id: ObjectOrNameOrIdType) -> int:
    if isinstance(object_or_name_or_id, str):
        object_body_id = self.get_object_by_name(
            self.get_object_name(object_or_name_or_id)
        ).body_id
    elif isinstance(object_or_name_or_id, MlSpacesObject):
        object_body_id = object_or_name_or_id.body_id
    else:
        object_body_id = object_or_name_or_id

    return object_body_id
get_object_by_name 
get_object_by_name(object_name: str) -> MlSpacesObject | None

Return the top-level object with the specified name, or None if not found.

Source code in molmo_spaces/env/object_manager.py 
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
def get_object_by_name(self, object_name: str) -> MlSpacesObject | None:
    """Return the top-level object with the specified name, or None if not found."""
    # static or articulable object
    # if articulable and has joints, return the articulation object
    # if static, return the static object

    cache_in_use = self._model_cache

    is_articulable = False

    if "articulable" not in cache_in_use[object_name]:
        is_articulable = self.is_object_articulable(object_name)
        if self._caching_enabled:
            cache_in_use[object_name]["articulable"] = is_articulable
        else:
            cache_in_use.pop(object_name, None)

    if is_articulable or (self._caching_enabled and cache_in_use[object_name]["articulable"]):
        return MlSpacesArticulationObject(data=self.data, object_name=object_name)

    return MlSpacesObject(data=self.data, object_name=object_name)
get_object_hypernyms 
get_object_hypernyms(target: ObjectOrNameOrIdType, context_synsets: Collection[str]) -> list[str]
Source code in molmo_spaces/env/object_manager.py 
376
377
378
379
380
def get_object_hypernyms(
    self, target: ObjectOrNameOrIdType, context_synsets: Collection[str]
) -> list[str]:
    object_hypernyms = self.default_object_context_synsets(target)
    return sorted(object_hypernyms & set(context_synsets))
get_object_name 
get_object_name(object_or_name_or_id: ObjectOrNameOrIdType) -> str
Source code in molmo_spaces/env/object_manager.py 
235
236
237
238
239
240
241
242
243
244
245
246
def get_object_name(self, object_or_name_or_id: ObjectOrNameOrIdType) -> str:
    if isinstance(object_or_name_or_id, str):
        object_name = object_or_name_or_id
    elif isinstance(object_or_name_or_id, MlSpacesObject):
        object_name = object_or_name_or_id.name
    else:
        object_name = self.model.body(object_or_name_or_id).name
        assert object_name is not None, (
            f"Object with ID {object_or_name_or_id} does not have a name"
        )

    return object_name
get_objects_of_type 
get_objects_of_type(object_types: Collection[str]) -> list[MlSpacesObject]

Return top-level scene objects for which at least one valid identifier matches any in object_types. If pickup is True, objects need to have a free joint to be returned. Returns MlSpacesObject instances (not MujocoBody) built from top-level bodies. If empty list, return any, if None, accept None

Source code in molmo_spaces/env/object_manager.py 
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
def get_objects_of_type(self, object_types: Collection[str]) -> list[MlSpacesObject]:
    """Return top-level scene objects for which at least one valid identifier matches any in object_types.
    If pickup is True, objects need to have a free joint to be returned.
    Returns MlSpacesObject instances (not MujocoBody) built from top-level bodies.
    If empty list, return any, if None, accept None
    """
    results: list[MlSpacesObject] = []

    for b in self.top_level_bodies():
        name = self.get_object_name(b)
        if not name:
            continue

        if self.is_excluded(name):
            continue

        if self.is_structural(name):
            continue

        if self.has_some_valid_identifier(name, object_types):
            results.append(self.get_object_by_name(name))

    return sorted(results, key=lambda x: x.name)
get_objects_that_are_on_top_of_object 
get_objects_that_are_on_top_of_object(object_or_name_or_id: ObjectOrNameOrIdType, pickup_types: Collection[str], z_above_min: float = 0.02, z_above_max: float = 0.6, constrain_to_pickupable: bool = True) -> list[MlSpacesObject]
Source code in molmo_spaces/env/object_manager.py 
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
def get_objects_that_are_on_top_of_object(
    self,
    object_or_name_or_id: ObjectOrNameOrIdType,
    pickup_types: Collection[str],
    z_above_min: float = 0.02,
    z_above_max: float = 0.60,
    constrain_to_pickupable: bool = True,
) -> list[MlSpacesObject]:
    cache_in_use = self._model_cache

    oname = self.get_object_name(object_or_name_or_id)

    if "objects_on_top" not in cache_in_use[oname]:
        # Compute host placement region
        obj = self.get_object(object_or_name_or_id)
        oid = obj.body_id
        region = self.compute_placement_region(object_or_name_or_id)
        xy_min, xy_max, top_z = region["xy_min"], region["xy_max"], float(region["top_z"])
        # Candidate roots: direct children of world, non-structural
        candidates = [
            b
            for b in self.top_level_bodies()
            if b != oid and not (self.is_structural(b) or self.is_excluded(b))
        ]

        roots_seen: set[int] = set()
        results: list[MlSpacesObject] = []
        for b in candidates:
            root_b = MlSpacesObject.find_top_object_body_id(self.model, b)
            if root_b in roots_seen:
                continue
            roots_seen.add(root_b)

            name = self.get_object_name(root_b)

            if constrain_to_pickupable:
                if not self.is_pickup_candidate(name, pickup_types):
                    continue

            pos = self.data.xpos[root_b]
            in_xy = (xy_min[0] <= pos[0] <= xy_max[0]) and (xy_min[1] <= pos[1] <= xy_max[1])
            above = (pos[2] >= top_z + z_above_min) and (pos[2] <= top_z + z_above_max)
            if in_xy and above:
                results.append(self.get_object_by_name(name))

        if self._caching_enabled:
            cache_in_use[oname]["objects_on_top"] = results
        else:
            cache_in_use.pop(oname, None)
            return results

    return cache_in_use[oname]["objects_on_top"]
get_parent_chain_names 
get_parent_chain_names(object_or_name_or_id: ObjectOrNameOrIdType) -> list[str]
Source code in molmo_spaces/env/object_manager.py 
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
def get_parent_chain_names(self, object_or_name_or_id: ObjectOrNameOrIdType) -> list[str]:
    # Rare use, so no cache

    try:
        parent_chain_names = [
            self.model.body(b).name for b in self.ancestors(object_or_name_or_id)
        ]
    except Exception as e:
        print(f"Error getting parent chain names: {e}")
        parent_chain_names = []

    return parent_chain_names
get_pickup_candidates 
get_pickup_candidates() -> list[MlSpacesObject]

Return top-level candidate small objects for pickup.

Source code in molmo_spaces/env/object_manager.py 
785
786
787
def get_pickup_candidates(self) -> list[MlSpacesObject]:
    """Return top-level candidate small objects for pickup."""
    return [o for o in self.list_top_level_objects() if self.has_free_joint(o)]
get_possible_object_types 
get_possible_object_types(object_or_name_or_id: ObjectOrNameOrIdType) -> list[str]
Source code in molmo_spaces/env/object_manager.py 
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
def get_possible_object_types(self, object_or_name_or_id: ObjectOrNameOrIdType) -> list[str]:
    object_name = self.get_object_name(object_or_name_or_id)

    if object_name not in self._object_name_to_possible_type_names:
        default_context = self.default_object_context_synsets(object_or_name_or_id)
        keys = set(
            self._extract_names_from_context(object_or_name_or_id, default_context).keys()
        )
        natural_names = set(
            sum(
                self._extract_names_from_context(
                    object_or_name_or_id, default_context
                ).values(),
                [],
            )
        )
        names = set(
            sum(
                [self._augment_natural_name(natural_name) for natural_name in natural_names],
                [],
            )
        )
        res = sorted(keys | names, key=lambda x: (len(x), x))
        if self._name_caching_enabled:
            self._object_name_to_possible_type_names[object_name] = res
        else:
            return res

    return self._object_name_to_possible_type_names[object_name]
get_receptacles 
get_receptacles() -> list[MlSpacesObject]

Return top-level receptacle objects (tables, counters, etc.)

Source code in molmo_spaces/env/object_manager.py 
781
782
783
def get_receptacles(self) -> list[MlSpacesObject]:
    """Return top-level receptacle objects (tables, counters, etc.)"""
    return [o for o in self.list_top_level_objects() if self.has_receptacle_site(o)]
get_support_below 
get_support_below(object_or_name_or_id: ObjectOrNameOrIdType, receptacle_types: Collection[str], z_clearance_eps: float = 0.1) -> str | None

Return the name of the support surface below (receptacle or room floor).

Source code in molmo_spaces/env/object_manager.py 
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
def get_support_below(
    self,
    object_or_name_or_id: ObjectOrNameOrIdType,
    receptacle_types: Collection[str],
    z_clearance_eps: float = 1e-1,
) -> str | None:
    """Return the name of the support surface below (receptacle or room floor)."""
    if self.is_structural(object_or_name_or_id) or self.is_excluded(object_or_name_or_id):
        return None

    cache_in_use = self._data_cache

    oname = self.get_object_name(object_or_name_or_id)

    if "support_below" not in cache_in_use[oname]:
        obj = self.get_object(object_or_name_or_id)
        obj_bottom = self.object_bottom_z(object_or_name_or_id)
        oid = obj.body_id
        xy = obj.position[:2]

        best_name = None
        best_top = -np.inf

        for b in self.top_level_bodies():
            if b == oid:
                continue

            name = self.get_object_name(b)
            if not (
                self.is_receptacle(name, receptacle_types)
                or any(
                    identifier in {"room", "floor"}
                    for identifier in self.get_possible_object_types(b)
                )
            ):
                continue

            region = self.compute_placement_region(name)
            xy_min, xy_max, top_z = region["xy_min"], region["xy_max"], float(region["top_z"])
            in_xy = (xy_min[0] <= xy[0] <= xy_max[0]) and (xy_min[1] <= xy[1] <= xy_max[1])
            clearance = obj_bottom - top_z
            z_check = top_z <= obj_bottom + z_clearance_eps

            is_better = top_z > best_top

            # Debug log for each candidate
            log.debug(
                f"[get_support_below] {oname} <- {name}: "
                f"clearance={clearance:.4f}, in_xy={in_xy}, "
                f"z_check={z_check} (top_z={top_z:.4f} <= obj_bottom={obj_bottom:.4f} + eps={z_clearance_eps:.4f}), "
                f"is_better={is_better} (top_z={top_z:.4f} vs best={best_top:.4f})"
            )

            # TODO The problem here must be that for objaverse assets, only one body has all the possible shelves
            #  but It could also be due to a suboptimal placement region computation
            if in_xy and (top_z <= obj_bottom + z_clearance_eps) and (top_z > best_top):
                best_top = top_z
                best_name = name
                log.debug(f"[get_support_below] New best support: {best_name}")

        if self._caching_enabled:
            cache_in_use[oname]["support_below"] = best_name
        else:
            cache_in_use.pop(oname, None)
            return best_name

    return cache_in_use[oname]["support_below"]
has_free_joint 
has_free_joint(object_or_name_or_id: ObjectOrNameOrIdType) -> bool
Source code in molmo_spaces/env/object_manager.py 
670
671
672
673
674
675
676
677
def has_free_joint(self, object_or_name_or_id: ObjectOrNameOrIdType) -> bool:
    # Ensure the object has a free joint
    try:
        MlSpacesFreeJointBody(self.data, self.get_object_name(object_or_name_or_id))
    except AssertionError:
        return False

    return True
has_receptacle_site 
has_receptacle_site(object_or_name_or_id: ObjectOrNameOrIdType) -> bool
Source code in molmo_spaces/env/object_manager.py 
619
620
621
def has_receptacle_site(self, object_or_name_or_id: ObjectOrNameOrIdType) -> bool:
    # Ensure receptacles have some site
    return bool(self.object_metadata(object_or_name_or_id).get("name_map", {}).get("sites", {}))
has_some_valid_identifier 
has_some_valid_identifier(object_or_name_or_id: ObjectOrNameOrIdType, valid_identifiers: Collection[str])

If empty valid_identifiers, equivalent to accept any. If None, equivalent to accept none

Source code in molmo_spaces/env/object_manager.py 
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
def has_some_valid_identifier(
    self, object_or_name_or_id: ObjectOrNameOrIdType, valid_identifiers: Collection[str]
):
    """
    If empty valid_identifiers, equivalent to accept any.
    If None, equivalent to accept none
    """

    if valid_identifiers is None:
        return False
    elif len(valid_identifiers) == 0:
        return True

    # Ensure at least one identifier within the given valid identifiers
    for identifier in self.get_possible_object_types(object_or_name_or_id):
        if identifier in valid_identifiers:
            return True

    return False
infer_room_name 
infer_room_name(object_or_name_or_id: ObjectOrNameOrIdType, receptacle_types: Collection[str]) -> str | None
Source code in molmo_spaces/env/object_manager.py 
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
def infer_room_name(
    self, object_or_name_or_id: ObjectOrNameOrIdType, receptacle_types: Collection[str]
) -> str | None:
    # Note: If passing [] for receptacle types, only finds room or floor
    cache_in_use = self._data_cache

    oname = self.get_object_name(object_or_name_or_id)

    receptacle_key = hashlib.md5(f"{sorted(receptacle_types) or []}".encode()).hexdigest()
    cache_key = f"room_{receptacle_key}"

    if cache_key not in cache_in_use[oname]:

        def dfs(cur_object_or_name_or_id):
            support_name = self.get_support_below(cur_object_or_name_or_id, receptacle_types)
            if support_name is None:
                return None

            # If the support is a room or floor, return room or map floor->room
            for t in self.get_possible_object_types(support_name):
                # This should be terminal
                if t == "room":
                    return support_name

                # This should be terminal
                if t == "floor":
                    # Find enclosing room by nearest/first room in parent chain
                    ancestors = self.ancestors(cur_object_or_name_or_id)
                    for a in ancestors:
                        aname = self.get_object_name(a)
                        for tt in self.get_possible_object_types(aname):
                            if tt == "room":
                                return aname
                    raise ValueError("BUG? Floor has no room ancestor")

            # else: keep searching
            return dfs(support_name)

        res = dfs(object_or_name_or_id)
        if self._caching_enabled:
            cache_in_use[oname][cache_key] = res
        else:
            cache_in_use.pop(oname, None)
            return res

    return cache_in_use[oname][cache_key]
invalidate_all_caches 
invalidate_all_caches() -> None
Source code in molmo_spaces/env/object_manager.py 
102
103
104
105
106
107
108
109
110
111
112
def invalidate_all_caches(self) -> None:
    self._model_cache.clear()
    self._data_cache.clear()
    self._object_name_to_possible_type_names.clear()
    self._object_name_and_context_to_natural_names.clear()
    self._object_name_and_context_to_source_to_natural_names.clear()

    if "scene_metadata" in self.__dict__:
        del self.__dict__["scene_metadata"]
    if "_mesh_name_to_file" in self.__dict__:
        del self.__dict__["_mesh_name_to_file"]
invalidate_data_cache 
invalidate_data_cache() -> None
Source code in molmo_spaces/env/object_manager.py 
96
97
def invalidate_data_cache(self) -> None:
    self._data_cache.clear()
invalidate_model_cache 
invalidate_model_cache() -> None
Source code in molmo_spaces/env/object_manager.py 
 99
100
def invalidate_model_cache(self) -> None:
    self._model_cache.clear()
is_excluded 
is_excluded(object_or_name_or_id: ObjectOrNameOrIdType) -> bool
Source code in molmo_spaces/env/object_manager.py 
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
def is_excluded(self, object_or_name_or_id: ObjectOrNameOrIdType) -> bool:
    cache_in_use = self._model_cache

    oname = self.get_object_name(object_or_name_or_id)

    if "excluded" not in cache_in_use[oname]:
        is_excluded = (
            self._env.config.robot_config.robot_namespace in oname
        ) or not descendant_geoms(
            self.model, self.get_object(object_or_name_or_id).body_id, True
        )
        if self._caching_enabled:
            cache_in_use[oname]["excluded"] = is_excluded
        else:
            cache_in_use.pop(oname, None)
            return is_excluded

    return cache_in_use[oname]["excluded"]
is_object_articulable 
is_object_articulable(object_name: str) -> bool

If it has at least one hinge or slide joint, return True else return False

Source code in molmo_spaces/env/object_manager.py 
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
def is_object_articulable(self, object_name: str) -> bool:
    """
    If it has at least one hinge or slide joint, return True
    else return False
    """
    body_id = self.model.body(object_name).id

    # go through all joints
    for joint_id in range(self.model.njnt):
        # if root body is same as the body_id, then joint is part of object
        if self.model.body(self.model.joint(joint_id).bodyid[0]).rootid[0] == body_id:
            if self.model.joint(joint_id).type in [
                mujoco.mjtJoint.mjJNT_HINGE,
                mujoco.mjtJoint.mjJNT_SLIDE,
            ]:
                return True

    return False
is_pickup_candidate 
is_pickup_candidate(object_or_name_or_id: ObjectOrNameOrIdType, pickup_types: Collection[str]) -> bool

If pickup_types is None, match None If empty list, match any with free joint

Source code in molmo_spaces/env/object_manager.py 
679
680
681
682
683
684
685
686
687
688
def is_pickup_candidate(
    self, object_or_name_or_id: ObjectOrNameOrIdType, pickup_types: Collection[str]
) -> bool:
    """
    If pickup_types is None, match None
    If empty list, match any with free joint
    """
    return self.has_free_joint(object_or_name_or_id) and self.has_some_valid_identifier(
        object_or_name_or_id, pickup_types
    )
is_receptacle 
is_receptacle(object_or_name_or_id: ObjectOrNameOrIdType, receptacle_types: Collection[str]) -> bool

If receptacle_types is None, match None If empty list, match any with receptacle site

Source code in molmo_spaces/env/object_manager.py 
623
624
625
626
627
628
629
630
631
632
def is_receptacle(
    self, object_or_name_or_id: ObjectOrNameOrIdType, receptacle_types: Collection[str]
) -> bool:
    """
    If receptacle_types is None, match None
    If empty list, match any with receptacle site
    """
    return self.has_receptacle_site(object_or_name_or_id) and self.has_some_valid_identifier(
        object_or_name_or_id, receptacle_types
    )
is_structural 
is_structural(object_or_name_or_id: ObjectOrNameOrIdType) -> bool
Source code in molmo_spaces/env/object_manager.py 
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
def is_structural(self, object_or_name_or_id: ObjectOrNameOrIdType) -> bool:
    cache_in_use = self._model_cache

    oname = self.get_object_name(object_or_name_or_id)

    if "structural" not in cache_in_use[oname]:
        is_structural = self.has_some_valid_identifier(
            object_or_name_or_id, self.STRUCTURAL_TYPES
        )
        if self._caching_enabled:
            cache_in_use[oname]["structural"] = is_structural
        else:
            cache_in_use.pop(oname, None)
            return is_structural

    return cache_in_use[oname]["structural"]
list_top_level_objects 
list_top_level_objects() -> list[MlSpacesObject]

List all non-structural top-level objects as Object instances.

Source code in molmo_spaces/env/object_manager.py 
770
771
772
773
774
775
776
777
778
779
def list_top_level_objects(self) -> list[MlSpacesObject]:
    """List all non-structural top-level objects as Object instances."""
    objs: list[MlSpacesObject] = []
    for b in self.top_level_bodies():
        name = self.get_object_name(b)
        if not name or self.is_structural(name) or self.is_excluded(name):
            continue
        objs.append(self.get_object_by_name(name))

    return sorted(objs, key=lambda x: x.name)
most_concrete_synset staticmethod 
most_concrete_synset(all_hypernyms) -> Any
Source code in molmo_spaces/env/object_manager.py 
369
370
371
372
373
374
@staticmethod
def most_concrete_synset(all_hypernyms) -> Any:
    for current in all_hypernyms:
        if not any(is_hypernym_of(other, current) for other in all_hypernyms - {current}):
            return current
    raise ValueError(f"No most concrete element among {all_hypernyms}?!")
object_bottom_z 
object_bottom_z(object_or_name_or_id: ObjectOrNameOrIdType) -> float
Source code in molmo_spaces/env/object_manager.py 
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
def object_bottom_z(self, object_or_name_or_id: ObjectOrNameOrIdType) -> float:
    cache_in_use = self._data_cache

    oname = self.get_object_name(object_or_name_or_id)

    if "bottom_z" not in cache_in_use[oname]:
        # Bottom Z from aggregated AABB minima
        oid = self.get_object_name(object_or_name_or_id)
        body_ids = {oid, *self.descendants(oid)}
        bottom_z = np.inf
        for geom_id in range(self.model.ngeom):
            try:
                if int(self.model.geom_bodyid[geom_id]) in body_ids:
                    aabb_min, _ = geom_aabb(self.model, self.data, [geom_id])
                    bottom_z = min(bottom_z, float(aabb_min[2]))
            except Exception as e:
                print(f"Error getting object bottom z: {e}")
                continue

        if not np.isfinite(bottom_z):
            bottom_z = float(self.get_object(object_or_name_or_id).position[2])

        if self._caching_enabled:
            cache_in_use[oname]["bottom_z"] = bottom_z
        else:
            cache_in_use.pop(oname, None)
            return bottom_z

    return cache_in_use[oname]["bottom_z"]
object_metadata 
object_metadata(object_or_name_or_id: ObjectOrNameOrIdType) -> dict
Source code in molmo_spaces/env/object_manager.py 
114
115
116
117
118
119
def object_metadata(self, object_or_name_or_id: ObjectOrNameOrIdType) -> dict:
    return (
        (self.scene_metadata or {})
        .get("objects", {})
        .get(self.get_object_name(object_or_name_or_id), {})
    )
object_summary_str 
object_summary_str(object_or_name_or_id: ObjectOrNameOrIdType, receptacle_types: Collection[str]) -> str
Source code in molmo_spaces/env/object_manager.py 
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
def object_summary_str(
    self, object_or_name_or_id: ObjectOrNameOrIdType, receptacle_types: Collection[str]
) -> str:
    # Rare use, so no cache
    obj = self.get_object(object_or_name_or_id)
    name = obj.name
    category = self.get_annotation_category(obj)
    synset = self.get_annotation_synset(obj)
    pos = obj.position
    support = self.get_support_below(obj, receptacle_types)
    room = self.infer_room_name(obj, receptacle_types)
    on_str = f"on {support}" if support else "on <unknown>"
    room_str = f"in {room}" if room else "in <unknown>"
    return f"{name} (category={category} synset={synset}) center=({pos[0]:.3f},{pos[1]:.3f},{pos[2]:.3f}) {on_str}, {room_str}"
objects_on_bench 
objects_on_bench(bench_geom_ids: list[int], angle_threshold: float = radians(30), fallback_thres=0.01, attempt_contact: bool = True) -> list[MlSpacesObject]
Source code in molmo_spaces/env/object_manager.py 
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
def objects_on_bench(
    self,
    bench_geom_ids: list[int],
    angle_threshold: float = np.radians(30),
    fallback_thres=0.01,  # 1 cm
    attempt_contact: bool = True,
) -> list[MlSpacesObject]:
    return self.objects_on_receptacle(
        objs_to_check=self.list_top_level_objects(),
        bench_geom_ids=bench_geom_ids,
        angle_threshold=angle_threshold,
        fallback_thres=fallback_thres,
        attempt_contact=attempt_contact,
    )
objects_on_receptacle 
objects_on_receptacle(objs_to_check: list[MlSpacesObject], bench_geom_ids: list[int], angle_threshold: float = radians(30), fallback_thres=0.01, attempt_contact: bool = True) -> list[MlSpacesObject]
Source code in molmo_spaces/env/object_manager.py 
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
def objects_on_receptacle(
    self,
    objs_to_check: list[MlSpacesObject],
    bench_geom_ids: list[int],
    angle_threshold: float = np.radians(30),
    fallback_thres=0.01,  # 1 cm
    attempt_contact: bool = True,
) -> list[MlSpacesObject]:
    from shapely.geometry import Point, Polygon

    from molmo_spaces.utils.mujoco_scene_utils import body_aabb

    valid_names = {obj.name for obj in objs_to_check}

    data = self.data
    model = data.model

    cos_threshold = np.cos(angle_threshold)

    object_names = set()

    if attempt_contact:
        # Based on contact with the bench geom id:
        for c in data.contact:
            root_body1, root_body2 = model.body_rootid[model.geom_bodyid[c.geom]]
            if (c.geom[0] in bench_geom_ids) ^ (c.geom[1] in bench_geom_ids):
                other_body_id = root_body1 if c.geom[0] not in bench_geom_ids else root_body2
                normal = c.frame[:3] / np.linalg.norm(c.frame[:3])
                if c.geom[1] in bench_geom_ids:
                    normal = -normal
                body_aabb_center, _ = body_aabb(model, data, other_body_id)
                if c.pos[2] < body_aabb_center[2] and normal[2] >= cos_threshold:
                    cname = model.body(other_body_id).name
                    if cname in valid_names:
                        object_names.add(cname)

    contactless_object_names = set()
    seen_poly_z = set()

    # Fallback: list with objects with aabbs overlapping >= 50% in xy with the bench and "just above" in z
    for geom_id in bench_geom_ids:
        # Take full body
        bc, be = body_aabb(model, data, model.body_rootid[model.geom_bodyid[geom_id]])

        # Avoid recomputing if all geom ids are part of the same body
        cur_str = f"{np.round(bc, 3).tolist() + np.round(be, 3).tolist()}"
        if cur_str in seen_poly_z:
            continue
        seen_poly_z.add(cur_str)

        bench_poly = Polygon(
            [
                Point(bc[0] - be[0] / 2, bc[1] - be[1] / 2),
                Point(bc[0] + be[0] / 2, bc[1] - be[1] / 2),
                Point(bc[0] + be[0] / 2, bc[1] + be[1] / 2),
                Point(bc[0] - be[0] / 2, bc[1] + be[1] / 2),
            ]
        )
        bench_z = bc[2] + be[2] / 2

        # # Debug
        # for object_name in object_names:
        #     body_id = model.body(object_name).id
        #     obj_center, obj_ext = body_aabb(model, data, body_id)
        #     assert bench_poly.contains(Point(*obj_center[:2]))
        #     obj_base_z = obj_center[2] - obj_ext[2] / 2
        #     print(body_id, object_name, abs(bench_z - obj_base_z) <= fallback_thres)

        for obj in objs_to_check:
            object_name = obj.name

            if object_name in object_names:
                continue
            if object_name not in self.scene_metadata.get("objects", {}):
                continue
            if object_name in contactless_object_names:
                continue

            body_id = model.body(object_name).id
            obj_center, obj_ext = body_aabb(model, data, body_id)
            if bench_poly.contains(Point(*obj_center[:2])):
                obj_base_z = obj_center[2] - obj_ext[2] / 2
                # Check the base of the object is somewhere between the bbox center below and the fallback thres above
                if -be[2] / 2 <= obj_base_z - bench_z <= fallback_thres:
                    contactless_object_names.add(object_name)

    # Combine all objects in a single list
    object_list = [
        self.get_object_by_name(object_name)
        for object_name in sorted(object_names | contactless_object_names)
    ]

    return object_list
prefilter_with_clip staticmethod 
prefilter_with_clip(target_type: str | list[str], uids: list[str], min_sim: float = 0.25) -> list[str]
Source code in molmo_spaces/env/object_manager.py 
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
@staticmethod
def prefilter_with_clip(
    target_type: str | list[str], uids: list[str], min_sim: float = 0.25
) -> list[str]:
    all_uids = set(ObjectMeta.all_uids())
    kept_uids = [uid for uid in uids if uid in all_uids]
    if not kept_uids:
        return []

    if isinstance(target_type, str):
        expressions = [normalize_expression(target_type)]
    else:
        expressions = [normalize_expression(cur_type) for cur_type in target_type]
    try:
        text_features = compute_text_clip(expressions)
    except ValueError:
        log.warning("No image features, using all uids")
        return uids

    img_features = ObjectMeta.img_features(kept_uids)
    sims = clip_sim(img_features, text_features).max(axis=1).flatten()
    return np.array(kept_uids)[sims >= min_sim].tolist()
referral_expression_priority 
referral_expression_priority(object_or_name_or_id: ObjectOrNameOrIdType, context_object_or_name_or_ids: Collection[ObjectOrNameOrIdType]) -> list[tuple[float, float, str]]

Return prioritized referral expressions, each presented as a tuple with: - similarity score margin versus highest-scoring distractor in context - similarity score for target object - referral expression Priority is based on margin, similarity to target, expression length,

Source code in molmo_spaces/env/object_manager.py 
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
def referral_expression_priority(
    self,
    object_or_name_or_id: ObjectOrNameOrIdType,
    context_object_or_name_or_ids: Collection[ObjectOrNameOrIdType],
) -> list[tuple[float, float, str]]:
    """
    Return prioritized referral expressions, each presented as a tuple with:
      - similarity score margin versus highest-scoring distractor in context
      - similarity score for target object
      - referral expression
    Priority is based on margin, similarity to target, expression length,
    """
    # Rare use, so no cache

    target_name = self.get_object_name(object_or_name_or_id)
    context_names = {
        self.get_object_name(maybe_obj) for maybe_obj in context_object_or_name_or_ids
    }
    assert target_name in context_names, f"Target {target_name} must be in context"

    name_to_uid = {
        self.get_object_name(maybe_obj): self.object_metadata(maybe_obj)["asset_id"]
        for maybe_obj in context_object_or_name_or_ids
        if "asset_id" in self.object_metadata(maybe_obj)
    }

    assert target_name in name_to_uid, f"Target {target_name} has no metadata"

    asset_ids = sorted(set(name_to_uid.values()))

    img = []
    kept_asset_ids = []
    for asset_id in asset_ids:
        try:
            img.append(ObjectMeta.img_features(asset_id))
            kept_asset_ids.append(asset_id)
        except ValueError:
            log.warning(f"No image features for {asset_id}, ignoring.")

    if name_to_uid[target_name] not in kept_asset_ids:
        raise ValueError(f"Missing asset id {name_to_uid[target_name]} from {target_name}")

    img = np.concatenate(img, axis=0)
    asset_ids = kept_asset_ids

    descriptions = self.get_natural_object_names(
        object_or_name_or_id,
        self.get_context_synsets(context_object_or_name_or_ids),
    )
    try:
        sim = clip_sim(img, compute_text_clip(descriptions))
    except NameError:
        log.warning("No CLIP module, using dummy description scores.")
        # e.g. when you don't want to install it / no gpu (?)
        names = self.get_natural_object_names(object_or_name_or_id, [])
        return [(1.0, 1.0, name) for name in names]

    target_idx = asset_ids.index(name_to_uid[target_name])
    target_sims = sim[target_idx]

    if sim.shape[0] > 1:
        sorting = np.argsort(sim, axis=0)

        # last row has the indices of largest similarity
        nearest_sim = sim[sorting[sim.shape[0] - 1], np.arange(sim.shape[1])]
        runner_up_sim = sim[sorting[sim.shape[0] - 2], np.arange(sim.shape[1])]

        deltas = target_sims - np.where(target_sims == nearest_sim, runner_up_sim, nearest_sim)
    else:
        deltas = target_sims

    return sorted(
        [
            (delta, target_sim, description)
            for delta, target_sim, description in zip(deltas, target_sims, descriptions)
        ],
        key=lambda x: (x[0], x[1], len(x[2]), x[2]),
        reverse=True,
    )
sample_expression staticmethod 
sample_expression(priority: list[tuple[float, float, str]], temperature: float = 0.02) -> str

Sample a candidate expression using a softmax distribution over priority scores.

Parameters

priority : list[tuple[float, float, str]] A list of tuples (sim_margin, target_similarity, value): - sim_margin (float): A priority score (typically in [-1, 1]) used as the primary softmax logit. - target_similarity (float): A secondary score indicating similarity to a target representation. - value (str): The actual expression or token to be sampled. temperature : float, optional Softmax temperature. Lower values (< 0.1) make sampling more deterministic by amplifying score differences; higher values increase exploration. Default is 2e-2, producing a sharp distribution.

Returns

str The sampled expression/value from the given list.

Notes

The final sampling is categorical: p(i) = softmax(sim_margin_i / temperature).

Source code in molmo_spaces/env/object_manager.py 
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
@staticmethod
def sample_expression(
    priority: list[tuple[float, float, str]],
    temperature: float = 2e-2,
) -> str:
    """
    Sample a candidate expression using a softmax distribution over priority scores.

    Parameters
    ----------
    priority : list[tuple[float, float, str]]
        A list of tuples `(sim_margin, target_similarity, value)`:
        - sim_margin (float):    A priority score (typically in [-1, 1]) used as the
                                 primary softmax logit.
        - target_similarity (float): A secondary score indicating similarity to a
                                     target representation.
        - value (str):           The actual expression or token to be sampled.
    temperature : float, optional
        Softmax temperature. Lower values (< 0.1) make sampling more deterministic by
        amplifying score differences; higher values increase exploration.
        Default is `2e-2`, producing a sharp distribution.

    Returns
    -------
    str
        The sampled expression/value from the given list.

    Notes
    -----
    The final sampling is categorical:
        p(i) = softmax(sim_margin_i / temperature).
    """

    probs = ObjectManager.expression_probs(priority, temperature)
    return np.random.choice([p[-1] for p in priority], p=probs)
summarize_top_level_bodies 
summarize_top_level_bodies(receptacle_types: Collection[str], limit: int = 50) -> list[str]
Source code in molmo_spaces/env/object_manager.py 
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
def summarize_top_level_bodies(
    self, receptacle_types: Collection[str], limit: int = 50
) -> list[str]:
    # No widespread usage, so no cache
    lines: list[str] = []
    count = 0
    for b in self.top_level_bodies():
        if self.is_structural(b) or self.is_excluded(b):
            continue
        try:
            obj = self.get_object(b)
            lines.append(self.object_summary_str(obj, receptacle_types))
            count += 1
            if count >= limit:
                break
        except KeyboardInterrupt:
            raise
        except Exception as e:
            try:
                name = self.get_object(b).name
            except KeyboardInterrupt:
                raise
            print(f"Error summarizing top-level bodies (obj with name '{name}'): {e}")
            continue
    return lines
thresholded_expression_priority staticmethod 
thresholded_expression_priority(priority: list[tuple[float, float, str]], sim_margin_threshold: float = 0.03, target_sim_threshold: float = 0.1) -> list[tuple[float, float, str]]

Thresholding to filter out ambiguous expressions.

Parameters

priority : list[tuple[float, float, str]] A list of tuples (sim_margin, target_similarity, value): - sim_margin (float): A priority score (typically in [-1, 1]) used as the primary softmax logit. - target_similarity (float): A secondary score indicating similarity to a target representation. - value (str): The actual expression or token to be sampled. sim_margin_threshold : float, optional Minimum allowed sim_margin for an item to be considered when thresholding. target_sim_threshold : float, optional Minimum required target_similarity for thresholding.

Returns

list[tuple[float, float, str]] (filtered prioritization) The sampled expression/value from the filtered list.

Source code in molmo_spaces/env/object_manager.py 
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
@staticmethod
def thresholded_expression_priority(
    priority: list[tuple[float, float, str]],
    sim_margin_threshold: float = 0.03,
    target_sim_threshold: float = 0.1,
) -> list[tuple[float, float, str]]:
    """
    Thresholding to filter out ambiguous expressions.

    Parameters
    ----------
    priority : list[tuple[float, float, str]]
        A list of tuples `(sim_margin, target_similarity, value)`:
        - sim_margin (float):    A priority score (typically in [-1, 1]) used as the
                                 primary softmax logit.
        - target_similarity (float): A secondary score indicating similarity to a
                                     target representation.
        - value (str):           The actual expression or token to be sampled.
    sim_margin_threshold : float, optional
        Minimum allowed sim_margin for an item to be considered when thresholding.
    target_sim_threshold : float, optional
        Minimum required target_similarity for thresholding.

    Returns
    -------
    list[tuple[float, float, str]] (filtered prioritization)
        The sampled expression/value from the filtered list.
    """

    return [
        p for p in priority if p[0] >= sim_margin_threshold and p[1] >= target_sim_threshold
    ]
top_level_bodies 
top_level_bodies() -> list[int]

Return bodies whose parent is the world body.

Source code in molmo_spaces/env/object_manager.py 
690
691
692
693
694
695
696
697
698
699
700
701
702
703
def top_level_bodies(self) -> list[int]:
    """Return bodies whose parent is the world body."""
    cache_in_use = self._model_cache
    cache_key = "__scene__top_level_bodies__"

    if cache_key not in cache_in_use:
        bodies = MlSpacesObject.get_top_level_bodies(self.model)
        if self._caching_enabled:
            cache_in_use[cache_key] = bodies
        else:
            cache_in_use.pop(cache_key, None)
            return bodies

    return cast(list[int], cache_in_use[cache_key])
uid_to_annotation_for_type staticmethod 
uid_to_annotation_for_type(object_type: str) -> dict[str, dict]

Return list of uids in entire object library with object_type among the possible types

Note: for now, we are reusing the functionality based on use of scene metadata in this class. We also don't cache results, so best keep them cached in your caller

Source code in molmo_spaces/env/object_manager.py 
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
@staticmethod
def uid_to_annotation_for_type(object_type: str) -> dict[str, dict]:
    """
    Return list of uids in entire object library with object_type among the possible types

    Note: for now, we are reusing the functionality based on use of scene
    metadata in this class. We also don't cache results, so best keep them
    cached in your caller
    """
    object_type = object_type.lower()

    valid_uids = {}

    class DummyEnv:
        mj_datas = [None]

    om = ObjectManager(DummyEnv(), -1)  # type:ignore
    om.scene_metadata = {"objects": {}}

    for uid, anno in ObjectMeta.annotation().items():
        category = anno["category"]
        name = f"{category.lower()}_{hashlib.md5(uid.encode()).hexdigest()}_0_0_0"

        om.scene_metadata["objects"][name] = {
            "asset_id": uid,
            "category": category,
            "object_enum": "temp_object",
        }

        possible_types = om.get_possible_object_types(name)
        if object_type in possible_types:
            valid_uids[uid] = anno

        # Avoid wasting memory
        om.scene_metadata["objects"].pop(name)
        om._object_name_to_possible_type_names = {}
        om._object_name_and_context_to_source_to_natural_names = {}

    return valid_uids

rby1_sensors

Classes:

Name Description
RBY1GraspPoseSensor

Sensor for RBY1 grasp pose in 7D format (can be current TCP or planned grasp pose).
RBY1GraspStateSensor
RBY1RobotStateSensor

Sensor for RBY1 robot joint positions, velocities, and dual end-effector poses.
RBY1TCPPoseSensor

Sensor for RBY1 TCP (Tool Center Point) poses in 7D format for both arms.

Functions:

Name Description
get_rby1_door_opening_sensors

Get core sensors for RBY1 door opening data generation.

RBY1GraspPoseSensor 

RBY1GraspPoseSensor(uuid: str = 'rby1_grasp_pose', arm_side: str = 'left')

Bases: Sensor

Sensor for RBY1 grasp pose in 7D format (can be current TCP or planned grasp pose).

Parameters:

Name Type Description Default
uuid str

Unique identifier for this sensor

 'rby1_grasp_pose'
arm_side str

Which arm to track ("left" or "right")

 'left'

Methods:

Name Description
get_observation

Get grasp pose (using current TCP pose as proxy) for the specified arm.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
arm_side
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/rby1_sensors.py 
176
177
178
179
180
181
182
183
184
def __init__(self, uuid: str = "rby1_grasp_pose", arm_side: str = "left") -> None:
    """
    Args:
        uuid: Unique identifier for this sensor
        arm_side: Which arm to track ("left" or "right")
    """
    self.arm_side = arm_side
    observation_space = gyms.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32)
    super().__init__(uuid=uuid, observation_space=observation_space)
arm_side instance-attribute 
arm_side = arm_side
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get grasp pose (using current TCP pose as proxy) for the specified arm.

Source code in molmo_spaces/env/rby1_sensors.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get grasp pose (using current TCP pose as proxy) for the specified arm."""
    try:
        robot_view = env.robots[batch_index].robot_view

        # Get TCP pose as grasp pose proxy
        arm_group_name = f"{self.arm_side}_arm"
        tcp_pose_matrix = robot_view.get_move_group(arm_group_name).leaf_frame_to_robot
        tcp_pose_world = robot_view.base.pose @ tcp_pose_matrix

        return pose_mat_to_7d(tcp_pose_world).astype(np.float32)

    except Exception as e:
        print(f"Warning: Could not get RBY1 grasp pose for {self.arm_side} arm: {e}")
        return np.zeros(7, dtype=np.float32)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

RBY1GraspStateSensor 

RBY1GraspStateSensor(uuid: str = 'rby1_grasp_state', arm_side: str = 'left')

Bases: Sensor

Methods:

Name Description
get_observation
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
arm_side
finger1_name
finger2_name
is_dict bool
obj_name
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/rby1_sensors.py 
126
127
128
129
130
131
132
def __init__(self, uuid: str = "rby1_grasp_state", arm_side: str = "left") -> None:
    self.arm_side = arm_side
    self.finger1_name = f"robot_0/ee_finger_{self.arm_side[0]}1"  # e.g., robot_0/ee_finger_l1
    self.finger2_name = f"robot_0/ee_finger_{self.arm_side[0]}2"  # e.g., robot_0/ee_finger_l2
    self.obj_name = None  # TODO: if pick, pick and place, or open task then object name is self.get_task_objects()["pickup_obj"] if it is door opening it is self.get_task_objects()["door_handle"]
    observation_space = gyms.Box(low=0, high=1, shape=(1,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
arm_side instance-attribute 
arm_side = arm_side
finger1_name instance-attribute 
finger1_name = f'robot_0/ee_finger_{arm_side[0]}1'
finger2_name instance-attribute 
finger2_name = f'robot_0/ee_finger_{arm_side[0]}2'
is_dict class-attribute instance-attribute 
is_dict: bool = False
obj_name instance-attribute 
obj_name = None
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray
Source code in molmo_spaces/env/rby1_sensors.py 
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
160
161
162
163
164
165
166
167
168
169
170
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    # Lazy init
    if self.obj_name is None and task is not None:
        task_objects = task.get_task_objects(batch_index)
        if "door_handle" in task_objects:
            self.obj_name = task_objects["door_handle"]
        elif "pickup_obj" in task_objects:
            self.obj_name = task_objects["pickup_obj"]

    mj_model = env.mj_model
    mj_data = env.mj_datas[batch_index]

    finger1_contact = False
    finger2_contact = False

    # Check all contacts
    for i in range(mj_data.ncon):
        contact = mj_data.contact[i]
        body1_id = mj_model.geom_bodyid[contact.geom1]
        body2_id = mj_model.geom_bodyid[contact.geom2]
        body1_name = mujoco.mj_id2name(mj_model, mujoco.mjtObj.mjOBJ_BODY, body1_id)
        body2_name = mujoco.mj_id2name(mj_model, mujoco.mjtObj.mjOBJ_BODY, body2_id)

        # Check if finger1 is in contact with object
        if (body1_name == self.finger1_name and body2_name == self.obj_name) or (
            body2_name == self.finger1_name and body1_name == self.obj_name
        ):
            finger1_contact = True

        # Check if finger2 is in contact with object
        if (body1_name == self.finger2_name and body2_name == self.obj_name) or (
            body2_name == self.finger2_name and body1_name == self.obj_name
        ):
            finger2_contact = True

    grasping_object = finger1_contact and finger2_contact
    return np.array([int(grasping_object)], dtype=np.uint8)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

RBY1RobotStateSensor 

RBY1RobotStateSensor(uuid: str = 'rby1_robot_state', str_max_len: int = 4000)

Bases: Sensor

Sensor for RBY1 robot joint positions, velocities, and dual end-effector poses.

Methods:

Name Description
get_observation

Get RBY1 robot state observation for a specific environment in the batch.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict bool
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/rby1_sensors.py 
35
36
37
38
39
def __init__(self, uuid: str = "rby1_robot_state", str_max_len: int = 4000) -> None:
    self.str_max_len = str_max_len
    # Use bytes array for HDF5 compatibility
    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get RBY1 robot state observation for a specific environment in the batch.

Source code in molmo_spaces/env/rby1_sensors.py 
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
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get RBY1 robot state observation for a specific environment in the batch."""
    robot = env.robots[batch_index]
    robot_view = robot.robot_view

    # Get joint positions and velocities
    qpos_dict = robot_view.get_qpos_dict()
    # qvel_dict = robot_view.get_qvel_dict()

    # Get dual end-effector poses for RBY1
    left_ee_pose = robot_view.get_move_group("left_arm").leaf_frame_to_robot
    right_ee_pose = robot_view.get_move_group("right_arm").leaf_frame_to_robot

    # Define expected joint groups and their max lengths for RBY1
    expected_joint_groups = {
        "left_arm": 7,  # Left arm (7 joints)
        "right_arm": 7,  # Right arm (7 joints)
        "left_gripper": 2,  # Left gripper (2 joints)
        "right_gripper": 2,  # Right gripper (2 joints)
        "base": 3,  # Base (3 DOF)
        "torso": 6,  # Torso (6 DOF)
        "head": 2,  # Head (2 DOF)
    }

    # Pad qpos to consistent lengths
    padded_qpos = {}
    for group_name, max_length in expected_joint_groups.items():
        if group_name in qpos_dict and qpos_dict[group_name].size > 0:
            actual_data = qpos_dict[group_name]
            padded_array = np.zeros(max_length)
            padded_array[: len(actual_data)] = actual_data[:max_length]  # Truncate if too long
            padded_qpos[group_name] = padded_array.tolist()
        else:
            # Fill with zeros if group doesn't exist
            padded_qpos[group_name] = [0.0] * max_length

    # Create data dict with consistent structure and ordering
    data_dict = {
        "joint_positions": padded_qpos,
        "left_ee_pose": pose_mat_to_7d(left_ee_pose).tolist(),
        "right_ee_pose": pose_mat_to_7d(right_ee_pose).tolist(),
        "timestamp": float(env.mj_datas[batch_index].time),
    }

    # Convert to JSON string, then to bytes
    data_str = json.dumps(data_dict, separators=(",", ":"))
    data_bytes = data_str.encode("utf-8")

    # Pad or truncate to fixed length
    if len(data_bytes) > self.str_max_len:
        data_bytes = data_bytes[: self.str_max_len]
    else:
        data_bytes = data_bytes + b"\x00" * (self.str_max_len - len(data_bytes))

    return np.frombuffer(data_bytes, dtype=np.uint8)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

RBY1TCPPoseSensor 

RBY1TCPPoseSensor(uuid: str = 'rby1_tcp_pose', arm_side: str = 'left')

Bases: Sensor

Sensor for RBY1 TCP (Tool Center Point) poses in 7D format for both arms.

Parameters:

Name Type Description Default
uuid str

Unique identifier for this sensor

 'rby1_tcp_pose'
arm_side str

Which arm to track ("left" or "right")

 'left'

Methods:

Name Description
get_observation

Get TCP pose in world coordinates for the specified arm.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
arm_side
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/rby1_sensors.py 
101
102
103
104
105
106
107
108
109
def __init__(self, uuid: str = "rby1_tcp_pose", arm_side: str = "left") -> None:
    """
    Args:
        uuid: Unique identifier for this sensor
        arm_side: Which arm to track ("left" or "right")
    """
    self.arm_side = arm_side
    observation_space = gyms.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32)
    super().__init__(uuid=uuid, observation_space=observation_space)
arm_side instance-attribute 
arm_side = arm_side
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get TCP pose in world coordinates for the specified arm.

Source code in molmo_spaces/env/rby1_sensors.py 
111
112
113
114
115
116
117
118
119
120
121
122
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get TCP pose in world coordinates for the specified arm."""

    try:
        robot_view = env.robots[batch_index].robot_view
        # Get TCP pose relative to robot base
        gripper_mg_id = f"{self.arm_side}_gripper"
        tcp_pose_matrix = robot_view.get_move_group(gripper_mg_id).leaf_frame_to_robot
        return pose_mat_to_7d(tcp_pose_matrix).astype(np.float32)
    except (AttributeError, KeyError) as e:
        print(f"Warning: Could not get TCP pose: {e}")
        return np.zeros(7, dtype=np.float32)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

get_rby1_door_opening_sensors 

get_rby1_door_opening_sensors(exp_config)

Get core sensors for RBY1 door opening data generation.

Parameters:

Name Type Description Default
exp_config

Experiment configuration object with sensor parameters

 required

Returns: List of initialized sensors

Source code in molmo_spaces/env/rby1_sensors.py 
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
def get_rby1_door_opening_sensors(exp_config):
    """Get core sensors for RBY1 door opening data generation.

    Args:
        exp_config: Experiment configuration object with sensor parameters
    Returns:
        List of initialized sensors
    """
    sensors = []
    for camera_spec in exp_config.camera_config.cameras:
        camera_name = camera_spec.name
        camera_name = camera_name.split("/")[-1]
        cam_params = CameraParameterSensor(
            camera_name=camera_name, uuid=f"sensor_param_{camera_name}"
        )
        sensors.append(cam_params)

        # RGB sensor
        cam_rgb = CameraSensor(
            camera_name=camera_name,
            img_resolution=exp_config.camera_config.img_resolution,
            uuid=camera_name,
        )
        sensors.append(cam_rgb)
        if camera_spec.record_depth:
            cam_depth = DepthSensor(
                camera_name=camera_name,
                img_resolution=exp_config.camera_config.img_resolution,
                uuid=f"{camera_name}_depth",
            )
            sensors.append(cam_depth)

    # Agent data sensors - RBY1 has more joints (dual arm + base + torso + head + grippers)
    sensors.append(RobotJointVelocitySensor(uuid="qvel", max_joints=25))
    sensors.append(RobotJointPositionSensor(uuid="qpos", max_joints=25))
    # Action sensors - RBY1 specific action spec
    sensors.append(
        LastActionSensor(
            dtype=exp_config.task_config.action_dtype,
        )
    )
    # Position action sensors
    sensors.append(LastCommandedJointPosSensor())
    sensors.append(LastCommandedRelativeJointPosSensor())
    sensors.append(LastCommandedEETwistSensor())
    sensors.append(LastCommandedEEPoseSensor())
    # Door opening specific sensors
    sensors.append(DoorStateSensor(uuid="door_state"))

    # Door opening policy phase sensor
    sensors.append(PolicyPhaseSensor(uuid="policy_phase"))

    # TCP poses for both arms (RBY1 is dual-arm) - use separate sensors for each arm
    sensors.append(RBY1TCPPoseSensor(uuid="left_tcp_pose", arm_side="left"))
    sensors.append(RBY1TCPPoseSensor(uuid="right_tcp_pose", arm_side="right"))
    sensors.append(RobotBasePoseSensor(uuid="robot_base_pose"))

    # Grasp state for both arms
    sensors.append(RBY1GraspStateSensor(uuid="rby1_left_grasp_state", arm_side="left"))
    sensors.append(RBY1GraspStateSensor(uuid="rby1_right_grasp_state", arm_side="right"))

    # Environment state sensors
    sensors.append(EnvStateSensor(uuid="env_states"))
    sensors.append(TaskInfoSensor(uuid="task_info"))

    # sensors.append(RBY1RobotStateSensor(uuid="robot_state"))
    # TODO: just make sure the door and handle are in the tracked objects and have this take care of it
    sensors.append(
        ObjectPoseSensor(
            object_names=exp_config.task_config.tracked_object_names or [],
            uuid="object_poses",
        )
    )

    # Object tracking sensor for image points (door handle, etc.)
    sensors.append(ObjectImagePointsSensor(exp_config=exp_config))

    return sensors

sensors

Classes:

Name Description
ActorStateSensor

Sensor for actor (object) states in numerical format.
ArticulationStateSensor

Sensor for articulation (robot) state in numerical format.
DoorStateSensor

Sensor for door state including joint angle and opening percentage.
EnvStateSensor

Sensor for complete MuJoCo environment state.
GraspPoseSensor

Sensor for the planned grasp pose in 7D format.
GraspStateSensor

Sensor for grasp state. For each gripper, track whether it is touching (and possibly holding) the object.
LastActionSensor

Sensor for robot actions (absolute values only).
LastCommandedEEPoseSensor
LastCommandedEETwistSensor
LastCommandedJointPosSensor
LastCommandedRelativeJointPosSensor
ObjectEndPoseSensor

Sensor for target/end object pose in 7D format (x, y, z, qw, qx, qy, qz).
ObjectImagePointsSensor

Sensor for tracking object pixel coordinates across multiple cameras.
ObjectPoseSensor

Sensor for object poses relative to robot base.
ObjectStartPoseSensor

Sensor for initial object pose in 7D format (x, y, z, qw, qx, qy, qz).
PolicyNumRetriesSensor

Sensor for tracking the number of retries of the object manipulation policy.
PolicyPhaseSensor

Sensor for tracking the current phase of a planner policy.
RobotBasePoseSensor

Sensor for robot base pose in 7D format.
RobotJointPositionSensor

Sensor for robot joint positions as numerical array.
RobotJointVelocitySensor

Sensor for robot joint velocities as numerical array.
RobotStateSensor

Sensor for robot joint positions, velocities, and end-effector pose.
TCPPoseSensor

Sensor for TCP (Tool Center Point / End Effector) pose in 7D format.
TaskInfoSensor

Sensor for task information.

Functions:

Name Description
get_core_sensors

Get core sensors for Franka pick-place data generation.
get_nav_task_sensors

Get sensors for navigation to object task.

Attributes:

Name Type Description
log

log module-attribute 

log = getLogger(__name__)

ActorStateSensor 

ActorStateSensor(actor_names: list[str], uuid: str = 'actor_states')

Bases: Sensor

Sensor for actor (object) states in numerical format.

Methods:

Name Description
get_observation

Get actor states as numerical array.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
actor_names
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
466
467
468
469
470
471
472
def __init__(self, actor_names: list[str], uuid: str = "actor_states") -> None:
    self.actor_names = actor_names
    # Each actor has 13D state: pos(3) + quat(4) + vel(6)
    observation_space = gyms.Box(
        low=-np.inf, high=np.inf, shape=(len(actor_names), 13), dtype=np.float32
    )
    super().__init__(uuid=uuid, observation_space=observation_space)
actor_names instance-attribute 
actor_names = actor_names
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get actor states as numerical array.

Source code in molmo_spaces/env/sensors.py 
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get actor states as numerical array."""
    try:
        from molmo_spaces.env.data_views import create_mlspaces_body

        data = env.mj_datas[batch_index]
        actor_states = []

        for actor_name in self.actor_names:
            try:
                body = create_mlspaces_body(data, actor_name)

                # Create 13D state: position(3) + quaternion(4) + velocity(6)
                actor_state = np.concatenate(
                    [
                        body.position,  # 3D position
                        body.quaternion,  # 4D quaternion [w,x,y,z]
                        body.velocities[:6],  # 6D velocity (3 linear + 3 angular)
                    ]
                )
                actor_states.append(actor_state)

            except Exception:
                # Use zeros for missing/invalid actors
                actor_states.append(np.zeros(13))

        return np.array(actor_states, dtype=np.float32)

    except (AttributeError, KeyError) as e:
        print(f"Warning: Could not get actor states: {e}")
        return np.zeros((len(self.actor_names), 13), dtype=np.float32)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

ArticulationStateSensor 

ArticulationStateSensor(state_dim: int = 31, uuid: str = 'articulation_states')

Bases: Sensor

Sensor for articulation (robot) state in numerical format.

Methods:

Name Description
get_observation

Get articulation state as numerical array.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict bool
observation_space Space
state_dim
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
510
511
512
513
def __init__(self, state_dim: int = 31, uuid: str = "articulation_states") -> None:
    self.state_dim = state_dim
    observation_space = gyms.Box(low=-np.inf, high=np.inf, shape=(state_dim,), dtype=np.float32)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
state_dim instance-attribute 
state_dim = state_dim
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get articulation state as numerical array.

Source code in molmo_spaces/env/sensors.py 
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get articulation state as numerical array."""
    try:
        robot_view = env.robots[batch_index].robot_view

        # Collect all joint positions and velocities
        all_qpos = []
        all_qvel = []
        for mg_name in robot_view.move_group_ids():
            mg = robot_view.get_move_group(mg_name)
            all_qpos.extend(mg.joint_pos)
            all_qvel.extend(mg.joint_vel)

        # Combine and pad to target dimension
        combined_state = np.concatenate([all_qpos, all_qvel])
        padded_state = np.zeros(self.state_dim, dtype=np.float32)
        actual_length = min(len(combined_state), self.state_dim)
        padded_state[:actual_length] = combined_state[:actual_length]

        return padded_state

    except (AttributeError, KeyError) as e:
        print(f"Warning: Could not get articulation state: {e}")
        return np.zeros(self.state_dim, dtype=np.float32)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

DoorStateSensor 

DoorStateSensor(uuid: str = 'door_state', str_max_len: int = 1000)

Bases: Sensor

Sensor for door state including joint angle and opening percentage.

Methods:

Name Description
get_observation

Get door state as encoded JSON.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/sensors.py 
703
704
705
706
707
708
709
710
711
712
def __init__(
    self,
    uuid: str = "door_state",
    str_max_len: int = 1000,
) -> None:
    self.str_max_len = str_max_len
    self.is_dict = True
    # Use bytes array for HDF5 compatibility
    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict instance-attribute 
is_dict = True
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray | dict[str, ndarray]

Get door state as encoded JSON.

Source code in molmo_spaces/env/sensors.py 
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
def get_observation(
    self, env, task, batch_index: int = 0, *args, **kwargs
) -> np.ndarray | dict[str, np.ndarray]:
    """Get door state as encoded JSON."""
    try:
        # Get door state from task
        # TODO: Add handle bbox dimensions
        door_state_data = {
            "joint_angle": [0.0],
            "opening_percentage": [0.0],
            "handle_position": [0.0, 0.0, 0.0],
            "handle_extents": [0.0, 0.0, 0.0],
            "door_position": [0.0, 0.0, 0.0],
            "is_open": [False],
        }

        if hasattr(task, "door_object") and task.door_object is not None:
            # Get door joint angle
            if hasattr(task, "current_door_joint_state"):
                door_state_data["joint_angle"] = [float(task.current_door_joint_state.item())]

            # Calculate opening percentage
            if hasattr(task, "exp_config") and hasattr(
                task.config.task_config, "articulated_joint_range"
            ):
                joint_range = task.config.task_config.articulated_joint_range
                current_angle = door_state_data["joint_angle"]
                opening_percentage = (current_angle - joint_range[0]) / (
                    joint_range[1] - joint_range[0]
                )
                if isinstance(opening_percentage, np.ndarray):
                    opening_percentage = opening_percentage.item()
                door_state_data["opening_percentage"] = [
                    float(np.clip(opening_percentage, 0.0, 1.0))
                ]

            # Get door handle position
            if hasattr(task, "get_door_handle_position"):
                try:
                    handle_pos = task.get_door_handle_position()
                    door_state_data["handle_position"] = handle_pos.tolist()
                except AttributeError:
                    pass

            # Get door handle bbox extents
            if hasattr(task, "get_door_handle_extents"):
                try:
                    handle_extents = task.get_door_handle_extents()
                    door_state_data["handle_extents"] = handle_extents.tolist()
                except AttributeError:
                    pass

            # Get door position
            if hasattr(task, "get_door_joint_position"):
                try:
                    door_pos = task.get_door_joint_position()
                    door_state_data["door_position"] = door_pos.tolist()
                except AttributeError:
                    pass

            # Check if door is considered open
            if hasattr(task, "door_opened"):
                door_state_data["is_open"] = [bool(task.door_opened)]

        return door_state_data

    except (AttributeError, KeyError) as e:
        print(f"Warning: Could not get door state: {e}")
        # Return empty structure
        empty_data = {
            "joint_angle": 0.0,
            "opening_percentage": 0.0,
            "handle_position": [0.0, 0.0, 0.0],
            "handle_extents": [0.0, 0.0, 0.0],
            "door_position": [0.0, 0.0, 0.0],
            "is_open": False,
        }
        return empty_data
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

EnvStateSensor 

EnvStateSensor(uuid: str = 'env_states', str_max_len: int = 50000)

Bases: Sensor

Sensor for complete MuJoCo environment state.

Methods:

Name Description
get_observation

Get complete environment state from MuJoCo.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/sensors.py 
381
382
383
384
385
386
def __init__(self, uuid: str = "env_states", str_max_len: int = 50000) -> None:
    self.str_max_len = str_max_len
    self.is_dict = True
    # Use bytes array for HDF5 compatibility
    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict instance-attribute 
is_dict = True
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get complete environment state from MuJoCo.

Source code in molmo_spaces/env/sensors.py 
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get complete environment state from MuJoCo."""
    try:
        import mujoco

        from molmo_spaces.env.data_views import create_mlspaces_body

        model = env.mj_model
        data = env.mj_datas[batch_index]

        # Collect all body states
        actors = {}
        articulations = {}

        # Get robot state first (for articulations)
        robot_view = env.robots[batch_index].robot_view
        all_qpos = []
        all_qvel = []
        for mg_name in robot_view.move_group_ids():
            mg = robot_view.get_move_group(mg_name)
            all_qpos.extend(mg.joint_pos)
            all_qvel.extend(mg.joint_vel)

        # Pad robot state to 31 dimensions
        robot_state = np.zeros(31)
        combined_state = np.concatenate([all_qpos, all_qvel])
        actual_length = min(len(combined_state), 31)
        robot_state[:actual_length] = combined_state[:actual_length]
        articulations["panda"] = robot_state.tolist()

        # Get ALL bodies for actors (not hardcoded)
        for body_id in range(model.nbody):
            body_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, body_id)
            if body_name is None:
                continue

            # Skip robot bodies (they go in articulations)
            if any(
                robot_name in body_name.lower() for robot_name in ["panda", "franka", "robot"]
            ):
                continue

            try:
                body = create_mlspaces_body(data, body_name)

                # Create 13D state: position(3) + quaternion(4) + velocity(6)
                body_state = np.concatenate(
                    [
                        body.position,  # 3D position
                        body.quaternion,  # 4D quaternion [w,x,y,z]
                        body.velocities[:6],  # 6D velocity (3 linear + 3 angular)
                    ]
                )
                actors[body_name] = body_state.tolist()

            except Exception:
                # Skip bodies that can't be processed
                continue

        # Create complete environment state
        env_state_data = {"actors": actors, "articulations": articulations}

        # Convert to JSON string with consistent formatting
        return env_state_data

    except (AttributeError, KeyError) as e:
        print(f"Warning: Could not get environment state: {e}")
        # Return empty structure
        empty_data = {
            "actors": {},
            "articulations": {"panda": np.zeros(31).tolist()},
        }
        return empty_data
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

GraspPoseSensor 

GraspPoseSensor(uuid: str = 'grasp_pose')

Bases: Sensor

Sensor for the planned grasp pose in 7D format.

Methods:

Name Description
get_observation

Get grasp pose (using current TCP pose as proxy).
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
1083
1084
1085
def __init__(self, uuid: str = "grasp_pose") -> None:
    observation_space = gyms.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get grasp pose (using current TCP pose as proxy).

Source code in molmo_spaces/env/sensors.py 
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get grasp pose (using current TCP pose as proxy)."""
    if task._registered_policy is None:
        log.warning("No registered policy, cannot get grasp pose.")
        return np.zeros(7, dtype=np.float32)
    else:
        return np.array(
            pose_mat_to_7d(task._registered_policy.target_poses["grasp"]),
            dtype=np.float32,
        )
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

GraspStateSensor 

GraspStateSensor(object_name: str, uuid: str | None = None, str_max_len: int = 2000)

Bases: Sensor

Sensor for grasp state. For each gripper, track whether it is touching (and possibly holding) the object. The held state is calculated with a heuristic, and only tracks whether the object is only touching the gripper, not whether the object is actually stably supported by the gripper. More sophisticated heuristics may be added in the future.

Methods:

Name Description
get_observation
reset

Attributes:

Name Type Description
is_dict
object_name
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/sensors.py 
551
552
553
554
555
556
557
558
559
560
561
562
def __init__(self, object_name: str, uuid: str | None = None, str_max_len: int = 2000) -> None:
    self.object_name = object_name
    self.str_max_len = str_max_len
    self.is_dict = True
    if uuid is None:
        uuid = f"grasp_state_{object_name}"

    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)

    self._object_geoms: set[int] | None = None
    self._gripper_geoms: dict[str, set[int]] | None = None
is_dict instance-attribute 
is_dict = True
object_name instance-attribute 
object_name = object_name
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs) -> dict
Source code in molmo_spaces/env/sensors.py 
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
def get_observation(
    self, env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs
) -> dict:
    model = env.mj_model
    robot_view = env.robots[batch_index].robot_view

    if self._gripper_geoms is None:
        self._gripper_geoms = {}
        for mg_id in robot_view.get_gripper_movegroup_ids():
            mg = robot_view.get_move_group(mg_id)
            self._gripper_geoms[mg_id] = descendant_geoms(
                env.mj_model, mg.root_body_id, visual_only=False
            )

    if self._object_geoms is None:
        object_body = create_mlspaces_body(env.mj_datas[batch_index], self.object_name)
        self._object_geoms = set(
            descendant_geoms(model, object_body.body_id, visual_only=False)
        )

    held = True
    gripper_touching = {k: False for k in self._gripper_geoms}

    for cid in range(env.mj_datas[batch_index].ncon):
        c = env.mj_datas[batch_index].contact[cid]

        # skip contacts between the object and itself and contacts not involving the object
        if (c.geom1 in self._object_geoms) == (c.geom2 in self._object_geoms):
            continue

        other_geom = c.geom2 if c.geom1 in self._object_geoms else c.geom1
        for gripper_id, gripper_geoms in self._gripper_geoms.items():
            if other_geom in gripper_geoms:
                gripper_touching[gripper_id] = True
                break
        else:
            # object is in contact with a non-gripper geom, so it is not held
            held = False

    grasp_state = {}
    for gripper_id, touching in gripper_touching.items():
        grasp_state[gripper_id] = {
            "touching": touching,
            "held": held and touching,
        }

    return grasp_state
reset 
reset() -> None
Source code in molmo_spaces/env/sensors.py 
564
565
566
def reset(self) -> None:
    self._object_geoms = None
    self._gripper_geoms = None

LastActionSensor 

LastActionSensor(dtype: str = 'float32', uuid: str = 'actions/commanded_action', str_max_len: int = 2000)

Bases: Sensor

Sensor for robot actions (absolute values only).

Parameters:

Name Type Description Default
dtype str

Target dtype for all action components

 'float32'
uuid str

Sensor UUID

 'actions/commanded_action'

Methods:

Name Description
get_observation

Get action dictionary with absolute values.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
dtype
is_dict
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/sensors.py 
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
def __init__(
    self,
    dtype: str = "float32",
    uuid: str = "actions/commanded_action",
    str_max_len: int = 2000,
) -> None:
    """
    Args:
        dtype: Target dtype for all action components
        uuid: Sensor UUID
    """
    self.dtype = getattr(np, dtype)
    self.str_max_len = str_max_len
    self.is_dict = True
    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
dtype instance-attribute 
dtype = getattr(numpy, dtype)
is_dict instance-attribute 
is_dict = True
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> dict[str, list]

Get action dictionary with absolute values.

Source code in molmo_spaces/env/sensors.py 
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> dict[str, list]:
    """Get action dictionary with absolute values."""
    current_action = getattr(task, "last_action", None)

    if current_action is None:
        # Return dummy if no action has been set yet (padding value)
        action_dict = {}
    else:
        action_dict = {}
        for component_name, data in current_action.items():
            if not isinstance(data, np.ndarray):
                data = np.array(data, dtype=self.dtype)
            else:
                data = data.astype(self.dtype)
            if data.ndim == 0:
                data = data.reshape(1)
            action_dict[component_name] = data.tolist()

    return action_dict
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

LastCommandedEEPoseSensor 

LastCommandedEEPoseSensor(uuid: str = 'actions/ee_pose')

Bases: Sensor

Methods:

Name Description
get_observation
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
227
228
229
def __init__(self, uuid: str = "actions/ee_pose") -> None:
    super().__init__(uuid=uuid, observation_space=gyms.Box(0, 255, (1,), dtype=np.uint8))
    self.is_dict = True
is_dict instance-attribute 
is_dict = True
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs)
Source code in molmo_spaces/env/sensors.py 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def get_observation(
    self, env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs
):
    # sentinel done action when the task is terminal
    if task.is_terminal():
        return {}

    robot = env.robots[batch_index]
    prev_cmd_jp = _cmd_joint_pos(robot)
    prev_cmd_poses = robot.kinematics.fk(prev_cmd_jp, np.eye(4), rel_to_base=True)
    prev_cmd_posquats = {
        name: pose_mat_to_7d(pose).tolist() for name, pose in prev_cmd_poses.items()
    }
    return prev_cmd_posquats
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

LastCommandedEETwistSensor 

LastCommandedEETwistSensor(uuid: str = 'actions/ee_twist')

Bases: Sensor

Methods:

Name Description
get_observation
reset

Attributes:

Name Type Description
is_dict
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
248
249
250
251
252
def __init__(self, uuid: str = "actions/ee_twist") -> None:
    super().__init__(uuid=uuid, observation_space=gyms.Box(0, 255, (1,), dtype=np.uint8))
    self.is_dict = True
    self._prev_poses = None
    self._tracked_keys: set[str] | None = None  # move group ids that are position-commanded
is_dict instance-attribute 
is_dict = True
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs)
Source code in molmo_spaces/env/sensors.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
def get_observation(
    self, env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs
):
    # sentinel done action when the task is terminal
    if task.is_terminal():
        return {}

    robot = env.robots[batch_index]
    prev_cmd_jp = _cmd_joint_pos(robot)
    if self._tracked_keys is None:
        self._tracked_keys = set(prev_cmd_jp.keys())
    curr_poses = {
        name: robot.robot_view.get_move_group(name).leaf_frame_to_robot
        for name in self._tracked_keys
    }

    # on the first step, return dummy action
    if self._prev_poses is None:
        self._prev_poses = curr_poses
        return {name: np.zeros(6).tolist() for name in self._tracked_keys}

    prev_cmd_poses = robot.kinematics.fk(prev_cmd_jp, np.eye(4), rel_to_base=True)
    prev_cmd_twists: dict[str, list[float]] = {}
    for name, prev_pose in self._prev_poses.items():
        prev_cmd_pose = prev_cmd_poses[name]
        cmd_twist = np.concatenate(transform_to_twist(np.linalg.inv(prev_pose) @ prev_cmd_pose))
        prev_cmd_twists[name] = cmd_twist.tolist()
    self._prev_poses = curr_poses
    return prev_cmd_twists
reset 
reset() -> None
Source code in molmo_spaces/env/sensors.py 
284
285
def reset(self) -> None:
    self._prev_poses = None

LastCommandedJointPosSensor 

LastCommandedJointPosSensor(uuid: str = 'actions/joint_pos')

Bases: Sensor

Methods:

Name Description
get_observation
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
179
180
181
def __init__(self, uuid: str = "actions/joint_pos") -> None:
    super().__init__(uuid=uuid, observation_space=gyms.Box(0, 255, (1,), dtype=np.uint8))
    self.is_dict = True
is_dict instance-attribute 
is_dict = True
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs)
Source code in molmo_spaces/env/sensors.py 
183
184
185
186
187
188
189
190
def get_observation(
    self, env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs
):
    # sentinel done action when the task is terminal
    if task.is_terminal():
        return {}
    robot = env.robots[batch_index]
    return {k: v.tolist() for k, v in _cmd_joint_pos(robot).items()}
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

LastCommandedRelativeJointPosSensor 

LastCommandedRelativeJointPosSensor(uuid: str = 'actions/joint_pos_rel')

Bases: Sensor

Methods:

Name Description
get_observation
reset

Attributes:

Name Type Description
is_dict
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
194
195
196
197
def __init__(self, uuid: str = "actions/joint_pos_rel") -> None:
    super().__init__(uuid=uuid, observation_space=gyms.Box(0, 255, (1,), dtype=np.uint8))
    self.is_dict = True
    self._prev_jp = None
is_dict instance-attribute 
is_dict = True
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs)
Source code in molmo_spaces/env/sensors.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
def get_observation(
    self, env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs
):
    # sentinel done action when the task is terminal
    if task.is_terminal():
        return {}

    robot = env.robots[batch_index]
    # on the first step, return dummy action
    if self._prev_jp is None:
        self._prev_jp = robot.robot_view.get_qpos_dict()
        return {name: np.zeros_like(jp).tolist() for name, jp in self._prev_jp.items()}

    prev_cmd_jp = _cmd_joint_pos(robot)
    cmd_rel_jp: dict[str, list[float]] = {}
    for name, jp in prev_cmd_jp.items():
        # some move groups (e.g. grippers) have different action and state dimensions,
        # it doesn't make much sense to compute relative actions in these cases.
        if jp.shape == self._prev_jp[name].shape:
            cmd_rel_jp[name] = (jp - self._prev_jp[name]).tolist()
    self._prev_jp = robot.robot_view.get_qpos_dict()
    return cmd_rel_jp
reset 
reset() -> None
Source code in molmo_spaces/env/sensors.py 
222
223
def reset(self) -> None:
    self._prev_jp = None

ObjectEndPoseSensor 

ObjectEndPoseSensor(object_name: str, uuid: str | None = None)

Bases: Sensor

Sensor for target/end object pose in 7D format (x, y, z, qw, qx, qy, qz).

Methods:

Name Description
get_observation

Get target object pose.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict bool
object_name
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
680
681
682
683
684
685
686
def __init__(self, object_name: str, uuid: str | None = None) -> None:
    self.object_name = object_name
    if uuid is None:
        uuid = f"obj_end_{object_name}"

    observation_space = gyms.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict class-attribute instance-attribute 
is_dict: bool = False
object_name instance-attribute 
object_name = object_name
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get target object pose.

Source code in molmo_spaces/env/sensors.py 
688
689
690
691
692
693
694
695
696
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get target object pose."""
    if task.config.task_type in ["pick", "open", "close"]:
        goal_pose = np.array(task.config.task_config.pickup_obj_goal_pose, dtype=np.float32)
        return goal_pose
    else:
        # TODO(max): fix this
        goal_pose = np.zeros(7, dtype=np.float32)
        return goal_pose
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

ObjectImagePointsSensor 

ObjectImagePointsSensor(exp_config, camera_names: list[str] | None = None, uuid: str = 'object_image_points', max_points: int = 10, erosion_iterations: int = 2, exclude_follower_cameras: bool = True)

Bases: Sensor

Sensor for tracking object pixel coordinates across multiple cameras.

Detects task objects in camera views and returns sampled pixel coordinates normalized to 0-1 range. Returns HDF5-native numpy arrays (not JSON).

Output structure (nested dict of numpy arrays): { "": { "": { "points": np.ndarray (max_points, 2) - normalized (x, y) coords, NaN-padded "num_points": np.ndarray (1,) - number of valid points }, ... }, ... }

Parameters:

Name Type Description Default
exp_config

Experiment configuration with camera_config

 required
camera_names list[str] | None

Optional list of camera names to track. If None, uses all cameras from config.

 None
uuid str

Unique sensor identifier

 'object_image_points'
max_points int

Maximum number of points to sample per camera

 10
erosion_iterations int

Number of erosion iterations for segmentation mask

 2
exclude_follower_cameras bool

If True, exclude cameras with 'follower' in their name

 True

Methods:

Name Description
get_observation

Get pixel coordinates of task objects in each camera view.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
camera_names
camera_specs
erosion_iterations
is_dict
max_points
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
def __init__(
    self,
    exp_config,
    camera_names: list[str] | None = None,
    uuid: str = "object_image_points",
    max_points: int = 10,
    erosion_iterations: int = 2,
    exclude_follower_cameras: bool = True,
) -> None:
    """
    Args:
        exp_config: Experiment configuration with camera_config
        camera_names: Optional list of camera names to track. If None, uses all cameras from config.
        uuid: Unique sensor identifier
        max_points: Maximum number of points to sample per camera
        erosion_iterations: Number of erosion iterations for segmentation mask
        exclude_follower_cameras: If True, exclude cameras with 'follower' in their name
    """
    # Build camera spec lookup from config
    all_camera_specs = {
        camera_spec.name: camera_spec for camera_spec in exp_config.camera_config.cameras
    }

    # Filter to requested cameras or use all
    if camera_names is not None:
        self.camera_specs = {}
        for cam_name in camera_names:
            if cam_name not in all_camera_specs:
                raise ValueError(
                    f"Camera '{cam_name}' not found in camera config. Available cameras: {list(all_camera_specs.keys())}"
                )
            self.camera_specs[cam_name] = all_camera_specs[cam_name]
    else:
        self.camera_specs = all_camera_specs

    # Exclude follower cameras (they follow robot/workspace and aren't useful for object tracking)
    if exclude_follower_cameras:
        self.camera_specs = {
            name: spec
            for name, spec in self.camera_specs.items()
            if "follower" not in name.lower()
        }

    self.camera_names = list(self.camera_specs.keys())
    self.img_width, self.img_height = exp_config.camera_config.img_resolution
    self.max_points = max_points
    self.erosion_iterations = erosion_iterations

    # Use nested dict structure - NOT is_dict (which triggers JSON serialization)
    # The nested dict will be handled natively by batch_observations and save_utils
    self.is_dict = False

    # Observation space is a nested dict structure (for reference, not enforced)
    # Each camera produces (max_points, 2) points array + (1,) num_points array
    observation_space = gyms.Dict({})  # Dynamic based on task objects
    super().__init__(uuid=uuid, observation_space=observation_space)
camera_names instance-attribute 
camera_names = list(keys())
camera_specs instance-attribute 
camera_specs = {}
erosion_iterations instance-attribute 
erosion_iterations = erosion_iterations
is_dict instance-attribute 
is_dict = False
max_points instance-attribute 
max_points = max_points
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> dict[str, dict[str, dict[str, ndarray]]]

Get pixel coordinates of task objects in each camera view.

All points are normalized to 0-1 range based on image resolution. Returns HDF5-native numpy arrays (not JSON strings).

Returns:

Type Description
dict[str, dict[str, dict[str, ndarray]]]

Dictionary with structure:
dict[str, dict[str, dict[str, ndarray]]]

{ "": { "": { "points": np.ndarray (max_points, 2) - normalized (x, y), NaN-padded "num_points": np.ndarray (1,) - count of valid points }, ... }, ...
dict[str, dict[str, dict[str, ndarray]]]

}
dict[str, dict[str, dict[str, ndarray]]]

Where object_key is determined by task.get_task_objects() (e.g., 'pickup_obj',
dict[str, dict[str, dict[str, ndarray]]]

'place_receptacle', 'door_handle').
Source code in molmo_spaces/env/sensors.py 
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
def get_observation(
    self,
    env,
    task,
    batch_index: int = 0,
    *args,
    **kwargs,
) -> dict[str, dict[str, dict[str, np.ndarray]]]:
    """Get pixel coordinates of task objects in each camera view.

    All points are normalized to 0-1 range based on image resolution.
    Returns HDF5-native numpy arrays (not JSON strings).

    Returns:
        Dictionary with structure:
        {
            "<object_key>": {
                "<camera_name>": {
                    "points": np.ndarray (max_points, 2) - normalized (x, y), NaN-padded
                    "num_points": np.ndarray (1,) - count of valid points
                },
                ...
            },
            ...
        }
        Where object_key is determined by task.get_task_objects() (e.g., 'pickup_obj',
        'place_receptacle', 'door_handle').
    """
    # Get object names from task's get_task_objects() method (preferred)
    object_names = {}
    if hasattr(task, "get_task_objects"):
        object_names = task.get_task_objects(batch_index)

    # Fall back to legacy config-based lookup for backward compatibility
    if not object_names:
        if hasattr(task, "config") and hasattr(task.config, "task_config"):
            task_config = task.config.task_config
            if hasattr(task_config, "pickup_obj_name") and task_config.pickup_obj_name:
                object_names["pickup_obj"] = task_config.pickup_obj_name
            if (
                hasattr(task_config, "place_receptacle_name")
                and task_config.place_receptacle_name
            ):
                object_names["place_receptacle"] = task_config.place_receptacle_name

    # Initialize result dict with empty numpy arrays for all object/camera combinations
    result = {
        obj_key: {camera: self._create_empty_points_data() for camera in self.camera_names}
        for obj_key in object_names
    }

    # If no objects found, return empty results
    if not object_names:
        log.warning("No task objects found from get_task_objects() or task config")
        return result

    # Check if environment supports segmentation masks
    if not hasattr(env, "get_segmentation_mask_of_object"):
        log.warning("Environment does not support segmentation masks")
        return result

    # Process each object
    for obj_key, obj_name in object_names.items():
        for camera_name in self.camera_names:
            try:
                # Get segmentation mask for the target object
                segmentation_mask = env.get_segmentation_mask_of_object(
                    obj_name, camera_name=camera_name, batch_index=batch_index
                )

                if segmentation_mask is None or not np.any(segmentation_mask > 0):
                    # Keep empty result (already initialized with NaN/0)
                    continue

                # Skip erosion for thin objects (door handles, grippers) or cameras configured to skip
                camera_spec = self.camera_specs.get(camera_name)
                skip_erosion = "handle" in obj_key or (camera_spec and camera_spec.skip_erosion)
                if skip_erosion:
                    mask_to_use = segmentation_mask
                else:
                    # Erode mask to get more stable interior points
                    eroded_mask = erode_segmentation_mask(
                        segmentation_mask, iterations=self.erosion_iterations
                    )
                    # Fall back to original mask if erosion removed all points
                    if eroded_mask is not None and np.any(eroded_mask > 0):
                        mask_to_use = eroded_mask
                    else:
                        mask_to_use = segmentation_mask

                # Find the points where the object is visible
                if mask_to_use is not None and np.any(mask_to_use > 0):
                    points = np.argwhere(mask_to_use > 0)

                    # Sample random subset up to max_points
                    num_points = min(len(points), self.max_points)
                    if len(points) > self.max_points:
                        indices = np.random.choice(len(points), self.max_points, replace=False)
                        points = points[indices]

                    # Switch from (row, col) to (x, y) format
                    switched_points = points[:, [1, 0]].astype(np.float32)

                    # Check if this specific camera is warped
                    is_warped = camera_spec.is_warped if camera_spec else False

                    # Get distortion map if camera is warped
                    distortion_map = None
                    if is_warped:
                        raise NotImplementedError(
                            "Distortion map not implemented - what are you doing here?"
                        )

                    # Normalize points (handles both distortion correction and 0-1 normalization)
                    normalized_points = normalize_points(
                        switched_points,
                        self.img_width,
                        self.img_height,
                        distortion_map=distortion_map,
                    )

                    # Round to 4 decimal places
                    rounded_points = np.round(normalized_points, 4)

                    # Create padded array (NaN for unused slots)
                    points_array = np.full((self.max_points, 2), np.nan, dtype=np.float32)
                    points_array[:num_points] = rounded_points

                    result[obj_key][camera_name] = {
                        "points": points_array,
                        "num_points": np.array([num_points], dtype=np.int32),
                    }

            except NotImplementedError:
                log.warning(
                    f"Segmentation mask retrieval not yet implemented for {camera_name}"
                )
                # Keep empty result
            except Exception as e:
                log.exception(
                    f"Error processing camera {camera_name} for object {obj_name}: {e}"
                )
                # Keep empty result

    return result
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

ObjectPoseSensor 

ObjectPoseSensor(object_names: list[str], uuid: str = 'object_poses', str_max_len: int = 2000)

Bases: Sensor

Sensor for object poses relative to robot base.

Methods:

Name Description
get_observation

Get object poses relative to robot base for a specific environment.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict
object_names
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/sensors.py 
143
144
145
146
147
148
149
150
151
152
153
154
def __init__(
    self,
    object_names: list[str],
    uuid: str = "object_poses",
    str_max_len: int = 2000,
) -> None:
    self.object_names = object_names
    self.str_max_len = str_max_len
    self.is_dict = True
    # Use bytes array for HDF5 compatibility
    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict instance-attribute 
is_dict = True
object_names instance-attribute 
object_names = object_names
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get object poses relative to robot base for a specific environment.

Source code in molmo_spaces/env/sensors.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get object poses relative to robot base for a specific environment."""
    data = env.mj_datas[batch_index]
    robot_view = env.robots[batch_index].robot_view

    # Ensure consistent structure - always include all tracked objects
    object_poses = {}
    valid_objects = [x for x in self.object_names if x is not None]
    for obj_name in sorted(valid_objects):  # Sort for consistent ordering
        # try:
        obj_body = create_mlspaces_body(data, obj_name)
        # Get pose relative to robot base
        obj_pose_rel = np.linalg.inv(robot_view.base.pose) @ obj_body.pose
        object_poses[obj_name] = obj_pose_rel.tolist()
        # except (AttributeError, KeyError) as e:
        #    # Always use identity matrix for missing objects (consistent structure)
        #    object_poses[obj_name] = np.eye(4).tolist()

    # Convert to JSON string with consistent formatting
    return object_poses
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

ObjectStartPoseSensor 

ObjectStartPoseSensor(object_name: str, uuid: str | None = None)

Bases: Sensor

Sensor for initial object pose in 7D format (x, y, z, qw, qx, qy, qz).

Methods:

Name Description
get_observation

Get initial object pose.
reset

Reset stored initial pose.

Attributes:

Name Type Description
is_dict bool
object_name
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
646
647
648
649
650
651
652
653
654
655
def __init__(self, object_name: str, uuid: str | None = None) -> None:
    self.object_name = object_name
    if uuid is None:
        uuid = f"obj_start_{object_name}"

    observation_space = gyms.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32)
    super().__init__(uuid=uuid, observation_space=observation_space)

    # Store initial pose when first called
    self._initial_pose = None
is_dict class-attribute instance-attribute 
is_dict: bool = False
object_name instance-attribute 
object_name = object_name
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get initial object pose.

Source code in molmo_spaces/env/sensors.py 
657
658
659
660
661
662
663
664
665
666
667
668
669
670
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get initial object pose."""
    if task.config.task_type in ["pick", "open", "close"]:
        return np.array(task.config.task_config.pickup_obj_start_pose, dtype=np.float32)
    else:
        try:
            data = env.mj_datas[batch_index]
            # Store initial pose on first call (typically during reset)
            if self._initial_pose is None:
                obj_body = create_mlspaces_body(data, self.object_name)
                self._initial_pose = pose_mat_to_7d(obj_body.pose)
            return self._initial_pose.astype(np.float32)
        except (AttributeError, KeyError) as e:
            raise ValueError(f"Could not get initial pose for object {self.object_name}") from e
reset 
reset() -> None

Reset stored initial pose.

Source code in molmo_spaces/env/sensors.py 
672
673
674
def reset(self) -> None:
    """Reset stored initial pose."""
    self._initial_pose = None

PolicyNumRetriesSensor 

PolicyNumRetriesSensor(uuid: str = 'policy_num_retries')

Bases: Sensor

Sensor for tracking the number of retries of the object manipulation policy.

Methods:

Name Description
get_observation

Return the number of retries of the object manipulation policy.
reset

Attributes:

Name Type Description
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
1054
1055
1056
1057
def __init__(self, uuid: str = "policy_num_retries") -> None:
    observation_space = gyms.Box(low=0, high=255, shape=(1,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
    self._logged_warning = False
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> int

Return the number of retries of the object manipulation policy.

Source code in molmo_spaces/env/sensors.py 
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> int:
    """Return the number of retries of the object manipulation policy."""
    policy = task._registered_policy
    if policy is None:
        if not self._logged_warning:
            log.warning("No registered policy, cannot track retries.")
            self._logged_warning = True
        return 0
    elif hasattr(policy, "retry_count"):
        return policy.retry_count
    else:
        if not self._logged_warning:
            log.warning(f"Policy {type(policy)} does not support tracking retries.")
            self._logged_warning = True
        return 0
reset 
reset() -> None
Source code in molmo_spaces/env/sensors.py 
1075
1076
1077
def reset(self) -> None:
    super().reset()
    self._logged_warning = False

PolicyPhaseSensor 

PolicyPhaseSensor(uuid: str = 'policy_phase')

Bases: Sensor

Sensor for tracking the current phase of a planner policy.

Methods:

Name Description
get_observation

Return the current phase of the policy as a string encoded as uint8 array.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
1030
1031
1032
def __init__(self, uuid: str = "policy_phase") -> None:
    observation_space = gyms.Box(low=0, high=255, shape=(1,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> int

Return the current phase of the policy as a string encoded as uint8 array.

Source code in molmo_spaces/env/sensors.py 
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> int:
    """Return the current phase of the policy as a string encoded as uint8 array."""
    if task._registered_policy is not None:
        phase_name = task._registered_policy.get_phase()
        all_phases = task._registered_policy.get_all_phases()
        try:
            phase_num = all_phases[phase_name]
        except KeyError:
            log.warning(f"Unknown phase {phase_name}, options are {all_phases.keys()}")
            phase_num = -1
    else:
        log.warning("No registered policy, cannot get policy phase. Using default phase -1.")
        phase_num = -1

    return int(phase_num)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

RobotBasePoseSensor 

RobotBasePoseSensor(uuid: str = 'robot_base_pose')

Bases: Sensor

Sensor for robot base pose in 7D format.

Methods:

Name Description
get_observation

Get robot base pose.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
123
124
125
def __init__(self, uuid: str = "robot_base_pose") -> None:
    observation_space = gyms.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get robot base pose.

Source code in molmo_spaces/env/sensors.py 
127
128
129
130
131
132
133
134
135
136
137
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get robot base pose."""
    try:
        robot_view = env.robots[batch_index].robot_view
        base_pose = robot_view.base.pose

        return pose_mat_to_7d(base_pose).astype(np.float32)

    except (AttributeError, KeyError) as e:
        print(f"Warning: Could not get robot base pose: {e}")
        return np.zeros(7, dtype=np.float32)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

RobotJointPositionSensor 

RobotJointPositionSensor(uuid: str = 'qpos', max_joints: int = 9, str_max_len: int = 2000)

Bases: Sensor

Sensor for robot joint positions as numerical array.

Methods:

Name Description
get_observation

Get robot joint positions as numerical array.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict
max_joints
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/sensors.py 
332
333
334
335
336
337
def __init__(self, uuid: str = "qpos", max_joints: int = 9, str_max_len: int = 2000) -> None:
    self.max_joints = max_joints
    self.str_max_len = str_max_len
    self.is_dict = True
    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict instance-attribute 
is_dict = True
max_joints instance-attribute 
max_joints = max_joints
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray | dict[str, ndarray]

Get robot joint positions as numerical array.

Source code in molmo_spaces/env/sensors.py 
339
340
341
342
343
344
345
346
347
348
349
350
def get_observation(
    self, env, task, batch_index: int = 0, *args, **kwargs
) -> np.ndarray | dict[str, np.ndarray]:
    """Get robot joint positions as numerical array."""
    robot = env.robots[batch_index]
    robot_view = robot.robot_view

    # Get joint positions
    qpos_dict = robot_view.get_qpos_dict()

    qpos_data = {k: v.astype(np.float32).tolist() for k, v in qpos_dict.items()}
    return qpos_data
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

RobotJointVelocitySensor 

RobotJointVelocitySensor(uuid: str = 'qvel', max_joints: int = 9, str_max_len: int = 2000)

Bases: Sensor

Sensor for robot joint velocities as numerical array.

Methods:

Name Description
get_observation

Get robot joint velocities as numerical array.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict
max_joints
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/sensors.py 
356
357
358
359
360
361
def __init__(self, uuid: str = "qvel", max_joints: int = 9, str_max_len: int = 2000) -> None:
    self.max_joints = max_joints
    self.str_max_len = str_max_len
    self.is_dict = True
    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict instance-attribute 
is_dict = True
max_joints instance-attribute 
max_joints = max_joints
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray | dict[str, ndarray]

Get robot joint velocities as numerical array.

Source code in molmo_spaces/env/sensors.py 
363
364
365
366
367
368
369
370
371
372
373
374
def get_observation(
    self, env, task, batch_index: int = 0, *args, **kwargs
) -> np.ndarray | dict[str, np.ndarray]:
    """Get robot joint velocities as numerical array."""
    robot = env.robots[batch_index]
    robot_view = robot.robot_view

    # Get joint velocities
    qvel_dict = robot_view.get_qvel_dict()

    qvel_data = {k: v.astype(np.float32).tolist() for k, v in qvel_dict.items()}
    return qvel_data
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

RobotStateSensor 

RobotStateSensor(uuid: str = 'robot_state', str_max_len: int = 2000)

Bases: Sensor

Sensor for robot joint positions, velocities, and end-effector pose.

Methods:

Name Description
get_observation

Get robot state observation for a specific environment in the batch.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/sensors.py 
47
48
49
50
51
52
def __init__(self, uuid: str = "robot_state", str_max_len: int = 2000) -> None:
    self.str_max_len = str_max_len
    self.is_dict = True
    # Use bytes array for HDF5 compatibility
    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict instance-attribute 
is_dict = True
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get robot state observation for a specific environment in the batch.

Source code in molmo_spaces/env/sensors.py 
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
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get robot state observation for a specific environment in the batch."""
    robot = env.robots[batch_index]
    robot_view = robot.robot_view

    # Get joint positions and velocities
    qpos_dict = robot_view.get_qpos_dict()
    # qvel_dict = robot_view.get_qvel_dict()

    # Get end-effector pose
    gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0]
    ee_pose = robot_view.get_move_group(gripper_mg_id).leaf_frame_to_robot

    # Ensure consistent structure with fixed-length arrays
    # Define expected joint groups and their max lengths for padding
    expected_joint_groups = {
        "arm": 7,  # Typical arm has 7 joints
        "gripper": 2,  # Typical gripper has 2 joints
        "base": 3,  # Base might have 3 DOF
    }

    # Pad qpos to consistent lengths
    padded_qpos = {}
    for group_name, max_length in expected_joint_groups.items():
        if group_name in qpos_dict and qpos_dict[group_name].size > 0:
            actual_data = qpos_dict[group_name]
            padded_array = np.zeros(max_length)
            padded_array[: len(actual_data)] = actual_data[:max_length]  # Truncate if too long
            padded_qpos[group_name] = padded_array.tolist()
        else:
            # Fill with zeros if group doesn't exist
            padded_qpos[group_name] = [0.0] * max_length

    # Create data dict with consistent structure and ordering
    data = {
        "qpos": padded_qpos,
        # "qvel": padded_qvel,  # Add when needed
        "ee_pose": ee_pose.tolist(),
    }

    # Convert to JSON string and then to bytes with consistent formatting
    return data
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

TCPPoseSensor 

TCPPoseSensor(uuid: str = 'tcp_pose')

Bases: Sensor

Sensor for TCP (Tool Center Point / End Effector) pose in 7D format.

Methods:

Name Description
get_observation

Get TCP pose in robot frame.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors.py 
101
102
103
def __init__(self, uuid: str = "tcp_pose") -> None:
    observation_space = gyms.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get TCP pose in robot frame.

Source code in molmo_spaces/env/sensors.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get TCP pose in robot frame."""
    try:
        robot_view = env.robots[batch_index].robot_view
        # Get TCP pose relative to robot base
        gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0]
        tcp_pose_matrix = robot_view.get_move_group(gripper_mg_id).leaf_frame_to_robot

        return pose_mat_to_7d(tcp_pose_matrix).astype(np.float32)

    except (AttributeError, KeyError) as e:
        print(f"Warning: Could not get TCP pose: {e}")
        return np.zeros(7, dtype=np.float32)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

TaskInfoSensor 

TaskInfoSensor(uuid: str = 'task_info', str_max_len: int = 4000)

Bases: Sensor

Sensor for task information.

Methods:

Name Description
get_observation
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
is_dict
observation_space Space
str_max_len
uuid str
Source code in molmo_spaces/env/sensors.py 
620
621
622
623
624
def __init__(self, uuid: str = "task_info", str_max_len: int = 4000) -> None:
    self.str_max_len = str_max_len
    self.is_dict = True
    observation_space = gyms.Box(low=0, high=255, shape=(str_max_len,), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
is_dict instance-attribute 
is_dict = True
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len instance-attribute 
str_max_len = str_max_len
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs) -> dict
Source code in molmo_spaces/env/sensors.py 
635
636
637
638
639
640
def get_observation(
    self, env: BaseMujocoEnv, task: BaseMujocoTask, batch_index: int = 0, *args, **kwargs
) -> dict:
    info = task.get_info()[batch_index]
    self._sanitize(info)
    return info
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

get_core_sensors 

get_core_sensors(exp_config)

Get core sensors for Franka pick-place data generation.

Parameters:

Name Type Description Default
exp_config

Experiment configuration object with sensor parameters

 required

Returns:

Type Description

List of initialized sensors
Source code in molmo_spaces/env/sensors.py 
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
def get_core_sensors(exp_config):
    """Get core sensors for Franka pick-place data generation.

    Args:
        exp_config: Experiment configuration object with sensor parameters

    Returns:
        List of initialized sensors
    """
    sensors = []

    # Get camera names dynamically from camera config instead of hardcoded list
    for camera_spec in exp_config.camera_config.cameras:
        camera_name = camera_spec.name

        # Camera parameter sensor
        cam_params = CameraParameterSensor(
            camera_name=camera_name, uuid=f"sensor_param_{camera_name}"
        )
        sensors.append(cam_params)

        # RGB sensor
        cam_rgb = CameraSensor(
            camera_name=camera_name,
            img_resolution=exp_config.camera_config.img_resolution,
            uuid=camera_name,
        )
        sensors.append(cam_rgb)

        # Depth sensor (conditional based on camera config)
        if camera_spec.record_depth:
            cam_depth = DepthSensor(
                camera_name=camera_name,
                img_resolution=exp_config.camera_config.img_resolution,
                uuid=f"{camera_name}_depth",
            )
            sensors.append(cam_depth)

        # # Segmentation sensor
        # cam_seg = SegmentationSensor(camera_name=camera_name, img_resolution=config.img_resolution, uuid=f"{camera_name}_seg")
        # sensors.append(cam_seg)

    # Robot State sensors
    sensors.append(RobotJointPositionSensor(uuid="qpos", max_joints=9))
    sensors.append(RobotJointVelocitySensor(uuid="qvel", max_joints=9))
    sensors.append(TCPPoseSensor(uuid="tcp_pose"))
    sensors.append(RobotBasePoseSensor(uuid="robot_base_pose"))

    # Environment state sensors
    sensors.append(EnvStateSensor(uuid="env_states"))

    # Task pose sensors
    sensors.append(
        ObjectStartPoseSensor(
            object_name=exp_config.task_config.pickup_obj_name, uuid="obj_start_pose"
        )
    )
    sensors.append(
        ObjectEndPoseSensor(
            object_name=exp_config.task_config.place_target_name, uuid="obj_end_pose"
        )
    )
    sensors.append(
        GraspStateSensor(
            object_name=exp_config.task_config.pickup_obj_name,
            uuid="grasp_state_pickup_obj",
        )
    )
    # TODO: this kind of hacky hardcoded conditionals should be refactored.
    # Tasks should register their own task-specific sensors.
    if (
        hasattr(exp_config.task_config, "place_receptacle_name")
        and exp_config.task_config.place_receptacle_name
    ):
        sensors.append(
            GraspStateSensor(
                object_name=exp_config.task_config.place_receptacle_name,
                uuid="grasp_state_place_receptacle",
            )
        )
    if isinstance(exp_config.robot_config, RBY1Config):
        from molmo_spaces.env.rby1_sensors import RBY1GraspStateSensor

        sensors.append(RBY1GraspStateSensor(uuid="rby1_left_grasp_state", arm_side="left"))
        sensors.append(RBY1GraspStateSensor(uuid="rby1_right_grasp_state", arm_side="right"))

    sensors.append(TaskInfoSensor(uuid="task_info"))

    sensors.append(GraspPoseSensor(uuid="grasp_pose"))

    # Policy sensors
    sensors.append(PolicyPhaseSensor(uuid="policy_phase"))
    sensors.append(PolicyNumRetriesSensor(uuid="policy_num_retries"))

    # Action sensors
    sensors.append(
        LastActionSensor(
            dtype=exp_config.task_config.action_dtype,
        )
    )
    sensors.append(LastCommandedJointPosSensor())
    sensors.append(LastCommandedRelativeJointPosSensor())
    sensors.append(LastCommandedEETwistSensor())
    sensors.append(LastCommandedEEPoseSensor())

    # Object tracking sensors
    sensors.append(ObjectImagePointsSensor(exp_config=exp_config))

    # Legacy sensors for debugging
    sensors.append(RobotStateSensor(uuid="robot_state"))
    sensors.append(
        ObjectPoseSensor(
            object_names=exp_config.task_config.tracked_object_names
            or [exp_config.task_config.pickup_obj_name, exp_config.task_config.place_target_name],
            uuid="object_poses",
        )
    )

    return sensors

get_nav_task_sensors 

get_nav_task_sensors(exp_config)

Get sensors for navigation to object task.

Parameters:

Name Type Description Default
exp_config

Experiment configuration object with sensor parameters

 required

Returns:

Type Description

List of initialized sensors
Source code in molmo_spaces/env/sensors.py 
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
def get_nav_task_sensors(exp_config):
    """Get sensors for navigation to object task.

    Args:
        exp_config: Experiment configuration object with sensor parameters

    Returns:
        List of initialized sensors
    """
    sensors = []

    # Get camera names dynamically from camera config instead of hardcoded list
    for camera_spec in exp_config.camera_config.cameras:
        camera_name = camera_spec.name

        # Camera parameter sensor
        cam_params = CameraParameterSensor(
            camera_name=camera_name, uuid=f"sensor_param_{camera_name}"
        )
        sensors.append(cam_params)

        # RGB sensor
        cam_rgb = CameraSensor(
            camera_name=camera_name,
            img_resolution=exp_config.camera_config.img_resolution,
            uuid=camera_name,
        )
        sensors.append(cam_rgb)
    # Robot State sensors
    sensors.append(RobotBasePoseSensor(uuid="robot_base_pose"))
    sensors.append(RobotJointPositionSensor(uuid="qpos", max_joints=25))

    # Environment state sensors
    sensors.append(EnvStateSensor(uuid="env_states"))

    # Task pose sensors - determine pickup object names
    if (
        hasattr(exp_config.task_config, "pickup_obj_candidates")
        and exp_config.task_config.pickup_obj_candidates is not None
        and len(exp_config.task_config.pickup_obj_candidates) > 0
    ):
        pickup_obj_names = exp_config.task_config.pickup_obj_candidates
    elif exp_config.task_config.pickup_obj_name:
        pickup_obj_names = [exp_config.task_config.pickup_obj_name]
    else:
        pickup_obj_names = []

    if pickup_obj_names:  # Only add sensor if there are objects to track
        sensors.append(ObjectPoseSensor(object_names=pickup_obj_names, uuid="pickup_obj_pose"))

    # Action sensors
    sensors.append(
        LastActionSensor(
            dtype=exp_config.task_config.action_dtype,
        )
    )

    return sensors

sensors_cameras

Classes:

Name Description
CameraParameterSensor

Sensor for camera parameters (intrinsics and extrinsics).
CameraSensor

Sensor for RGB camera images from MuJoCo.
DepthSensor

Sensor for depth images from MuJoCo.
SegmentationSensor

Sensor for segmentation images from MuJoCo, outputs video-compatible arrays.

CameraParameterSensor 

CameraParameterSensor(camera_name: str = 'camera', uuid: str | None = None, img_resolution: tuple[int, int] = (480, 480))

Bases: Sensor

Sensor for camera parameters (intrinsics and extrinsics).

Methods:

Name Description
get_observation

Get camera parameters for a specific environment.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
camera_name
img_resolution
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors_cameras.py 
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
def __init__(
    self,
    camera_name: str = "camera",
    uuid: str | None = None,
    img_resolution: tuple[int, int] = (480, 480),
) -> None:
    self.img_resolution = img_resolution
    self.camera_name = camera_name

    if uuid is None:
        uuid = f"camera_params_{camera_name}"

    observation_space = gyms.Dict(
        {
            "extrinsic_cv": gyms.Box(low=-np.inf, high=np.inf, shape=(3, 4), dtype=np.float32),
            "cam2world_gl": gyms.Box(low=-np.inf, high=np.inf, shape=(4, 4), dtype=np.float32),
            "intrinsic_cv": gyms.Box(low=-np.inf, high=np.inf, shape=(3, 3), dtype=np.float32),
        }
    )
    super().__init__(uuid=uuid, observation_space=observation_space)
camera_name instance-attribute 
camera_name = camera_name
img_resolution instance-attribute 
img_resolution = img_resolution
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> dict

Get camera parameters for a specific environment.

Source code in molmo_spaces/env/sensors_cameras.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> dict:
    """Get camera parameters for a specific environment."""
    camera = env.camera_manager.registry[self.camera_name]
    world2cam = camera.get_pose()
    # Create extrinsic_cv (Computer Vision convention - world2cam)
    extrinsic_cv = np.linalg.inv(world2cam)[:3, :]  # 3x4 matrix
    cam2world_gl = world2cam

    height, width = self.img_resolution
    fovy_degrees = camera.fov

    # Convert field of view to focal length
    focal_length = (height / 2.0) / np.tan(np.radians(fovy_degrees / 2.0))

    # Create intrinsic matrix (assuming square pixels and centered principal point)
    intrinsic_cv = np.array(
        [[focal_length, 0, width / 2.0], [0, focal_length, height / 2.0], [0, 0, 1]],
        dtype=np.float32,
    )

    # Ensure consistent structure and ordering
    data = {
        "cam2world_gl": cam2world_gl.tolist(),
        "extrinsic_cv": extrinsic_cv.tolist(),
        "intrinsic_cv": intrinsic_cv.tolist(),
    }
    return data
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

CameraSensor 

CameraSensor(camera_name: str = 'camera', img_resolution: tuple[int, int] = (480, 480), uuid: str | None = None)

Bases: Sensor

Sensor for RGB camera images from MuJoCo.

Methods:

Name Description
get_observation

Get camera image from environment rendering.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
camera_name
img_resolution
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors_cameras.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def __init__(
    self,
    camera_name: str = "camera",
    img_resolution: tuple[int, int] = (480, 480),
    uuid: str | None = None,
) -> None:
    self.camera_name = camera_name
    self.img_resolution = img_resolution

    if uuid is None:
        uuid = f"camera_{camera_name}"

    # Define observation space for RGB images
    width, height = img_resolution
    observation_space = gyms.Box(low=0, high=255, shape=(height, width, 3), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
camera_name instance-attribute 
camera_name = camera_name
img_resolution instance-attribute 
img_resolution = img_resolution
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get camera image from environment rendering.

Source code in molmo_spaces/env/sensors_cameras.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get camera image from environment rendering."""

    # Use camera-specific frame access for multi-camera support
    # if hasattr(env, 'render_rgb_frame') and callable(env.render_rgb_frame):
    frame = env.render_rgb_frame(self.camera_name)

    if frame is not None:
        return frame

    # Return black image if no rendering available
    width, height = self.img_resolution
    return np.zeros((height, width, 3), dtype=np.uint8)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

DepthSensor 

DepthSensor(camera_name: str = 'camera', img_resolution: tuple[int, int] = (480, 480), uuid: str | None = None)

Bases: Sensor

Sensor for depth images from MuJoCo.

Returns raw metric depth in meters as float32. Encoding to RGB for video storage happens at save time. See molmo_spaces.utils.depth_utils for encoding/decoding functions.

Methods:

Name Description
get_observation

Get depth image from environment rendering.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
camera_name
img_resolution
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors_cameras.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def __init__(
    self,
    camera_name: str = "camera",
    img_resolution: tuple[int, int] = (480, 480),
    uuid: str | None = None,
) -> None:
    self.camera_name = camera_name
    self.img_resolution = img_resolution

    if uuid is None:
        uuid = f"depth_{camera_name}"

    # Define observation space for raw depth (float32 in meters)
    width, height = img_resolution
    observation_space = gyms.Box(low=0.0, high=10.0, shape=(height, width), dtype=np.float32)
    super().__init__(uuid=uuid, observation_space=observation_space)
camera_name instance-attribute 
camera_name = camera_name
img_resolution instance-attribute 
img_resolution = img_resolution
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get depth image from environment rendering.

Source code in molmo_spaces/env/sensors_cameras.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get depth image from environment rendering."""
    # Use camera-specific frame access for multi-camera support
    if hasattr(env, "render_depth_frame") and callable(env.render_depth_frame):
        frame = env.render_depth_frame(self.camera_name)
        if frame is not None:
            return frame

    # Fallback to default camera for backward compatibility
    if hasattr(env, "depth_frame") and env.depth_frame is not None:
        return env.depth_frame

    # Return zero depth if no rendering available
    width, height = self.img_resolution
    return np.zeros((height, width), dtype=np.float32)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None

SegmentationSensor 

SegmentationSensor(camera_name: str = 'camera', img_resolution: tuple[int, int] = (480, 480), uuid: str | None = None)

Bases: Sensor

Sensor for segmentation images from MuJoCo, outputs video-compatible arrays.

Methods:

Name Description
get_observation

Get segmentation image from environment rendering.
reset

Reset the sensor to its initial state.

Attributes:

Name Type Description
camera_name
img_resolution
is_dict bool
observation_space Space
str_max_len int
uuid str
Source code in molmo_spaces/env/sensors_cameras.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def __init__(
    self,
    camera_name: str = "camera",
    img_resolution: tuple[int, int] = (480, 480),
    uuid: str | None = None,
) -> None:
    self.camera_name = camera_name
    self.img_resolution = img_resolution

    if uuid is None:
        uuid = f"segmentation_{camera_name}"

    # Define observation space for uint8 images with channel dimension
    width, height = img_resolution
    observation_space = gyms.Box(low=0, high=255, shape=(height, width, 1), dtype=np.uint8)
    super().__init__(uuid=uuid, observation_space=observation_space)
camera_name instance-attribute 
camera_name = camera_name
img_resolution instance-attribute 
img_resolution = img_resolution
is_dict class-attribute instance-attribute 
is_dict: bool = False
observation_space instance-attribute 
observation_space: Space = observation_space
str_max_len class-attribute instance-attribute 
str_max_len: int = 2000
uuid instance-attribute 
uuid: str = uuid
get_observation 
get_observation(env, task, batch_index: int = 0, *args, **kwargs) -> ndarray

Get segmentation image from environment rendering.

Source code in molmo_spaces/env/sensors_cameras.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def get_observation(self, env, task, batch_index: int = 0, *args, **kwargs) -> np.ndarray:
    """Get segmentation image from environment rendering."""
    # Use camera-specific frame access for multi-camera support
    if hasattr(env, "segmentation_frame") and callable(env.segmentation_frame):
        frame = env.segmentation_frame(self.camera_name)
        if frame is not None:
            return frame

    # Fallback to default camera for backward compatibility
    if hasattr(env, "segmentation_frame") and env.segmentation_frame is not None:
        return env.segmentation_frame

    # Return zero segmentation if no rendering available
    width, height = self.img_resolution
    return np.zeros((height, width, 1), dtype=np.uint8)
reset 
reset() -> None

Reset the sensor to its initial state.

Source code in molmo_spaces/env/abstract_sensors.py 
53
54
55
def reset(self) -> None:
    """Reset the sensor to its initial state."""
    return None