Add generation caching in TextEnvironment and fix bugs in TextEnvironment

docs/source/text_environments.md

@@ -114,6 +114,8 @@ Let's decompose the settings:
 | `max_tool_response`| The tool response is truncated to this number to avoid running out of model context.|
 | `max_length`       |  The maximum number of tokens to allow in an episode. |
 | `generation_kwargs`| Generation settings used by the language model. |
+| `use_cache` | Cache keys and values between segment generation. When using caching, [`TextEnvironment`] is not suited for training use, i.e. backpropagation through the generated graph. Use with trl Trainers is of course possible. Furthermore, caching requires, that there be no calculation dependencies between examples at inference time. When using `BatchNorm`, the model should thus be in eval model. Caching is not guaranteed to produce identical results compared to not using caching and you should test for yourself, if it is suited to your needs, model and `generation_kwargs`. Compatibility with Encoder-Decoder architectures is untested. `use_cache` may currently be incompatible with torch.compile due to a possible issue in the transformers library's generate method. See this comment in [_get_initial_cache_position](https://github.com/huggingface/transformers/blob/2e752ead46a8845e8a160d2043c1336447895690/src/transformers/generation/utils.py#L1582).|
+| `save_logits` | Whether to save logits for the generated tokens in the returned histories. Mainly intended to help the user test caching for their use case. Backpropagation through logits is not supported. |
 
 You can customize the environment to your needs and add custom tools and settings. Let's see how you can use the environment to have the model interact with the available tools!
 
@@ -169,6 +171,7 @@ The following table summarises the available attributes of the `TextEnvironment`
 | `token_masks`      | The token masks can be used to ignore system generated tokens by masking them. |
 | `completed`        | Indicates if the interaction with the environment has completed. |
 | `truncated`        | Indicates if the interaction with the environment has completed because max length was reached. |
+| `logits` | A list containing a tensor for each non-system segment. If no logits were saved, then `logits` is an empty list.  The tensors contain the logits for the tokens generated by the model. See `save_logits` in [`TextEnvironment`]. |
 
 With these attributes you can reconstruct every interaction of the model with the `TextEnvironment`. The `TextHistory` also lets you visualize the text history. Let's have a look!