[Feature] [TQC] [2/3] Add truncated quantile critics

[gtnv]

Summary

Add a reusable TQCLoss, a SAC-shaped HalfCheetah example, and an objective benchmark. The loss belongs in torchrl.objectives: pooled target truncation and quantile regression must follow TorchRL's functional-parameter, target-update, key-remapping, and value-estimation contracts.

Testing + notes

Environment: H100 80 GB, PyTorch 2.11.0+cu130.

• Numerical contracts: 7 passed; public doctest: 1 passed.
• Explicit critic lists pass with SoftUpdate, [2, 3] batches, and vectorized or looped evaluation.
• Eager, CUDA-graph, and compiled prioritized-memmap 50k-frame runs exit 0.
• Prioritized replay received varying finite [256] priorities for 3,000 CUDA-graph updates.
• Full objectives: 111 failed, 8021 passed, 2327 skipped. Every failure is an existing CUDA tolerance in test_values.py; the representative failure reproduces unchanged on upstream main.
• All pre-commit hooks pass.

Testing + Some Notes

Good news: The TQC implementation works great. Bad news: there was a different bug from an old issue which affected TQC performance that I will fix in a follow-up.

Three compiled 1M-frame seeds completed. Final-five evaluation means were 12,175 / 10,501 / 13,282. Seed 42 recovered from an entropy/critic explosion; seed 43 ended during one (-588 final return). This matches the open TanhNormal instability in #2199, so this PR does not add a TQC-specific workaround.

Performance

5 passed, 1 expected skip in 33.41 s.

loss path	median	IQR
eager forward	7.298 ms	0.137 ms
eager backward	9.584 ms	0.149 ms
compiled forward	1.262 ms	0.024 ms
compiled backward	2.528 ms	0.026 ms
reduce-overhead forward	8.674 ms	6.736 ms

cc @vmoens

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4078

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 15 Awaiting Approval

As of commit 81a5470 with merge base d7659c7 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux Wheels (gh)
• Build Linux Wheels (gh)
• Build M1 Wheels (gh)
• Build Windows Wheels (gh)
• Continuous Benchmark (PR) (gh)
• Generate documentation (gh)
• Habitat Tests on Linux (gh)
• Libs Tests on Linux (gh)
• Lint (gh)
• LLM Tests on Linux (gh)
• Push Binary Nightly (gh)
• SOTA Tests on Linux (gh)
• Tutorials Tests on Linux (gh)
• Unit-tests on Linux (gh)
• Validate Test Partitioning (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

⚠️ PR Title Label Error

Unknown or invalid prefix [Feat].

Current title: [Feat] [TQC] [2/3] Add truncated quantile critics

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
[Algorithm]	new algo	[Algorithm] Add new RL objective
[BE]	BE	[BE] Improve error messages
[Benchmark] or [Benchmarks]	Benchmarks	[Benchmark] Add collector benchmark
[BugFix]	BugFix	[BugFix] Fix memory leak in collector
[Example] or [Examples]	Examples	[Example] Add training script
[Feature]	Feature	[Feature] Add new optimizer
[Doc] or [Docs]	Documentation	[Doc] Update installation guide
[Refactor]	Refactoring	[Refactor] Clean up module imports
[CI]	CI	[CI] Fix workflow permissions
[Test] or [Tests]	Tests	[Tests] Add unit tests for buffer
[Trainer] or [Trainers]	Trainers	[Trainer] Add trainer config
[Environment] or [Environments]	Environments	[Environments] Add Gymnasium support
[Data]	Data	[Data] Fix replay buffer sampling
[LLM]	llm/	[LLM] Add reward model integration
[Minor]	small change	[Minor] Fix typo in error message
[Performance] or [Perf]	Performance	[Performance] Optimize tensor ops
[BC-Breaking]	bc breaking	[BC-Breaking] Remove deprecated API
[Deprecation]	Deprecation	[Deprecation] Mark old function
[Algorithm] or [Algorithms]	new algo	[Algorithm] Add new objective
[Quality]	Quality	[Quality] Fix typos and add codespell
[Versioning]	versioning	[Versioning] Bump release version
[WIP]	WIP	[WIP] Draft implementation

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

[github-actions]

⚠️ PR Title Label Error

Unknown or invalid prefix [Feat].

Current title: [Feat] [TQC] [2/3] Add truncated quantile critics

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
[Algorithm]	new algo	[Algorithm] Add new RL objective
[BE]	BE	[BE] Improve error messages
[Benchmark] or [Benchmarks]	Benchmarks	[Benchmark] Add collector benchmark
[BugFix]	BugFix	[BugFix] Fix memory leak in collector
[Example] or [Examples]	Examples	[Example] Add training script
[Feature]	Feature	[Feature] Add new optimizer
[Doc] or [Docs]	Documentation	[Doc] Update installation guide
[Refactor]	Refactoring	[Refactor] Clean up module imports
[CI]	CI	[CI] Fix workflow permissions
[Test] or [Tests]	Tests	[Tests] Add unit tests for buffer
[Trainer] or [Trainers]	Trainers	[Trainer] Add trainer config
[Environment] or [Environments]	Environments	[Environments] Add Gymnasium support
[Data]	Data	[Data] Fix replay buffer sampling
[LLM]	llm/	[LLM] Add reward model integration
[Minor]	small change	[Minor] Fix typo in error message
[Performance] or [Perf]	Performance	[Performance] Optimize tensor ops
[BC-Breaking]	bc breaking	[BC-Breaking] Remove deprecated API
[Deprecation]	Deprecation	[Deprecation] Mark old function
[Algorithm] or [Algorithms]	new algo	[Algorithm] Add new objective
[Quality]	Quality	[Quality] Fix typos and add codespell
[Versioning]	versioning	[Versioning] Bump release version
[WIP]	WIP	[WIP] Draft implementation

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

[vmoens]

Reviewed current head 5d318a12e2 after refreshing it against origin/main (it was already up to date), including the current CI discussion. The pooled target truncation, full-mixture actor value, explicit quantile-shape checks, and TD0 atom expansion look internally consistent, and the module doctest passes locally (1 passed).

I do see a merge blocker in the validation/integration surface: this introduces a new public objective and algorithm, but there is no test/test_tqc.py; the only added pytest coverage is the speed benchmark. Please add unit tests for the numerical truncation/quantile-regression contract and gradients, nested set_keys(), explicit critic lists plus deactivate_vmap, multidimensional batches, invalid drop counts, terminal skipping, and priority output. Per the repository rules, this also needs a TQCLossConfig with constructor parity and the new algorithm's sota-check / test-linux-sota smoke registration. The existing gym/sklearn CI failures shown by Dr. CI appear outside this diff, but these missing required checks currently leave the core behavior unguarded.

[gtnv]

@vmoens test is in this one: #4079 but other than that I added the TQCLossConfig. PTAL & TY.