[Feature] example_data for NonTensor spec

[vmoens]

Stack from ghstack (oldest at bottom):

• [Feature,Refactor] Chess improvements: fen, pgn, pixels, san #2702
• [Feature] UnaryTransform for input entries #2700
• [BugFix] patch rand_action in TransformedEnv to read the base_env method #2699
• -> [Feature] example_data for NonTensor spec #2698

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 7 Pending, 4 Unrelated Failures

As of commit 69c6122 with merge base 256a700 ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Unit-tests on Windows / unittests-cpu / windows-job (gh)
  Process completed with exit code 1.

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

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

• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cpu (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cpu_
• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cuda11_8 (gh) (trunk failure)
• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cuda12_4 (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cu124_
• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cuda12_6 (gh) (trunk failure)

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

[vmoens]

@kurtamohler I think it would also be useful to have a spec that specified a set of available choices.
I was thinking to do that in NonTensor first (like NonTensor(choices["a", "b", "c"]), where rand() would pick one of the three values) but that seem restrictive (we'd need another one if you can select some tensor values).

One solution I'm contemplating is to have a Choice spec that works like this:

choice = Choice(
    torch.stack([
        NonTensor("a"),
        NonTensor("b"),
        NonTensor("c"),
    ])
)
choice.rand() # gives NonTensorData("a"), for instance

Using this you could do

choice = Choice(
    torch.stack([
        Bounded(0, 1),
        Bounded(5, 10),
    ])
)
choice.rand() # gives a tensor with value in [0, 5, 6, 7, 8, 9]

[kurtamohler]

Choice sounds like a useful spec. I could implement that if you'd like

I like this example_data feature--it partially addresses something I saw when writing Hash/UnaryTransform, which caused me to make this stipulation about None inputs:

rl/torchrl/envs/transforms/transforms.py

Lines 4418 to 4419 in c5f1565

	fn (Callable): the function to use as the unary operation. If it accepts
	a non-tensor input, it must also accept ``None``.

The only reason that stipulation exists is that when UnaryTransform calls spec.zero() to make a fake input to determine the output spec, if the spec has a NonTensor, the result for it is None. But after this PR, the example_data gets chosen instead of None.

Just a thought--maybe example_data could have a sensible default other than None for some dtypes? Like for str, it could be "" or even "example_data"