[Dataset Format Standard Modification] Record Unflattened Spaces in Datasets #77

rodrigodelazcano · 2023-05-24T14:10:05Z

Description

This PR is the continuation of issue #57. The PR removes the automatic flattening of observation/action spaces in StepDataCallback to store Dict and Tuple spaces in their original format. This will allow to create custom dataset spaces different from the environment.

Next tasks

Find solution to record Tuple spaces (maybe name each hdf5 group/dataset as element_{id})
Encode custom action/observation spaces for datasets (the spaces can be different from the environment's)
Fix pytest
Documentation update

Checklist:

I have run the pre-commit checks with pre-commit run --all-files (see CONTRIBUTING.md instructions to set it up)
I have run pytest -v and no errors are present.
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
I solved any possible warnings that pytest -v has generated that are related to my code to the best of my knowledge.
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

elliottower · 2023-05-31T03:22:39Z

minari/data_collector/data_collector.py

typos in the comments and debugging print statements

elliottower · 2023-05-31T03:23:56Z

minari/data_collector/data_collector.py

@@ -250,8 +265,8 @@ def reset(
        assert STEP_DATA_KEYS.issubset(step_data.keys())

        # If last episode in global buffer has saved steps, we need to check if it was truncated or terminated
-        # If not, then we need to auto-truncate the episode
-        if len(self._buffer[-1]["actions"]) > 0:
+        # If not (empty dicitionary), then we need to auto-truncate the episode.


comments should ideally be more understandable than this, I'd rephrase to "if the dictionary is not empty, "

elliottower · 2023-05-31T03:24:41Z

minari/data_collector/data_collector.py

-        # If not, then we need to auto-truncate the episode
-        if len(self._buffer[-1]["actions"]) > 0:
+        # If not (empty dicitionary), then we need to auto-truncate the episode.
+        if self._buffer[-1]:


Feel like there's probably a cleaner way to test for this but I'd have to look into the code a bit more deeply to be sure

I'm not sure, but I'm open to suggestions here.

Looking back this seems okay to do, the comment above explains it clearly enough

minari/data_collector/data_collector.py

…ll as a test for tuple action spaces and a combo env with nested dict and tuple action spaces

…ened spaces

…rvations for create_dataset_from_buffers, this may be inefficient and need refactoring

… observation and action space of data now saved in dataset.

…tten-spaces

…pss+1 observations were being loaded when calling get_episodes

balisujohn · 2023-06-08T07:39:41Z

This is ready for review again @younik @rodrigodelazcano, though I still need to write the doc update for it.

elliottower · 2023-06-13T16:29:42Z

docs/content/dataset_standards.md

 | `env_spec`              | `str`      | json string of the Gymnasium environment spec.|
-| `dataset_name`          | `str`      | Name tag of the Minari dataset. |
+| `dataset_id`            | `str`      | Identifier of the Minari dataset. |


Definitely a better var name to use so this is good

I changed it to match our datasets, which seemed to use already dataset_id instead of dataset_name

elliottower · 2023-06-13T16:33:12Z

docs/content/dataset_standards.md

 | `code_permalink`        | `str`      | Link to a repository with the code used to generate the dataset.|
 | `author`                | `str`      | Author's name that created the dataset. |
 | `author_email`          | `str`      | Email of the author that created the dataset.|
 | `algorithm_name`        | `str`      | Name of the expert policy used to create the dataset. |
+| `action_space`          | `str`      | Serialized Gymnasium action space describing actions in dataset. |
+| `observation_space`     | `str`      | Serialized Gymnasium observation space describing observations in dataset. |


I didn’t look I’m too much detail but maybe it’s worth noting how they get serialized? Pickled I’m assuming if it’s like the rest of our stuff, but we may want to look into transforming them into safetensors as pickle objects can be security vulnerabilities (someone can maliciously upload stuff by modifying the source code and running it thru a local install no matter what kinds of precautions we take, if we use pickle)

We wrote custom json serialization/de-serialization strategy which doesn't provide a straightforward path to eval() attacks as far as I know. That's not to say there isn't any way someone could figure out how to do an eval attack. I will add a note to the documentation describing the serialization format.

elliottower · 2023-06-13T16:34:58Z

minari/data_collector/callbacks/step_data.py

    This callback can be overridden to add extra environment information in each step or
    edit the observation, action, reward, termination, truncation, or info returns.
    """

-    def __init__(self, env: gym.Env):


Haven’t looked in depth but shouldn’t there still be an init even if it doesn’t flatten things? I guess it’s just a callback though so maybe there’s no need

It seems to work ok without one, for example in the test here https://github.com/Farama-Foundation/Minari/blob/8ff86b9fab6229169eb7516f839c725223e4711a/tests/data_collector/callbacks/test_step_data_callback.py I think we don't need to set the self.env property anymore

elliottower · 2023-06-13T16:36:21Z

minari/data_collector/data_collector.py

        # actions, observations, rewards, terminations, truncations, and infos
        assert STEP_DATA_KEYS.issubset(step_data.keys())
+        # Check that the saved observation and action belong to the dataset's observation/action spaces
+        assert self.dataset_observation_space.contains(step_data["observations"])
+        assert self.dataset_action_space.contains(step_data["actions"])


In general we want to have error messages if the assert fails just for simpler debugging, so something like assert(x), f“x is not true: {x}”

I added some more descriptive messages to the asserts in data_collector.py

elliottower · 2023-06-13T16:41:04Z

minari/data_collector/data_collector.py

@@ -425,7 +462,16 @@ def save_to_disk(self, path: str, dataset_metadata: Optional[Dict] = None):
        for key, value in dataset_metadata.items():
            self._tmp_f.attrs[key] = value

-        self._buffer.append({key: [] for key in STEP_DATA_KEYS})
+        assert "observation_space" not in dataset_metadata.keys()
+        assert "action_space" not in dataset_metadata.keys()


Again it’s a bit of a pain but best to add the assertion error statements just so when it fails you can see the value that was incorrect (and make it something more informative than assertion error 2 !=3, I’ve seen that before and it makes it such a headache to even figure out what went wrong

elliottower · 2023-06-13T16:42:00Z

minari/dataset/minari_dataset.py

@@ -309,20 +331,20 @@ def update_dataset_from_buffer(self, buffer: List[dict]):
                                                                                        element compared to the number of action steps {len(eps_buff['actions'])} \
                                                                                        The initial and final observation must be included"
                seed = eps_buff.pop("seed", None)
-                eps_group = clear_episode_buffer(
+                episode_group = clear_episode_buffer(


Good change imo, eps could be misunderstood as epsilon and it’s not that long to say episode so might as well

elliottower · 2023-06-13T16:42:48Z

minari/dataset/minari_storage.py

-            self._action_space = env.action_space
-
-            env.close()
+            # ww will default to using the reconstructed observation and action spaces from the dataset


Typo in “ww” (I’m assuming at least)

elliottower · 2023-06-13T16:43:48Z

minari/dataset/minari_storage.py

+                )
+                self._action_space = deserialize_space(f.attrs["action_space"])
+            else:
+                # checking if the base library of the environment is present in the environment


Pain to standardize them all but ideally we probably want most of the comments to be uppercase at the beginning and look consistent, not a huge deal but your commenting above was great so maybe just making sure it all looks good

Fixed this instance

elliottower · 2023-06-13T16:45:20Z

tests/data_collector/callbacks/test_step_data_callback.py

+    episodes = data.get_episodes(episode_indices)
+    # verify we have the right number of episodes, available at the right indices
+    assert data.total_episodes == len(episodes)
+    # verify the actions and observations are in the appropriate action space and observation space, and that the episode lengths are correct


Lowercase comment (minor, can ignore if you want to)

fixed this instance

elliottower

Looks pretty much good to me besides a few very minor writing or code style things

balisujohn · 2023-06-13T21:15:45Z

@younik
#77 (comment)
In regards to this comment, I refactored those functions into get_episodes, which calls _get_episode_from_h5, which calls _h5_group_to_dict_recursive

balisujohn · 2023-06-13T21:20:15Z

I addressed the reviews on this PR. I will also refactor the tests to reduce the duplication of helper function/env/space definitions, then I think it's ready for final review.

balisujohn · 2023-06-13T22:02:51Z

Refactor for the test dependencies is done :^), last thing I wanna check is that the old hosted datasets can still be sampled from ok temporarily until we convert them. (Also I will add a note about serialization format to the documentation.)

younik

Two minor comments on tests; for me is good to merge, good job!

younik · 2023-06-13T22:08:48Z

tests/test_serialization.py

+@pytest.mark.parametrize(
+    "space",
+    [
+        gym.spaces.Box(low=-1, high=4, shape=(2,), dtype=np.float32),


Can you delegate to test_common?

Oh sounds good, I'll do it for test_serialization as well.

younik · 2023-06-13T22:17:18Z

tests/test_common.py

@@ -1,5 +1,238 @@
+from typing import Iterable


My idea for test_common was to factor repetitions in tests (like a list of spaces that we want to test or environments) similar to https://github.com/Farama-Foundation/Gymnasium/blob/main/tests/testing_env.py

So good to have here the environments that we can reuse in multiple tests; for the test themselves I will keep the structure of the package as before (with test_minari_storage.pyand test_minari_dataset.py)

To avoid confusion probably a better name of test_common.py is just common.py

yeah that makes sense to avoid the test file name format, I'll make that change.

balisujohn · 2023-06-13T22:23:10Z

One heads up is that I think this PR will break loading the Dict-spaced flattened datasets from the database, In particular pointmaze, since they are no longer in compliance with our dataset standard (unless we reintroduce optional flattening). I will test if this is the case later today when I have fast internet, but I wanted to indicate that I think this failure is likely. I added a dataset integrity test to test_dataset_download.py and it works okay for the Box-spaced environments it downloads.

…ndencies file name to common.py, removed depdency duplication in serialization.py, added a dataset integrity check to test_download_dataset_from_farama_server

balisujohn · 2023-06-13T23:42:06Z

OK, ready for final review. Assuming aren't blocked by the fact we will likely break outdated flattened dict pointmaze dataset loading, until the pointmaze datasets are replaced with up to date datasets.

younik · 2023-06-14T15:02:18Z

minari/serialization.py

+    elif isinstance(space, gym.spaces.Tuple):
+        result = {"type": "Tuple", "subspaces": []}
+        for subspace in space.spaces:
+            result["subspaces"].append(serialize_space(subspace, to_string=False))
+    if to_string:
+        return json.dumps(result)
+    else:
+        return result
+


Missing an else branch
When you have an unsupported space, this will throw an uninformative error at line 43 (result is not initialized); better to have an informative error in the else branch

Addressed, and also added a test to check that a TypeError is correctly raised if unsupported space types are serialized.

…corresponding test

* change episodes structure + refactor * fix pre-commit

rodrigodelazcano added 4 commits May 24, 2023 09:41

unflatten StepDataCallback

aa45e15

datacollector update

0f3b5c2

remove flatten metadata

5cd6421

add fix removed with rebase

91c3b99

rodrigodelazcano marked this pull request as draft May 24, 2023 14:11

balisujohn added 3 commits May 27, 2023 23:17

fixed tests, added draft for testing saving unflattened dict spaces

1fcd31a

added test and tentative support for dict valued action spaces

01929a0

Merge remote-tracking branch 'origin' into unflatten-spaces

d39bd01

balisujohn force-pushed the unflatten-spaces branch from 461eac7 to d39bd01 Compare May 31, 2023 02:02

fixed registartion path for test dict env

22dde61

elliottower reviewed May 31, 2023

View reviewed changes

minari/data_collector/data_collector.py

Copy link

Member

elliottower May 31, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

typos in the comments and debugging print statements

elliottower reviewed May 31, 2023

View reviewed changes

minari/data_collector/data_collector.py Show resolved Hide resolved

balisujohn added 7 commits June 2, 2023 02:52

added experimental support for unflattened tuple action spaces, as we…

da9f474

…ll as a test for tuple action spaces and a combo env with nested dict and tuple action spaces

dummy env registration entrypoint change, hopefully fixes online tests

1df1304

added tuple space reconstruction support and tests for it for unflatt…

949c5de

…ened spaces

added support experimental support for non-flattened actions and obse…

4c0733a

…rvations for create_dataset_from_buffers, this may be inefficient and need refactoring

added tests for unflattened spaces with nested discrete spaces

f553e73

added action and observation space serialization and deserialization,…

8ae05b0

… observation and action space of data now saved in dataset.

Merge branch 'main' of github.com:Farama-Foundation/Minari into unfla…

7843af1

…tten-spaces

balisujohn changed the title ~~[Standard modification] Record unflatten spaces in datasets~~ [Standard modification] Record Unflattened Spaces in Datasets Jun 6, 2023

balisujohn changed the title ~~[Standard modification] Record Unflattened Spaces in Datasets~~ [Dataset Format Standard Modification] Record Unflattened Spaces in Datasets Jun 6, 2023

balisujohn added 2 commits June 7, 2023 17:58

fixed bug where total_timesteps observations instead of total_timeste…

76554ad

…pss+1 observations were being loaded when calling get_episodes

changes to address review

1ada485

balisujohn requested a review from younik June 8, 2023 17:35

balisujohn mentioned this pull request Jun 9, 2023

Code review notes #79

Open

balisujohn added 2 commits June 9, 2023 03:48

updated doc

7cf0928

added more detailed description of new data format in doc

596c3a7

elliottower reviewed Jun 13, 2023

View reviewed changes

elliottower approved these changes Jun 13, 2023

View reviewed changes

elliottower marked this pull request as ready for review June 13, 2023 16:46

balisujohn added 2 commits June 13, 2023 17:04

changes to address review

6908ea1

removed Tuple annotations when saving a Minari dataset to HDF5

e3698c1

balisujohn added 2 commits June 13, 2023 17:21

Merge branch 'main' into unflatten-spaces

899156d

refactored tests to reduce helper function duplication

6d74e63

younik approved these changes Jun 13, 2023

View reviewed changes

added note about space serialization to doc, changed test shared depe…

b04eaa6

…ndencies file name to common.py, removed depdency duplication in serialization.py, added a dataset integrity check to test_download_dataset_from_farama_server

younik mentioned this pull request Jun 14, 2023

[WIP] Datacollector refactoring #91

Closed

younik reviewed Jun 14, 2023

View reviewed changes

balisujohn and others added 4 commits June 14, 2023 18:41

added TypeError when attempting to serialize unsupported Spaces, and …

77e8f13

…corresponding test

Merge branch 'main' into unflatten-spaces

b1ea332

change episodes structure + refactor (#2)

3cfed59

* change episodes structure + refactor * fix pre-commit

fix first observation

5230da5

younik merged commit 3579ad5 into Farama-Foundation:main Jun 16, 2023
12 checks passed

balisujohn mentioned this pull request Jun 17, 2023

added check to ensure observation and action at each index is in the … #92

Merged

7 tasks

[Dataset Format Standard Modification] Record Unflattened Spaces in Datasets #77

[Dataset Format Standard Modification] Record Unflattened Spaces in Datasets #77

Conversation

rodrigodelazcano commented May 24, 2023 • edited by balisujohn Loading

Description

Next tasks

Checklist:

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

balisujohn commented Jun 8, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

elliottower Jun 13, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

elliottower left a comment

Choose a reason for hiding this comment

balisujohn commented Jun 13, 2023

balisujohn commented Jun 13, 2023

balisujohn commented Jun 13, 2023 • edited Loading

younik left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

balisujohn commented Jun 13, 2023 • edited Loading

balisujohn commented Jun 13, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rodrigodelazcano commented May 24, 2023 •

edited by balisujohn

Loading

elliottower Jun 13, 2023 •

edited

Loading

balisujohn commented Jun 13, 2023 •

edited

Loading

balisujohn commented Jun 13, 2023 •

edited

Loading