# Have one iteration of algo.train() shut down after one episode

**URL:** <https://discuss.ray.io/t/have-one-iteration-of-algo-train-shut-down-after-one-episode/21985>\
**Category:** RLlib\
**Created:** [March 7, 2025, 2:59pm UTC](https://discuss.ray.io/t/have-one-iteration-of-algo-train-shut-down-after-one-episode/21985 "2025-03-07T14:59:25Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![sunghyun.chung](https://avatars.discourse-cdn.com/v4/letter/s/41988e/32.png) [@sunghyun.chung](https://discuss.ray.io/u/sunghyun.chung)\
**Post date:** [March 7, 2025, 2:59pm UTC](https://discuss.ray.io/t/have-one-iteration-of-algo-train-shut-down-after-one-episode/21985/1 "2025-03-07T14:59:25Z")

</div>

**1. Severity of the issue: (select one)**  
 None: I’m just curious or want clarification.  
 Low: Annoying but doesn’t hinder my work.  
[O] Medium: Significantly affects my productivity but can find a workaround.  
 High: Completely blocks me.

**2. Environment:**

- Ray version: 2.43
- Python version: 3.12.9
- OS: Windows 11
- Cloud/Infrastructure:
- Other libs/tools (if relevant): torch 2.5, np 2.2

**3. What happened vs. what you expected:**

- Expected: algo.train() method to terminate when an episode is terminated.
- Actual: After the episode is terminated, resets and starts collect again until the total step size hits 1000.

This is my current code:

```auto
my_dqn_config = (
    DQNConfig()
    .environment(
        env='my_env,
        env_config=my_config,
    )
    .training(
        replay_buffer_config=replay_buffer_config,
        # train_batch_size_per_learner=32,
        # num_epochs=3,
        # shuffle_batch_per_epoch=True,
        model={
            "fcnet_hiddens": [256, 256]
            },
        num_steps_sampled_before_learning_starts=100
    )
    .learners(
        # TODO(@chungs4): This will fail without setting an environment
        # ray.init(runtime_env={"env_vars": {"USE_LIBUV": "0"}}) due to changes in Pytorch >= 2.4.0
        num_learners=1,
        # TODO(@chungs4): GPU Implementation with NCCL
        # num_gpus_per_learner=1
    )
    .env_runners(
        num_env_runners=1,
        batch_mode="complete_episodes"
    )
)

```

The training starts after collecting around 100 samples (due to num\_steps\_sampled\_before\_learning\_starts=100) and the batch is cut off by each episode (due to batch\_mode=“complete\_episodes”). However, I the training keeps running even though I want train() method to shut down after one episode.

---

<div class="post-metadata">

**Author:** ![christina](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/christina/32/7542_2.png) [@christina](https://discuss.ray.io/u/christina)\
**Post date:** [March 8, 2025, 2:40am UTC](https://discuss.ray.io/t/have-one-iteration-of-algo-train-shut-down-after-one-episode/21985/2 "2025-03-08T02:40:35Z")

</div>

Hi there and welcome to the Ray community!

Can you please post the `stop_config` that you are using with your training? There is an example of it here being used in RLllib : [Replay Buffers — Ray 2.43.0](https://docs.ray.io/en/latest/rllib/rllib-replay-buffers.html#basic-usage)

Thank you!!

---

<div class="post-metadata">

**Author:** ![mannyv](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mannyv/32/606_2.png) [@mannyv](https://discuss.ray.io/u/mannyv)\
**Post date:** [March 10, 2025, 2:12pm UTC](https://discuss.ray.io/t/have-one-iteration-of-algo-train-shut-down-after-one-episode/21985/3 "2025-03-10T14:12:28Z")

</div>

Hi @sunghyun.chung,

I think you are going to have a hard time getting it to stop after exactly one episode unless your episode has a fixed length.

Here is the logic for how rllib decides when to stop one training iteration.

> <https://github.com/ray-project/ray/blob/74f53dfc7bdccef48547c06b5ed8b9818a69340c/rllib/algorithms/algorithm.py#L4112>

You could create a custom DQN algorithm and overload the `should_stop` method.

The reason your are seeing 1000 timesteps is because of DQNs default `min_sample_timesteps_per_iteration`:

> <https://github.com/ray-project/ray/blob/74f53dfc7bdccef48547c06b5ed8b9818a69340c/rllib/algorithms/dqn/dqn.py#L156>

---

<div class="post-metadata">

**Author:** ![sunghyun.chung](https://avatars.discourse-cdn.com/v4/letter/s/41988e/32.png) [@sunghyun.chung](https://discuss.ray.io/u/sunghyun.chung)\
**Post date:** [March 12, 2025, 6:35pm UTC](https://discuss.ray.io/t/have-one-iteration-of-algo-train-shut-down-after-one-episode/21985/4 "2025-03-12T18:35:35Z")

</div>

Hey Christina. Thanks for the reply.

Rather than training, it was more of .reporting() method problem since it was directly relevant to how data is mounted on tensorboard. I got around by setting:

```auto
    .reporting(
        # I know that the episode will last at least 30+ no matter what.
        min_sample_timesteps_per_iteration=10
    )

```

1. Can I also set `min_sample_timesteps_per_iteration` using stop\_config?
2. I hard-coded the value `min_sample_timesteps_per_iteration=10` because the episode will run at least 30 timesteps so that the every train() would correspond to one episode. Is there a way (or a variable) to enforce 1-1 matching between one iteration and one episode without hardcoding as I did above?

---

<div class="post-metadata">

**Author:** ![sunghyun.chung](https://avatars.discourse-cdn.com/v4/letter/s/41988e/32.png) [@sunghyun.chung](https://discuss.ray.io/u/sunghyun.chung)\
**Post date:** [March 12, 2025, 6:37pm UTC](https://discuss.ray.io/t/have-one-iteration-of-algo-train-shut-down-after-one-episode/21985/5 "2025-03-12T18:37:47Z")

</div>

Thanks for the reply, mannyv.

It would have been a lot easier if I saw this comment earlier.  
I think overriding AlgorithmConfig or its subclass is a good idea. Thanks for the suggestion.

Meanwhile, what would happen if I set `min_sample_timesteps_per_iteration` as 0 or 1? Wouldn’t it clip iteration of train() call to match on episode? (since I am batching by an episode in env\_runner)