[Benchmark] Fix RB benchmarks

[vmoens]

cc @harnvo

I changed the names so that we can have a fresh look at the benchmarks on the benchmark tool.
I will try to remove the old traces that are wrong (need to look at how I can modify the xml files on gh-pages)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (7 Unrelated Failures)

As of commit 9e2ee73 with merge base 6d217c6 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Build M1 Wheels / pytorch/rl / wheel-py3_8-cpu (gh)
  ImportError: cannot import name 'MemoryMappedTensor' from 'tensordict' (/Users/ec2-user/runner/_work/_temp/conda_environment_7299990076/lib/python3.8/site-packages/tensordict/__init__.py)
• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Workflow failed! Cannot read properties of undefined (reading 'ops')
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Workflow failed! Cannot read properties of undefined (reading 'ops')
• Unit-tests on Linux CPU / tests (3.11) / linux-job (gh)
  test/test_libs.py::TestGym::test_vecenvs_env[CartPole-v1]

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Habitat Tests on Linux / tests (3.9, 11.6) / linux-job (gh)
  test/test_libs.py::TestHabitat::test_habitat_render[False-HabitatPick-v0]
• RLHF Tests on Linux / unittests (3.9, 12.1) / linux-job (gh)
  test/test_rlhf.py::TestRollout::test_rollout_from_data
• Unit-tests on Linux CPU / tests (3.8) / linux-job (gh)
  test/test_libs.py::TestGym::test_vecenvs_env[CartPole-v1]

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

[harnvo]

Hi, just a suggestion. The partial fix proposed in my issue is just for easy illustration but might not be appropriate for a merge.

Somehow pytest do not recognize the class name anymore after partial(TensorDictPrioritizedReplayBuffer, alpha=1, beta=0.9). (In practice you would see pb6, pb7, ... or something as their class name. ) This adds confusion when benchmarking (especially when you would like to add more tests with partial!!).

This is because partical change the __name__ of a class, and I would like to suggest solving this issue by doing the following:

_TensorDictPrioritizedReplayBuffer = partial(TensorDictPrioritizedReplayBuffer, alpha=1, beta=0.9)
# preserve the name of the class even after partial
_TensorDictPrioritizedReplayBuffer.__name__ = TensorDictPrioritizedReplayBuffer.__name__

and then change the partial(TensorDictPrioritizedReplayBuffer, alpha=1, beta=0.9) within parameterize into _TensorDictPrioritizedReplayBuffer

[vmoens]

@harnvo edited following your suggestion, would mean a lot if you could submit your review!

[harnvo]

typo: functools.partial

[harnvo]

lgtm

[harnvo]

It seems there exists a infinite loop here and the test_rb_iterate could not end the 4th test. (Possibly due to the circular sampling in SamplerWithoutReplacement)

The old iterate works fine and I think you should change this back.

[vmoens]

Yep, I wanted to avoid the cost of entering __iter__ but it's still better than an infinite loop!

[harnvo]

In this case, the following would work:

def iterate(rb_iter):
    try:
        next(rb_iter)
    except StopIteration:
        pass

and inside test_rb_iterate:

benchmark(iterate, iter(rb))

[vmoens]

The problem is that I was call iter on iter, not solving the problem that the first iter was empty. I will push a fix in a sec

[harnvo]

Works in my PC

[vmoens]

@harnvo it should run now!