# A Simple `Environment` for Guessing the Number

```bash
python -m tinker_cookbook.recipes.multiplayer_rl.guess_number.train
```

The `test/env/all/reward/total` should increase from ~40% to >=50% in 20 steps.

### Background: Guess the Number

In this task, we train an LLM to guess a hidden integer number between 0 and 1024.
- If the LLM guess correctly, the task is done
- If the LLM guesses incorrectly, the LLM will be given the chance to guess again and information whether the guessed number is too low or high.
- The interaction automatically ends after 10 guesses.
The LLM is rewarded with 1 if it guesses correctly, otherwise 0.

Here is one example game if the correct guess is 640:
```
[User]: Guess a number between 0 and 1024.
[LLM]: Guess: 512
[User]: Too low.
[LLM]: Guess: 768
[User]: Too high.
[LLM]: Guess: 640
[User]: Correct!
```

### Defining a `Guess-the-Number` Environment in Reinforcement Learning (RL)

In RL [1] (or more accurately, POMDP [2]), we need to specify the following components to define an environment:

| Component | Description |
|-----------|-------------|
| 1. **Action Space** | The tokens generated by the LLM |
| 2. **Observation Space** | The tokens that the LLM sees |
| 3. **Initial Observation** | The initial tokens that the LLM sees |
| 4. **Transition Function** | How the LLM-generated tokens determine what the user would say (e.g. Correct, Too high/low) |
| 5. **Reward Function** | Whether the LLM has output the correct guess |

The action space and observation space are the same for most LLM applications.

### Implementing the `Environment` Object

To customize your own training environment, you need to write a file like `recipes.multiplayer_rl.guess_number.env`.
We have already implemented the abstract class `Env` for you, so you can focus on implementing the initial observation, transition function, and reward function.

We will start by explaining the `step` function, which is the "core" of the environment: it determines
- how the LLM-generated action will influence the environment,
- what is the reward of this action,
- whether the game will end with this action, and
- what is the LLM's next observation.

In our implementation, we first parse the integer tokens into messages, and access the LLM generation in Python string format with `message["content"]`.
Then we use a helper function `_get_user_turn` to compute what the user would say, and what the reward should be. This helper function is straightforward to implement, because it is plain Python without using special libraries.

(What's the difference between a tokenizer and a renderer? A renderer performs one more step: it first converts messages (a list of dictionaries) into a raw string, which will then be converted to a list of integers by a tokenizer.)

In order to obtain the next observation that the LLM will see, we first update the conversation history `self.turns` by concatenating the messages that the policy and the LLM have just generated.
Then ``self._obs`` will return the LLM's next observation, which is the "tokenized conversation history" provided by the renderer.
We end the implementation by returning `StepResult`, which contains the next observation, the reward, whether the game should stop, and `self.stop_condition` (which is usually the renderer's stop_sequences).

Based on the knowledge above, it is straightforward to implement `initial_observation`.
We have now finished implementing the customized training environment for guessing numbers.

### Constructing the Environment Object

The RL training config takes `GuessNumberDatasetBuilder` as an argument,

* which constructs `GuessNumberDataset`,
* which constructs `GuessNumberEnvGroupBuilder`,
* which constructs a group of `GuessNumberEnv`.

While there are a lot of classes involved, most of them are just boilerplate. Crucial information, such as the answer and the renderer, are all constructed in `GuessNumberDatasetBuilder.__call__`, which are eventually passed to construct `GuessNumberEnv`.

### Next

We have covered a basic example of guessing numbers. For more sophisticated examples that involve more than one language model, feel free to read the README.md files under `recipes.multiplayer_rl.twenty_questions` and `recipes.multiplayer_rl.text_arena`.

### References

[1] Sutton, R. S., & Barto, A. G. (2018). Reinforcement learning: An introduction (2nd ed.). MIT Press. http://incompleteideas.net/book/the-book-2nd.html
[2] Kaelbling, L. P., Littman, M. L., & Cassandra, A. R. (1998). Planning and acting in partially observable stochastic domains. Artificial Intelligence, 101(1–2), 99–134. https://doi.org/10.1016/S0004-3702(98)00023-X
