apply suggestions from gemini

Author: chenyushuo

trinity/algorithm/sample_strategy/mix_sample_strategy.py

@@ -85,6 +85,7 @@ async def sample(self, step: int) -> Tuple[Experiences, Dict, List]:
             exp_list = usual_exp_list + expert_exp_list
             repr_samples = representative_sample(exp_list)
 
+        self.set_model_version_metric(exp_list, metrics)
         with Timer(metrics, "time/gather_experience"):
             exps = Experiences.gather_experiences(
                 experiences=exp_list,

trinity/algorithm/sample_strategy/sample_strategy.py

@@ -1,10 +1,10 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Tuple
+from typing import Dict, List, Tuple
 
 from trinity.algorithm.sample_strategy.utils import representative_sample
 from trinity.buffer import get_buffer_reader
 from trinity.common.config import BufferConfig
-from trinity.common.experience import Experiences
+from trinity.common.experience import Experience, Experiences
 from trinity.utils.annotations import Deprecated
 from trinity.utils.monitor import gather_metrics
 from trinity.utils.registry import Registry
@@ -17,6 +17,14 @@ class SampleStrategy(ABC):
     def __init__(self, buffer_config: BufferConfig, **kwargs) -> None:
         self.pad_token_id = buffer_config.pad_token_id
 
+    def set_model_version_metric(self, exp_list: List[Experience], metrics: Dict):
+        metric_list = [
+            {"model_version": exp.info["model_version"]}
+            for exp in exp_list
+            if "model_version" in exp.info
+        ]
+        metrics.update(gather_metrics(metric_list, "sample"))
+
     @abstractmethod
     async def sample(self, step: int) -> Tuple[Experiences, Dict, List]:
         """Sample data from buffer.
@@ -42,13 +50,12 @@ def __init__(self, buffer_config: BufferConfig, **kwargs):
         super().__init__(buffer_config)
         self.exp_buffer = get_buffer_reader(buffer_config.trainer_input.experience_buffer)  # type: ignore[arg-type]
 
-    async def sample(self, step: int, **kwargs) -> Tuple[Any, Dict, List]:
+    async def sample(self, step: int, **kwargs) -> Tuple[Experiences, Dict, List]:
         metrics = {}
         with Timer(metrics, "time/read_experience"):
             exp_list = await self.exp_buffer.read_async()
             repr_samples = representative_sample(exp_list)
-        metric_list = [{"model_version": exp.info["model_version"]} for exp in exp_list]
-        metrics.update(gather_metrics(metric_list, "sample"))
+        self.set_model_version_metric(exp_list, metrics)
         with Timer(metrics, "time/gather_experience"):
             exps = Experiences.gather_experiences(exp_list, self.pad_token_id)  # type: ignore
         return exps, metrics, repr_samples

trinity/trainer/trainer.py

@@ -125,7 +125,7 @@ async def _sample_data(self) -> Tuple[Experiences, Dict, List[Dict]]:
             List[Dict]: A list of representative samples for logging.
         """
         batch, metrics, repr_samples = await self.sample_strategy.sample(self.train_step_num + 1)
-        metrics["sample/task_count"] = len(set(f"{eid.batch}_{eid.task}" for eid in batch.eids))
+        metrics["sample/task_count"] = len(set(eid.tid for eid in batch.eids))
         return batch, metrics, repr_samples
 
     async def need_sync(self) -> bool: