Use `dataclass` instead of `TypedDict` for config
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
Currently, nemo-rl uses `TypedDict` to host the config options:
https://github.com/NVIDIA-NeMo/RL/blob/b238e41a6bad66324635ea6a65ef8bd8aeb8165e/nemo_rl/algorithms/grpo.py#L115-L177
We can consider to use `dataclass` to host these options like:
```python
@dataclass
class GRPOConfig:
num_prompts_per_step: int
num_generations_per_prompt: int
max_num_epochs: int
max_num_steps: int
# ...
```
It has several benifits, like, more friendly for type checking for `pyrefly`/`mypy`, and ease of development, e.g., you can directly access `config.logger.log_dir` with type hints in IDE/editor, instead of typing `config["logger"]["log_dir"]`. We may also add some runtime validation logic in the `__post_init__` dunder for dataclass.
The dataclass could be converted from omegaconf using `to_object` method.
Contributor guide
Assessment
This issue has not been assessed yet.