apply reviews

Author: chenyushuo

docs/sphinx_doc/source/tutorial/trinity_configs.md

@@ -112,7 +112,7 @@ algorithm:
 - `optimizer`: Optimizer configuration for actor.
   - `lr`: Learning rate for actor.
   - `warmup_style`: Warmup style for actor's learning rate.
-- `sample_strategy`: The sampling strategy used for loading experiences from experience buffer.
+- `sample_strategy`: The sampling strategy used for loading experiences from experience buffer. Supported types: `default`, `staleness_control`, `mix`.
 - `advantage_fn`: The advantage function used for computing advantages.
 - `kl_penalty_fn`: The KL penalty function used for computing KL penalty applied in reward.
 - `kl_loss_fn`: The KL loss function used for computing KL loss.

docs/sphinx_doc/source_zh/tutorial/trinity_configs.md

@@ -112,7 +112,7 @@ algorithm:
 - `optimizer`: Actor 优化器的参数。
   - `lr`: 优化器的学习率。
   - `warmup_style`: 学习率的预热策略。
-- `sample_strategy`: 从 experience buffer 加载 experience 时使用的采样策略。
+- `sample_strategy`: 从 experience buffer 加载 experience 时使用的采样策略。支持类型：`default`、`staleness_control`、`mix`。
 - `advantage_fn`: 用于计算优势值的函数。
 - `kl_penalty_fn`: 用于在奖励中计算 KL 惩罚的函数。
 - `kl_loss_fn`: 用于计算 KL 损失的函数。

tests/buffer/sample_strategy_test.py

@@ -94,17 +94,15 @@ async def _verify_sampling_model_versions(self, exps_list, expected_model_versio
         if current_task:
             await current_task
 
-    async def _flexible_verify_model_version(self, step, staleness_limit):
+    async def _flexible_verify_model_version(self, step, max_staleness):
         _, metrics, _ = await self.sample_strategy.sample(step=step)
         self.assertGreaterEqual(
             metrics["sample/model_version/min"],
-            step - staleness_limit,
+            step - max_staleness,
             f"Min model version mismatch at step {step}",
         )
 
-    async def _flexible_verify_sampling_model_versions(
-        self, exps_list, check_steps, staleness_limit
-    ):
+    async def _flexible_verify_sampling_model_versions(self, exps_list, check_steps, max_staleness):
         self._init_buffer_writer_and_sample_strategy()
 
         # Write experiences to buffer, while sample and validate model versions
@@ -115,7 +113,7 @@ async def _flexible_verify_sampling_model_versions(
                 if current_task:
                     await current_task
                 current_task = asyncio.create_task(
-                    self._flexible_verify_model_version(step, staleness_limit)
+                    self._flexible_verify_model_version(step, max_staleness)
                 )
                 await asyncio.sleep(0.1)
 
@@ -146,9 +144,9 @@ async def test_default_queue_default_sample_strategy(self):
         await self._verify_sampling_model_versions(exps_list, expected_model_versions_map)
 
     async def test_default_queue_staleness_control_sample_strategy(self):
-        staleness_limit = 3
+        max_staleness = 3
         self.config.algorithm.sample_strategy = "staleness_control"
-        self.config.algorithm.sample_strategy_args = {"staleness_limit": staleness_limit}
+        self.config.algorithm.sample_strategy_args = {"max_staleness": max_staleness}
         self.config.buffer.trainer_input.experience_buffer = ExperienceBufferConfig(
             name="default_queue_staleness_control",
             storage_type=StorageType.QUEUE.value,
@@ -161,15 +159,15 @@ async def test_default_queue_staleness_control_sample_strategy(self):
         steps = self._default_steps()
         expected_model_versions_map = {}
         for step in steps:
-            predict_version = max(step - staleness_limit, 0)
+            predict_version = max(step - max_staleness, 0)
             expected_model_versions_map[step] = [
                 predict_version + i // self.exp_write_batch_size
                 for i in range(self.config.buffer.train_batch_size)
             ]
 
         await self._verify_sampling_model_versions(exps_list, expected_model_versions_map)
 
-    def _simulate_priority_queue(self, steps, staleness_limit=float("inf")):
+    def _simulate_priority_queue(self, steps, max_staleness=float("inf")):
         expected_model_versions_map = {}
         buffer = deque()
         exp_pool = deque()
@@ -187,7 +185,7 @@ def _simulate_priority_queue(self, steps, staleness_limit=float("inf")):
                         exp_pool.extend(buffer.pop())
                     while len(exp_pool) > 0 and len(batch_versions) < train_batch_size:
                         exp_version = exp_pool.popleft()
-                        if exp_version < step - staleness_limit:
+                        if exp_version < step - max_staleness:
                             continue
                         batch_versions.append(exp_version)
                     if len(batch_versions) >= train_batch_size:
@@ -214,9 +212,9 @@ async def test_priority_queue_default_sample_strategy(self):
         await self._verify_sampling_model_versions(exps_list, expected_model_versions_map)
 
     async def test_priority_queue_staleness_control_sample_strategy(self):
-        staleness_limit = 2
+        max_staleness = 2
         self.config.algorithm.sample_strategy = "staleness_control"
-        self.config.algorithm.sample_strategy_args = {"staleness_limit": staleness_limit}
+        self.config.algorithm.sample_strategy_args = {"max_staleness": max_staleness}
         self.config.buffer.trainer_input.experience_buffer = ExperienceBufferConfig(
             name="priority_queue_staleness_control",
             storage_type=StorageType.QUEUE.value,
@@ -227,14 +225,14 @@ async def test_priority_queue_staleness_control_sample_strategy(self):
         # init testing data
         exps_list = self._default_exp_list()
         steps = self._default_steps()
-        expected_model_versions_map = self._simulate_priority_queue(steps, staleness_limit)
+        expected_model_versions_map = self._simulate_priority_queue(steps, max_staleness)
 
         await self._verify_sampling_model_versions(exps_list, expected_model_versions_map)
 
     async def test_sql_staleness_control_sample_strategy(self):
-        staleness_limit = 2
+        max_staleness = 2
         self.config.algorithm.sample_strategy = "staleness_control"
-        self.config.algorithm.sample_strategy_args = {"staleness_limit": staleness_limit}
+        self.config.algorithm.sample_strategy_args = {"max_staleness": max_staleness}
         self.config.buffer.trainer_input.experience_buffer = ExperienceBufferConfig(
             name="sql_staleness_control",
             storage_type=StorageType.SQL.value,
@@ -245,7 +243,7 @@ async def test_sql_staleness_control_sample_strategy(self):
         exps_list = self._default_exp_list()
         steps = self._default_steps()
 
-        await self._flexible_verify_sampling_model_versions(exps_list, steps, staleness_limit)
+        await self._flexible_verify_sampling_model_versions(exps_list, steps, max_staleness)
 
     def tearDown(self):
         asyncio.run(self.buffer_writer.release())

trinity/algorithm/sample_strategy/sample_strategy.py

@@ -79,10 +79,10 @@ def load_state_dict(self, state_dict: dict) -> None:
 class StalenessControlSampleStrategy(DefaultSampleStrategy):
     def __init__(self, buffer_config: BufferConfig, **kwargs):
         super().__init__(buffer_config)
-        self.staleness_limit = kwargs.get("staleness_limit", float("inf"))
+        self.max_staleness = kwargs.get("max_staleness", float("inf"))
 
     async def sample(self, step: int, **kwargs) -> Tuple[Experiences, Dict, List]:
-        min_model_version = max(step - self.staleness_limit, 0)
+        min_model_version = max(step - self.max_staleness, 0)
         metrics = {}
         with Timer(metrics, "time/read_experience"):
             exp_list = await self.exp_buffer.read_async(min_model_version=min_model_version)

trinity/buffer/schema/sql_schema.py

@@ -39,7 +39,7 @@ class ExperienceModel(Base):  # type: ignore
     message_list = Column(JSON, nullable=True)
     reward = Column(Float, nullable=True)
     # for step info
-    model_version = Column(Integer, nullable=True)
+    model_version = Column(Integer, nullable=True, index=True)
     # serialized experience object
     experience_bytes = Column(LargeBinary, nullable=True)
     consumed = Column(Integer, default=0, index=True)

trinity/buffer/storage/sql.py

@@ -158,12 +158,17 @@ def _read_priority(self, batch_size: int, min_model_version: int = 0) -> List[Ex
             with retry_session(
                 self.session, self.max_retry_times, self.max_retry_interval
             ) as session:
-                query = session.query(self.table_model_cls).order_by(
-                    asc(self.table_model_cls.consumed), desc(self.table_model_cls.id)
-                )
+                query = session.query(self.table_model_cls)
                 if min_model_version > 0:
                     query = query.filter(self.table_model_cls.model_version >= min_model_version)
-                experiences = query.limit(batch_size).with_for_update().all()
+                experiences = (
+                    query.order_by(
+                        asc(self.table_model_cls.consumed), desc(self.table_model_cls.id)
+                    )
+                    .limit(batch_size)
+                    .with_for_update()
+                    .all()
+                )
                 if len(experiences) != batch_size:
                     if latest_size != len(experiences):
                         latest_size = len(experiences)