Skip to content

Recording

Module for recording data to an h5 file.

Recorder

Records a dataset to an h5 file. Saves images defined to folder with _frames appended to the name stem.

Example 
# Define the observation and action spaces
observation_space = spaces.Dict(
    {"image": spaces.Box(low=0, high=255, shape=(224, 224, 3), dtype=np.uint8), "instruction": spaces.Discrete(10)}
)
action_space = spaces.Dict(
    {
        "gripper_position": spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32),
        "gripper_action": spaces.Discrete(2),
    }
)

state_space = spaces.Dict(
    {
        "position": spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32),
        "velocity": spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32),
    }
)

# Create a recorder instance
recorder = Recorder(
    name="test_recorder", observation_space=observation_space, action_space=action_space, state_space=state_space
)

# Generate some sample data
num_steps = 10
for i in range(num_steps):
    observation = {"image": np.ones((224, 224, 3), dtype=np.uint8), "instruction": i}
    action = {"gripper_position": np.zeros((3,), dtype=np.float32), "gripper_action": 1}
    state = {"position": np.random.rand(3).astype(np.float32), "velocity": np.random.rand(3).astype(np.float32)}
    recorder.record(observation, action, state=state)

# Save the statistics
recorder.save_stats()

# Close the recorder
recorder.close()

# Assert that the HDF5 file and directories are created
assert os.path.exists("test_recorder.h5")
assert os.path.exists("test_recorder_frames")
Source code in mbodied/data/recording.py 
 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
class Recorder:
    """Records a dataset to an h5 file. Saves images defined to folder with _frames appended to the name stem.

    Example:
      ```
      # Define the observation and action spaces
      observation_space = spaces.Dict(
          {"image": spaces.Box(low=0, high=255, shape=(224, 224, 3), dtype=np.uint8), "instruction": spaces.Discrete(10)}
      )
      action_space = spaces.Dict(
          {
              "gripper_position": spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32),
              "gripper_action": spaces.Discrete(2),
          }
      )

      state_space = spaces.Dict(
          {
              "position": spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32),
              "velocity": spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32),
          }
      )

      # Create a recorder instance
      recorder = Recorder(
          name="test_recorder", observation_space=observation_space, action_space=action_space, state_space=state_space
      )

      # Generate some sample data
      num_steps = 10
      for i in range(num_steps):
          observation = {"image": np.ones((224, 224, 3), dtype=np.uint8), "instruction": i}
          action = {"gripper_position": np.zeros((3,), dtype=np.float32), "gripper_action": 1}
          state = {"position": np.random.rand(3).astype(np.float32), "velocity": np.random.rand(3).astype(np.float32)}
          recorder.record(observation, action, state=state)

      # Save the statistics
      recorder.save_stats()

      # Close the recorder
      recorder.close()

      # Assert that the HDF5 file and directories are created
      assert os.path.exists("test_recorder.h5")
      assert os.path.exists("test_recorder_frames")
      ```
    """

    def __init__(
        self,
        name: str = "dataset.h5",
        observation_space: spaces.Dict | str | None = None,
        action_space: spaces.Dict | str | None = None,
        state_space: spaces.Dict | str | None = None,
        supervision_space: spaces.Dict | str | None = None,
        out_dir: str = "saved_datasets",
        image_keys_to_save: list = None,
    ):
        """Initialize the Recorder.

        Args:
          name (str): Name of the file.
          observation_space (spaces.Dict): Observation space.
          action_space (spaces.Dict): Action space.
          state_space (spaces.Dict): State space.
          supervision_space (spaces.Dict): Supervision space.
          out_dir (str, optional): Directory of the output file. Defaults to 'saved_datasets'.
          image_keys_to_save (list, optional): List of image keys to save. Defaults to ['image'].
        """
        logging.info("\nInitializing dataset recorder, recording to directory: %s", out_dir)
        if image_keys_to_save is None:
            image_keys_to_save = ["image"]
        self.out_dir = out_dir
        self.frames_dir = Path(out_dir) / (Path(name).stem + "_frames")
        self.frames_dir.mkdir(exist_ok=True, parents=True)

        filename = Path(out_dir) / Path(name).with_suffix(".h5")
        Path(out_dir).mkdir(exist_ok=True, parents=True)
        if Path.exists(filename):
            copy_and_delete_old(filename)
        self.file = h5py.File(filename, "a")

        self.name = name
        self.filename = filename

        self.observation_space = observation_space
        self.action_space = action_space
        self.state_space = state_space
        self.supervision_space = supervision_space
        self.root_keys, self.root_spaces = self.configure_root_spaces(
            observation=observation_space,
            action=action_space,
            state=state_space,
            supervision=supervision_space,
        )

        self.image_keys_to_save = image_keys_to_save
        self.index = 0

    def reset(self) -> None:
        """Reset the recorder."""
        self.file.close()
        copy_and_delete_old(self.filename)
        self.file = h5py.File(self.filename, "a")

    def configure_root_spaces(self, **spaces: spaces.Dict):
        """Configure the root spaces.

        Args:
            **spaces: Spaces to configure.
                observation_space (spaces.Dict): Observation space.
                action_space (spaces.Dict): Action space.
                state_space (spaces.Dict): State space.
                supervision_space (spaces.Dict): Supervision space.
        """
        root_keys = []
        root_spaces = []
        for name, space in spaces.items():
            if space is None:
                continue

            root_keys.append(name)
            root_spaces.append(space)
            group = self.file.create_group(name)
            logging.debug("creating group %s", name)
            create_dataset_for_space_dict(space, group)
        return root_keys, root_spaces

    def record_timestep(self, group: h5py.Group, sample: Any, index: int) -> None:
        """Record a timestep.

        Args:
          group (h5py.Group): Group to record to.
          sample (Any): Sample to record.
          index (int): Index to record at.
        """
        if isinstance(group, h5py.Dataset):
            if index >= group.shape[0]:
                group.resize((2 * index, *group.shape[1:]))
            if hasattr(sample, "value"):
                sample = sample.value
            group[index] = sample
            return
        logging.debug("group keys: %s", str(group.keys()))
        if not hasattr(sample, "dict"):
            sample = Sample(sample)
        for key, value in sample:
            if value is None:
                continue
            if hasattr(value, "array"):
                dataset = group[key]
                if index >= dataset.shape[0]:
                    dataset.resize((2 * index, *dataset.shape[1:]))
                dataset[index] = value.array
                if key in self.image_keys_to_save and hasattr(value, "save"):
                    value.save(self.frames_dir / f"{self.index}.png")
                continue
            logging.debug(" key: %s, value: %s", key, value)

            if key not in group:
                logging.warning("key %s not in group %s. Skipping key", key, group)
                continue
            if isinstance(value, dict | Sample):
                subgroup = group[key]
                self.record_timestep(subgroup, value, index)
                continue

            if group[key].attrs.get("tuple_length") is not None:
                value = Sample.pack_from(value).model_dump_json(round_trip=True)  # noqa: PLW2901

            dataset = group[key]
            if index >= dataset.shape[0]:
                dataset.resize((2 * index, *dataset.shape[1:]))
            dataset[index] = value

    def record(
        self,
        observation: Any | None = None,
        action: Any | None = None,
        state: Any | None = None,
        supervision: Any | None = None,
    ) -> None:
        """Record a timestep.

        Args:
          observation (Any): Observation to record.
          action (Any): Action to record.
          state (Any): State to record.
          supervision (Any): Supervision to record.
        """

        def recursive_setarray(sample):
            if not hasattr(sample, "dict"):
                sample = Sample(sample)
            for key, value in sample:
                if isinstance(value, Image):
                    setattr(sample, key, value.array)
                elif isinstance(value, dict | Sample):
                    setattr(sample, key, recursive_setarray(value))
            return sample

        if observation is not None:
            if not hasattr(observation, "dict"):
                observation = Sample(observation)
                observation = recursive_setarray(observation)  # Bug hacky fix for Image recording.
            if "observation" not in self.file:
                logging.warning("Recorder: observation not in file, creating new group")
                new_root_keys, new_root_spaces = self.configure_root_spaces(observation=observation.space())
                self.root_keys += new_root_keys
                self.root_spaces += new_root_spaces
            self.record_timestep(self.file["observation"], observation, self.index)
        if action is not None:
            if not hasattr(action, "dict"):
                action = Sample(action)
                action = recursive_setarray(action)  # Bug hacky fix for Image recording.
            if "action" not in self.file:
                logging.warning("Recorder: action not in file, creating new group")
                new_root_keys, new_root_spaces = self.configure_root_spaces(action=action.space())
                self.root_keys += new_root_keys
                self.root_spaces += new_root_spaces
            self.record_timestep(self.file["action"], action, self.index)
        if state is not None:
            if not hasattr(state, "dict"):
                state = Sample(state)
                state = recursive_setarray(state)  # Bug hacky fix for Image recording.
            if "state" not in self.file:
                logging.warning("Recorder: state not in file, creating new group")
                new_root_keys, new_root_spaces = self.configure_root_spaces(state=state.space())
                self.root_keys += new_root_keys
                self.root_spaces += new_root_spaces
            self.record_timestep(self.file["state"], state, self.index)
        if supervision is not None:
            if not hasattr(supervision, "dict"):
                supervision = Sample(supervision)
                supervision = recursive_setarray(supervision)  # Bug hacky fix for Image recording.
            if "supervision" not in self.file:
                logging.warning("Recorder: supervision not in file, creating new group")
                new_root_keys, new_root_spaces = self.configure_root_spaces(supervision=supervision.space())
                self.root_keys += new_root_keys
                self.root_spaces += new_root_spaces
            self.record_timestep(self.file["supervision"], supervision, self.index)

        self.index += 1
        self.file.attrs["size"] = self.index

    def close(self) -> None:
        """Closes the Recorder and send the data if train_config is set."""
        self.file.close()

__init__(name='dataset.h5', observation_space=None, action_space=None, state_space=None, supervision_space=None, out_dir='saved_datasets', image_keys_to_save=None)

Initialize the Recorder.

Parameters:

Name Type Description Default
name str

Name of the file.

 'dataset.h5'
observation_space Dict

Observation space.

 None
action_space Dict

Action space.

 None
state_space Dict

State space.

 None
supervision_space Dict

Supervision space.

 None
out_dir str

Directory of the output file. Defaults to 'saved_datasets'.

 'saved_datasets'
image_keys_to_save list

List of image keys to save. Defaults to ['image'].

 None
Source code in mbodied/data/recording.py 
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
def __init__(
    self,
    name: str = "dataset.h5",
    observation_space: spaces.Dict | str | None = None,
    action_space: spaces.Dict | str | None = None,
    state_space: spaces.Dict | str | None = None,
    supervision_space: spaces.Dict | str | None = None,
    out_dir: str = "saved_datasets",
    image_keys_to_save: list = None,
):
    """Initialize the Recorder.

    Args:
      name (str): Name of the file.
      observation_space (spaces.Dict): Observation space.
      action_space (spaces.Dict): Action space.
      state_space (spaces.Dict): State space.
      supervision_space (spaces.Dict): Supervision space.
      out_dir (str, optional): Directory of the output file. Defaults to 'saved_datasets'.
      image_keys_to_save (list, optional): List of image keys to save. Defaults to ['image'].
    """
    logging.info("\nInitializing dataset recorder, recording to directory: %s", out_dir)
    if image_keys_to_save is None:
        image_keys_to_save = ["image"]
    self.out_dir = out_dir
    self.frames_dir = Path(out_dir) / (Path(name).stem + "_frames")
    self.frames_dir.mkdir(exist_ok=True, parents=True)

    filename = Path(out_dir) / Path(name).with_suffix(".h5")
    Path(out_dir).mkdir(exist_ok=True, parents=True)
    if Path.exists(filename):
        copy_and_delete_old(filename)
    self.file = h5py.File(filename, "a")

    self.name = name
    self.filename = filename

    self.observation_space = observation_space
    self.action_space = action_space
    self.state_space = state_space
    self.supervision_space = supervision_space
    self.root_keys, self.root_spaces = self.configure_root_spaces(
        observation=observation_space,
        action=action_space,
        state=state_space,
        supervision=supervision_space,
    )

    self.image_keys_to_save = image_keys_to_save
    self.index = 0

close()

Closes the Recorder and send the data if train_config is set.

Source code in mbodied/data/recording.py 
335
336
337
def close(self) -> None:
    """Closes the Recorder and send the data if train_config is set."""
    self.file.close()

configure_root_spaces(**spaces)

Configure the root spaces.

Parameters:

Name Type Description Default
**spaces Dict

Spaces to configure. observation_space (spaces.Dict): Observation space. action_space (spaces.Dict): Action space. state_space (spaces.Dict): State space. supervision_space (spaces.Dict): Supervision space.

 {}
Source code in mbodied/data/recording.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
def configure_root_spaces(self, **spaces: spaces.Dict):
    """Configure the root spaces.

    Args:
        **spaces: Spaces to configure.
            observation_space (spaces.Dict): Observation space.
            action_space (spaces.Dict): Action space.
            state_space (spaces.Dict): State space.
            supervision_space (spaces.Dict): Supervision space.
    """
    root_keys = []
    root_spaces = []
    for name, space in spaces.items():
        if space is None:
            continue

        root_keys.append(name)
        root_spaces.append(space)
        group = self.file.create_group(name)
        logging.debug("creating group %s", name)
        create_dataset_for_space_dict(space, group)
    return root_keys, root_spaces

record(observation=None, action=None, state=None, supervision=None)

Record a timestep.

Parameters:

Name Type Description Default
observation Any

Observation to record.

 None
action Any

Action to record.

 None
state Any

State to record.

 None
supervision Any

Supervision to record.

 None
Source code in mbodied/data/recording.py 
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
def record(
    self,
    observation: Any | None = None,
    action: Any | None = None,
    state: Any | None = None,
    supervision: Any | None = None,
) -> None:
    """Record a timestep.

    Args:
      observation (Any): Observation to record.
      action (Any): Action to record.
      state (Any): State to record.
      supervision (Any): Supervision to record.
    """

    def recursive_setarray(sample):
        if not hasattr(sample, "dict"):
            sample = Sample(sample)
        for key, value in sample:
            if isinstance(value, Image):
                setattr(sample, key, value.array)
            elif isinstance(value, dict | Sample):
                setattr(sample, key, recursive_setarray(value))
        return sample

    if observation is not None:
        if not hasattr(observation, "dict"):
            observation = Sample(observation)
            observation = recursive_setarray(observation)  # Bug hacky fix for Image recording.
        if "observation" not in self.file:
            logging.warning("Recorder: observation not in file, creating new group")
            new_root_keys, new_root_spaces = self.configure_root_spaces(observation=observation.space())
            self.root_keys += new_root_keys
            self.root_spaces += new_root_spaces
        self.record_timestep(self.file["observation"], observation, self.index)
    if action is not None:
        if not hasattr(action, "dict"):
            action = Sample(action)
            action = recursive_setarray(action)  # Bug hacky fix for Image recording.
        if "action" not in self.file:
            logging.warning("Recorder: action not in file, creating new group")
            new_root_keys, new_root_spaces = self.configure_root_spaces(action=action.space())
            self.root_keys += new_root_keys
            self.root_spaces += new_root_spaces
        self.record_timestep(self.file["action"], action, self.index)
    if state is not None:
        if not hasattr(state, "dict"):
            state = Sample(state)
            state = recursive_setarray(state)  # Bug hacky fix for Image recording.
        if "state" not in self.file:
            logging.warning("Recorder: state not in file, creating new group")
            new_root_keys, new_root_spaces = self.configure_root_spaces(state=state.space())
            self.root_keys += new_root_keys
            self.root_spaces += new_root_spaces
        self.record_timestep(self.file["state"], state, self.index)
    if supervision is not None:
        if not hasattr(supervision, "dict"):
            supervision = Sample(supervision)
            supervision = recursive_setarray(supervision)  # Bug hacky fix for Image recording.
        if "supervision" not in self.file:
            logging.warning("Recorder: supervision not in file, creating new group")
            new_root_keys, new_root_spaces = self.configure_root_spaces(supervision=supervision.space())
            self.root_keys += new_root_keys
            self.root_spaces += new_root_spaces
        self.record_timestep(self.file["supervision"], supervision, self.index)

    self.index += 1
    self.file.attrs["size"] = self.index

record_timestep(group, sample, index)

Record a timestep.

Parameters:

Name Type Description Default
group Group

Group to record to.

 required
sample Any

Sample to record.

 required
index int

Index to record at.

 required
Source code in mbodied/data/recording.py 
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
def record_timestep(self, group: h5py.Group, sample: Any, index: int) -> None:
    """Record a timestep.

    Args:
      group (h5py.Group): Group to record to.
      sample (Any): Sample to record.
      index (int): Index to record at.
    """
    if isinstance(group, h5py.Dataset):
        if index >= group.shape[0]:
            group.resize((2 * index, *group.shape[1:]))
        if hasattr(sample, "value"):
            sample = sample.value
        group[index] = sample
        return
    logging.debug("group keys: %s", str(group.keys()))
    if not hasattr(sample, "dict"):
        sample = Sample(sample)
    for key, value in sample:
        if value is None:
            continue
        if hasattr(value, "array"):
            dataset = group[key]
            if index >= dataset.shape[0]:
                dataset.resize((2 * index, *dataset.shape[1:]))
            dataset[index] = value.array
            if key in self.image_keys_to_save and hasattr(value, "save"):
                value.save(self.frames_dir / f"{self.index}.png")
            continue
        logging.debug(" key: %s, value: %s", key, value)

        if key not in group:
            logging.warning("key %s not in group %s. Skipping key", key, group)
            continue
        if isinstance(value, dict | Sample):
            subgroup = group[key]
            self.record_timestep(subgroup, value, index)
            continue

        if group[key].attrs.get("tuple_length") is not None:
            value = Sample.pack_from(value).model_dump_json(round_trip=True)  # noqa: PLW2901

        dataset = group[key]
        if index >= dataset.shape[0]:
            dataset.resize((2 * index, *dataset.shape[1:]))
        dataset[index] = value

reset()

Reset the recorder.

Source code in mbodied/data/recording.py 
189
190
191
192
193
def reset(self) -> None:
    """Reset the recorder."""
    self.file.close()
    copy_and_delete_old(self.filename)
    self.file = h5py.File(self.filename, "a")