From f6abc93f3cd1d32edf410d5eaf587cb36739fea5 Mon Sep 17 00:00:00 2001 From: Tavish Date: Tue, 5 Aug 2025 19:09:39 +0800 Subject: [PATCH 1/7] fix rlds episode_metadata --- tensorflow_datasets/rlds/rlds_base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorflow_datasets/rlds/rlds_base.py b/tensorflow_datasets/rlds/rlds_base.py index 9e2f7ce1b53..6e34862392f 100644 --- a/tensorflow_datasets/rlds/rlds_base.py +++ b/tensorflow_datasets/rlds/rlds_base.py @@ -75,6 +75,8 @@ def build_info( episode_metadata = ds_config.episode_metadata_info if episode_metadata is None: episode_metadata = {} + else: + episode_metadata = {"episode_metadata": episode_metadata} step_info = { 'is_terminal': tf.bool, 'is_first': tf.bool, From 31a4c9c1c65750fe2867a1a1fe726846bd5346fe Mon Sep 17 00:00:00 2001 From: Tavish Date: Mon, 18 Aug 2025 09:41:15 +0800 Subject: [PATCH 2/7] fix robosuite example --- .../robosuite_panda_pick_place_can.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py b/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py index 88cfcf3a3f5..3d3d49e65e4 100644 --- a/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py +++ b/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py @@ -186,9 +186,11 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder): reward_info=np.float64, discount_info=np.float64, episode_metadata_info={ - 'agent_id': np.str_, - 'episode_index': np.int32, - 'episode_id': np.str_, + 'episode_metadata': { + 'agent_id': np.str_, + 'episode_index': np.int32, + 'episode_id': np.str_, + }, }, step_metadata_info={ 'tag:placed': np.bool_, From b99b86f88312158d23a56ed932b6eb3a45c835cd Mon Sep 17 00:00:00 2001 From: Tavish Date: Mon, 18 Aug 2025 17:17:39 +0800 Subject: [PATCH 3/7] fix robosuite example, again --- .../robosuite_panda_pick_place_can.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py b/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py index 3d3d49e65e4..b7322c75d83 100644 --- a/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py +++ b/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py @@ -117,9 +117,11 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder): reward_info=np.float64, discount_info=np.float64, episode_metadata_info={ - 'agent_id': np.str_, - 'episode_index': np.int32, - 'episode_id': np.str_, + 'episode_metadata': { + 'agent_id': np.str_, + 'episode_index': np.int32, + 'episode_id': np.str_, + }, }, step_metadata_info={ 'tag:placed': np.bool_, @@ -248,9 +250,11 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder): reward_info=np.float64, discount_info=np.float64, episode_metadata_info={ - 'agent_id': np.str_, - 'episode_index': np.int32, - 'episode_id': np.str_, + 'episode_metadata': { + 'agent_id': np.str_, + 'episode_index': np.int32, + 'episode_id': np.str_, + }, }, step_metadata_info={ 'tag:placed': np.bool_, From 642be13017ab0335338bdb677733e8d17cc23dc1 Mon Sep 17 00:00:00 2001 From: Tavish Date: Mon, 18 Aug 2025 17:18:46 +0800 Subject: [PATCH 4/7] fix generate example --- tensorflow_datasets/rlds/rlds_base.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tensorflow_datasets/rlds/rlds_base.py b/tensorflow_datasets/rlds/rlds_base.py index 6e34862392f..fa35ce70b4f 100644 --- a/tensorflow_datasets/rlds/rlds_base.py +++ b/tensorflow_datasets/rlds/rlds_base.py @@ -165,9 +165,19 @@ def _generate_examples_from_log_path( key_prefix = os.path.basename(log_path) with envlogger.Reader(log_path) as reader: for episode_dict in envlogger_reader.generate_episodes(reader): + assert "steps" in episode_dict, "steps must be in episode_dict" + if "episode_metadata" not in episode_dict: + episode_dict = { + "steps": episode_dict["steps"], + "episode_metadata": { + key: value + for key, value in episode_dict.items() + if key != "steps" + } + } # The example ID should be unique. episode_id = counter - if 'episode_id' in episode_dict: - episode_id = episode_dict['episode_id'] + if 'episode_id' in episode_dict["episode_metadata"]: + episode_id = episode_dict["episode_metadata"]['episode_id'] yield f'{key_prefix}/{episode_id}', episode_dict counter += 1 From dd88ddabd10ceecddcefba5d46763102fa9b9920 Mon Sep 17 00:00:00 2001 From: Tavish Date: Tue, 19 Aug 2025 09:42:45 +0800 Subject: [PATCH 5/7] fix empty metadata --- tensorflow_datasets/rlds/rlds_base.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/tensorflow_datasets/rlds/rlds_base.py b/tensorflow_datasets/rlds/rlds_base.py index fa35ce70b4f..41299d0ce9b 100644 --- a/tensorflow_datasets/rlds/rlds_base.py +++ b/tensorflow_datasets/rlds/rlds_base.py @@ -167,13 +167,18 @@ def _generate_examples_from_log_path( for episode_dict in envlogger_reader.generate_episodes(reader): assert "steps" in episode_dict, "steps must be in episode_dict" if "episode_metadata" not in episode_dict: + episode_metadata = { + key: value + for key, value in episode_dict.items() + if key != "steps" + } episode_dict = { - "steps": episode_dict["steps"], - "episode_metadata": { - key: value - for key, value in episode_dict.items() - if key != "steps" - } + "steps": episode_dict["steps"], + **( + {"episode_metadata": episode_metadata} + if episode_metadata + else {} + ), } # The example ID should be unique. episode_id = counter From 9d69c3527d7f0c669c2d267073cc39ea6965e3e0 Mon Sep 17 00:00:00 2001 From: Tavish Date: Tue, 19 Aug 2025 09:49:41 +0800 Subject: [PATCH 6/7] fix robosuite episode metadata --- .../robosuite_panda_pick_place_can.py | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py b/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py index b7322c75d83..88cfcf3a3f5 100644 --- a/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py +++ b/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py @@ -117,11 +117,9 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder): reward_info=np.float64, discount_info=np.float64, episode_metadata_info={ - 'episode_metadata': { - 'agent_id': np.str_, - 'episode_index': np.int32, - 'episode_id': np.str_, - }, + 'agent_id': np.str_, + 'episode_index': np.int32, + 'episode_id': np.str_, }, step_metadata_info={ 'tag:placed': np.bool_, @@ -188,11 +186,9 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder): reward_info=np.float64, discount_info=np.float64, episode_metadata_info={ - 'episode_metadata': { - 'agent_id': np.str_, - 'episode_index': np.int32, - 'episode_id': np.str_, - }, + 'agent_id': np.str_, + 'episode_index': np.int32, + 'episode_id': np.str_, }, step_metadata_info={ 'tag:placed': np.bool_, @@ -250,11 +246,9 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder): reward_info=np.float64, discount_info=np.float64, episode_metadata_info={ - 'episode_metadata': { - 'agent_id': np.str_, - 'episode_index': np.int32, - 'episode_id': np.str_, - }, + 'agent_id': np.str_, + 'episode_index': np.int32, + 'episode_id': np.str_, }, step_metadata_info={ 'tag:placed': np.bool_, From 6601155d713def940c20bb65b22fecec16595753 Mon Sep 17 00:00:00 2001 From: Tavish Date: Tue, 19 Aug 2025 17:52:21 +0800 Subject: [PATCH 7/7] fix empty episode_metadata --- tensorflow_datasets/rlds/rlds_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow_datasets/rlds/rlds_base.py b/tensorflow_datasets/rlds/rlds_base.py index 41299d0ce9b..c880ef15717 100644 --- a/tensorflow_datasets/rlds/rlds_base.py +++ b/tensorflow_datasets/rlds/rlds_base.py @@ -182,7 +182,7 @@ def _generate_examples_from_log_path( } # The example ID should be unique. episode_id = counter - if 'episode_id' in episode_dict["episode_metadata"]: + if 'episode_id' in episode_dict.get("episode_metadata", {}): episode_id = episode_dict["episode_metadata"]['episode_id'] yield f'{key_prefix}/{episode_id}', episode_dict counter += 1