fix(envs): Add source dataset and index to metadata (#388)

* add source dataset and index to metadata

* fix typo

* fix coach class and its test

reasoning_gym/code/bf.py

@@ -9,6 +9,8 @@ from ..data.wordle_words import wordle_words
 from ..factory import ProceduralDataset, register_dataset
 from .contrib.bfit.Compiler import Compiler, Minify
 
+DATASET_NAME = "bf"
+
 
 @dataclass
 class BFConfig:
@@ -53,6 +55,8 @@ class BFDataset(ProceduralDataset):
             "question": rng.choice(self._prompt_templates).format(bf_program=bf_program),
             "answer": result,
             "metadata": {
+                "source_dataset": DATASET_NAME,
+                "source_index": idx,
                 "bfit_code": bfit_code,
                 "bf_program": bf_program,
                 "difficulty": {"difficulty": self.config.difficulty},
@@ -160,4 +164,4 @@ class BFCurriculum(BaseCurriculum):
 
 
 # Register the dataset
-register_dataset("bf", BFDataset, BFConfig, BFCurriculum)
+register_dataset(DATASET_NAME, BFDataset, BFConfig, BFCurriculum)

reasoning_gym/code/codeio.py

@@ -49,6 +49,8 @@ Tip: Here is a reference code snippet for this question. You can refer to this c
 {3}
 """
 
+DATASET_NAME = "codeio"
+
 
 @dataclass
 class CodeIOConfig:
@@ -117,7 +119,12 @@ class CodeIODataset(ProceduralDataset):
         return {
             "question": question,
             "answer": solution,
-            "metadata": {"input_data": input_data, "output_data": output_data},
+            "metadata": {
+                "source_dataset": DATASET_NAME,
+                "source_index": idx,
+                "input_data": input_data,
+                "output_data": output_data,
+            },
         }
 
     def _json_to_tree(self, data, label="root"):
@@ -231,4 +238,4 @@ class CodeIODataset(ProceduralDataset):
 
 
 # Register the dataset
-register_dataset("codeio", CodeIODataset, CodeIOConfig)
+register_dataset(DATASET_NAME, CodeIODataset, CodeIOConfig)