[kaanutkuy]

Description

Fixes #4704.

One commit per problem described in the issue:

1. Stop repeating once the episode is done. Before each extra repeat, _step checks the done entries of the last step with _terminated_or_truncated (the helper EnvBase already uses, which handles nested done keys) and stops if any is set. This matches the loop in GymLikeEnv._step, which breaks on done.any().
2. Advance the state between repeats. Each repeat now steps from the state reached by the previous one, parent.step_mdp(tensordict.copy().set("next", next_tensordict)). Envs that keep their state internally behave as before; envs that read it from the tensordict (PendulumEnv, and from reading the code BraxWrapper and JumanjiWrapper) now take frame_skip real steps. The copy leaves the caller's tensordict unchanged. Entries the caller preset in "next", such as the next recurrent state written by an LSTM/GRU policy, are carried over to every repeat: step_mdp drops "next", and a TensorDictPrimer after the frame skip would otherwise reset them.

The class docstring now says that each repeat starts from the state reached so far and that the repetition stops early on done (for any element of a batched env).

Tests

Added to TestFrameSkipTransform. The expected values come from the plain env or from GymEnv's built-in frame skip, never from the transform.

• test_frame_skip_stops_at_done: the base env is done after 2 of 4 repeats; the transform must stop there with every done key set (flat and nested keys).
• test_transform_env_episode_end (gym): GymEnv(Pendulum, frame_skip=3) against the transform for 70 steps across the truncation at base step 200, which lands on the 2nd repeat of the 67th step.
• test_frame_skip_stateless_env (skip 2 and 3): PendulumEnv against the plain env rolled out skip steps, and the input tensordict must be unchanged.
• test_frame_skip_stateless_done_on_terminal_step: a small stateless env whose done is set only on the terminal step (flat and nested keys).
• test_frame_skip_keeps_next_preset: a preset ("next", "hidden") must survive the repeats with a TensorDictPrimer after the frame skip.

The tests of each commit fail on the code before it. The exception is the preset test, which guards the second commit's own change: it fails if the state advance is done without carrying "next" over.

Behavior notes

These are also raised in #4704.

• In a batched parent, the repeats stop for all elements as soon as any element is done. This is the rule GymLikeEnv uses (done.any()).
• A transform that keeps its state in the tensordict and sits before FrameSkipTransform, e.g. Compose(StepCounter(), FrameSkipTransform(3)), now advances once per base step, so step_count goes 3, 6, 9 instead of 1, 2, 3. Nothing in the repo uses that order, and the auto-appended transform comes first.
• The extra step_mdp and copy cost about 30 to 40 us per extra repeat in a local CPU timing with no episode ends. On cheap envs (ContinuousActionVecMockEnv, PendulumEnv, gym Pendulum-v1) that is +19% to +27% per step with frame_skip=2 and +29% to +40% with frame_skip=4; envs with a more expensive step pay a smaller share. I did not add a benchmark since this is a correctness fix, and I can look at reducing this if it matters.

Testing

• test/transforms, test/envs, test/test_collectors.py and test/libs/test_gym.py pass with gymnasium 1.4.0.
• test/transforms passes in a local copy of the olddeps setup: Python 3.10, torch 2.11, gym 0.13, numpy 1.26, tensordict main.
• TestFrameSkipTransform passes with tensordict main and with 0.15.5.
• pre-commit passes.
• Not run: Brax, Jumanji and MuJoCo Playground (no jax here).

Not changed here

These are pre-existing and behave the same on main; I can open separate issues for them.

• The inverse of a transform placed before FrameSkipTransform is applied a second time on each repeat.
• Only a single reward key is supported, so envs with several reward keys raise a KeyError.

Motivation and Context

See #4704: on main the transform keeps stepping after the episode ends (summing rewards from past the end and possibly reporting done=False), and on stateless envs it re-runs the same transition frame_skip times.

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• Bug fix (non-breaking change which fixes an issue)

Checklist

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

I used Claude Code (Opus 5.5 xHigh) to help investigate, implement and test this change, and I reviewed all of it.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 16 Awaiting Approval

As of commit 467a74d with merge base 003c5f4 ():

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)
• Examples Tests on Linux (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.