Fix CPU training slowdown by introducing train_freq

[google-labs-jules]

This PR addresses the issue where CPU-only, single-agent training throughput drops to ~1 step/sec after the replay buffer warms up.

Changes made:

• Added train_freq parameter to OptimizationConfig in config.py with a default of 1 for GPU use cases.
• Modified trainer.py to only invoke agent.optimize_model() if total_steps % cfg.opt.train_freq == 0.
• Documented the recommended settings (train_freq = 4 and target_update_freq = 2500) for CPU-only training in README.md.

---

PR created automatically by Jules for task 509645655834908047 started by @dzaczek

[google-labs-jules]

👋 Jules, reporting for duty! I'm here to lend a hand with this pull request.

When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down.

I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job!

For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with @jules. You can find this option in the Pull Request section of your global Jules UI settings. You can always switch back!

New to Jules? Learn more at jules.google/docs.

---

For security, I will only act on instructions from the user who triggered this task.

[dzaczek]

Hi @jules

The train_freq approach makes sense and is a reasonable way to prevent CPU optimization from blocking the environment loop. However, I think a few changes are needed before merging.

1. Do not reduce target_update_freq to 2500

Currently, target network updates are controlled by:

if total_steps % cfg.opt.target_update_freq == 0:
    agent.update_target()

total_steps counts environment steps, not optimizer updates.

With train_freq = 4 and target_update_freq = 2500, the target network would be updated every 2,500 environment steps, which corresponds to only 625 optimizer updates.

Compared with the previous configuration (train_freq = 1, target_update_freq = 10000), this makes target synchronization 16 times more frequent relative to optimizer updates and may negatively affect training stability.

I suggest keeping:

train_freq = 4
target_update_freq = 10000

This preserves the existing target-update interval in terms of environment steps.

If the goal is to preserve 10,000 optimizer updates between target synchronizations, target_update_freq would need to be 40000 when train_freq = 4.

An even clearer long-term solution would be to introduce a separate optimizer_steps counter and define target_update_freq explicitly in optimizer steps.

2. Validate train_freq

A value of 0 currently causes a modulo-by-zero exception, while negative values are invalid configuration.

Please add validation ensuring:

train_freq >= 1

For example, this could be handled in OptimizationConfig.__post_init__() or when starting the trainer.

3. Correct the README wording

The current text says that the change 'batches the experience.' The batch size remains unchanged at 128; the implementation skips optimizer execution on some environment steps.

A clearer description would be:

Setting train_freq = 4 runs one optimization step for every four environment steps, reducing CPU load while continuing to add every collected transition to the replay buffer.

4. Consider exposing this as a CLI option

Requiring users to modify config.py is inconvenient. It would be useful to support something like:

python trainer.py --train-freq 4

This is optional for this PR, but it would make CPU/GPU configuration much easier.

5. Add basic tests

Please add tests confirming that:

• train_freq = 1 calls optimize_model() on every environment step;
• train_freq = 4 calls it once every four steps;
• invalid values such as 0 are rejected;
• target network updates still follow their documented counter and frequency.

Overall, the optimization-throttling mechanism is good, but the recommended target_update_freq = 2500 should be corrected before merging.

What do you think about tgat

👋 Jules, reporting for duty! I'm here to lend a hand with this pull request.

When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down.

I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job!

For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with @jules. You can find this option in the Pull Request section of your global Jules UI settings. You can always switch back!

New to Jules? Learn more at jules.google/docs.

For security, I will only act on instructions from the user who triggered this task.

[google-labs-jules]

Hi @jules

The train_freq approach makes sense and is a reasonable way to prevent CPU optimization from blocking the environment loop. However, I think a few changes are needed before merging.

1. Do not reduce target_update_freq to 2500

Currently, target network updates are controlled by:

if total_steps % cfg.opt.target_update_freq == 0:
    agent.update_target()

total_steps counts environment steps, not optimizer updates.

With train_freq = 4 and target_update_freq = 2500, the target network would be updated every 2,500 environment steps, which corresponds to only 625 optimizer updates.

Compared with the previous configuration (train_freq = 1, target_update_freq = 10000), this makes target synchronization 16 times more frequent relative to optimizer updates and may negatively affect training stability.

I suggest keeping:

train_freq = 4
target_update_freq = 10000

This preserves the existing target-update interval in terms of environment steps.

If the goal is to preserve 10,000 optimizer updates between target synchronizations, target_update_freq would need to be 40000 when train_freq = 4.

An even clearer long-term solution would be to introduce a separate optimizer_steps counter and define target_update_freq explicitly in optimizer steps.

2. Validate train_freq

A value of 0 currently causes a modulo-by-zero exception, while negative values are invalid configuration.

Please add validation ensuring:

train_freq >= 1

For example, this could be handled in OptimizationConfig.__post_init__() or when starting the trainer.

3. Correct the README wording

The current text says that the change 'batches the experience.' The batch size remains unchanged at 128; the implementation skips optimizer execution on some environment steps.

A clearer description would be:

Setting train_freq = 4 runs one optimization step for every four environment steps, reducing CPU load while continuing to add every collected transition to the replay buffer.

4. Consider exposing this as a CLI option

Requiring users to modify config.py is inconvenient. It would be useful to support something like:

python trainer.py --train-freq 4

This is optional for this PR, but it would make CPU/GPU configuration much easier.

5. Add basic tests

Please add tests confirming that:

• train_freq = 1 calls optimize_model() on every environment step;
• train_freq = 4 calls it once every four steps;
• invalid values such as 0 are rejected;
• target network updates still follow their documented counter and frequency.

Overall, the optimization-throttling mechanism is good, but the recommended target_update_freq = 2500 should be corrected before merging.

What do you think about tgat

👋 Jules, reporting for duty! I'm here to lend a hand with this pull request.

When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down.

I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job!

For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with @jules. You can find this option in the Pull Request section of your global Jules UI settings. You can always switch back!

New to Jules? Learn more at jules.google/docs.

For security, I will only act on instructions from the user who triggered this task.

Thank you for the detailed feedback! I have implemented the following changes to address your points:

1. Target Update Freq: I've removed the suggestion to change target_update_freq to 2500 from the README, preserving the 10,000 environment steps interval.
2. Validate train_freq: Added an __post_init__ check in OptimizationConfig to ensure train_freq >= 1, rejecting invalid values like 0 or negative numbers.
3. Correct README wording: Updated the wording in the README to clarify that train_freq = 4 simply runs one optimization step for every four environment steps.
4. CLI Option: Added a --train-freq argument in trainer.py so users can easily set this without modifying config.py.
5. Basic tests: Added tests/test_train_freq.py which includes validation for the train_freq values and a simulation test to confirm the modulo behavior for the optimization calls and target updates.