{"id":15038707,"url":"https://github.com/allenai/rl4lms","last_synced_at":"2025-05-15T09:02:38.917Z","repository":{"id":64631511,"uuid":"526069335","full_name":"allenai/RL4LMs","owner":"allenai","description":"A modular RL library to fine-tune language models to human preferences","archived":false,"fork":false,"pushed_at":"2024-03-01T14:03:46.000Z","size":30548,"stargazers_count":2296,"open_issues_count":53,"forks_count":195,"subscribers_count":23,"default_branch":"main","last_synced_at":"2025-04-07T03:15:17.377Z","etag":null,"topics":["dialogue-generation","language-modeling","machine-translation","natural-language-processing","nlp","reinforcement-learning","summarization","table-to-text","text-generation"],"latest_commit_sha":null,"homepage":"https://rl4lms.apps.allenai.org/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/allenai.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-08-18T05:29:16.000Z","updated_at":"2025-04-07T01:15:15.000Z","dependencies_parsed_at":"2024-06-02T16:41:08.060Z","dependency_job_id":null,"html_url":"https://github.com/allenai/RL4LMs","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allenai%2FRL4LMs","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allenai%2FRL4LMs/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allenai%2FRL4LMs/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allenai%2FRL4LMs/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/allenai","download_url":"https://codeload.github.com/allenai/RL4LMs/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248890407,"owners_count":21178423,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["dialogue-generation","language-modeling","machine-translation","natural-language-processing","nlp","reinforcement-learning","summarization","table-to-text","text-generation"],"created_at":"2024-09-24T20:39:48.821Z","updated_at":"2025-04-14T13:44:50.215Z","avatar_url":"https://github.com/allenai.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cp align=\"center\"\u003e\n  \u003cimg src=\"RL4LMs_logo.png\" width=512px\u003e\n\u003c/p\u003e\n\n\u003ch1 align=\"center\"\u003e :robot: RL4LMs :rocket: \u003c/h1\u003e\n\n\u003ch3 align=\"center\"\u003e A modular RL library to fine-tune language models to human preferences \u003c/h3\u003e\n\n\u003cbr\u003e\n\nWe provide easily customizable building blocks for training language models including implementations of **on-policy algorithms**, **reward functions**, **metrics**, **datasets** and **LM based actor-critic policies**\n\nPaper Link: https://arxiv.org/abs/2210.01241\n\nWebsite Link: https://rl4lms.apps.allenai.org/\n\nThoroughly **tested** and **benchmarked** with over **2000 experiments** :fire: (GRUE benchmark :trophy:) on a comprehensive set of: \n- 7 different Natural Language Processing (NLP) Tasks:\n    - Summarization\n    - Generative Commonsense Reasoning\n    - IMDB Sentiment-based Text Continuation\n    - Table-to-text generation\n    - Abstractive Question Answering\n    - Machine Translation\n    - Dialogue Generation\n- Different types of NLG metrics (20+) which can be used as reward functions:\n    - Lexical Metrics (eg: ROUGE, BLEU, SacreBLEU, METEOR)\n    - Semantic Metrics (eg: BERTSCORE, BLEURT)\n    - Task specific metrics (eg: PARENT, CIDER, SPICE)\n    - Scores from pre-trained classifiers (eg: Sentiment scores)\n- On-policy algorithms of PPO, A2C, TRPO and novel **NLPO (Natural Language Policy Optimization)**\n- Actor-Critic Policies supporting causal LMs (eg. GPT-2/3) and seq2seq LMs (eg. T5, BART)\n\nAll of these building blocks can be customizable allowing users to train transformer-based LMs to optimize any arbitrary reward function on any dataset of their choice.\n\n## Recent updates (v0.2.0) on 23-Nov-22\n- Added daily dialog task\n- Fixed compatibility issues with some Seq2seq models such as BART, blendorbot etc\n- Implemented data parallel support\n- Refactored policy classes\n\n## Recent updates (v0.2.1) \n- Minor logging updates\n\n---\n# Install\n\n## Local Installation \n```bash\ngit clone https://github.com/allenai/RL4LMs.git\ncd RL4LMs\npip install -e .\n```\n\n## Docker\nWe provide also a Dockerfile for development using docker containers containing all the dependencies.\n```bash\ndocker build . -t rl4lms\n```\n\n## Additional dependencies\n\nOptionally, coreNLP libraries are required for certain metric computations (eg. SPICE) which can be downloaded through `cd rl4lms/envs/text_generation/caption_metrics/spice \u0026\u0026 bash get_stanford_models.sh`\n\n---\n# Quick Start - Train PPO/NLPO using pre-defined YAML configs\nWe provide a simple training API that can be invoked via train [script](https://github.com/allenai/RL4LMs/blob/main/scripts/training/train_text_generation.py) that allows to train PPO, NLPO or a supervised model by using a config file (YAML). \n\nFor example, to train T5-base on CNN/DM summarization on PPO using Rouge-1 as reward function, you can run:\n\n```bash\npython scripts/training/train_text_generation.py --config_path scripts/training/task_configs/summarization/t5_ppo.yml\n```\n\nConfig files for all tasks can be found [here](https://github.com/allenai/RL4LMs/tree/main/scripts/training/task_configs).\n\n## YAML file schema - Configuring building blocks\n\nConfig file contains details about hyper-parameter settings for building blocks which are described below:\n\n- **Dataset/Task**: Dataset containing samples with input prompts and reference sentences. Available datasets are found in the class `DataPoolRegistry` in [registry](https://github.com/allenai/RL4LMs/blob/main/rl4lms/envs/text_generation/registry.py). (See how to create your own dataset [here](#adding-dataset))\n\n  ```yaml\n  datapool:\n    id: cnn_daily_mail\n    args:\n      prompt_prefix: \"Summarize: \"\n  ```\n\n- **Tokenizer** - A pre-trained tokenizer that is used to (de)tokenize input and output sequences with settings for padding and truncation\n  ```yaml\n  tokenizer:\n    model_name: t5-base\n    padding_side: left\n    truncation_side: left\n    pad_token_as_eos_token: False\n  ``` \n- **Reward Function**: Reward function which computes token-level scores at each time step of MDP. Available reward functions can be found in the class `RewardFunctionRegistry`. (See how to create your own reward function [here](#adding-reward-function))\n\n  ```yaml\n  reward_fn:\n    id: rouge\n    args:\n      rouge_type: \"rouge1\"\n  ```\n\n- **Environment**: Configures a gym-style text generation [environment](https://github.com/allenai/RL4LMs/blob/main/rl4lms/envs/text_generation/env.py) which simulates MDP episodes. Rollouts are generated using train samples from dataset consisting of input and reference texts.\nFurther, we wrap our env with `SubProcVecEnv` from stable-baselines that processes `n_envs` episodes in parallel using multi-processing to compute step-wise rewards.  \nFurther configuration settings include: \n  - `max_episode_length` : max length of the episode \n  - `max_prompt_length` - maximum length of the input text to consider \n  - `terminate_on_eos` - whether to terminate the episode as soon as EOS action is performed \n  - `prompt_truncation_side` - truncation side for the prompt text \n  - `context_start_token` - id for context token (corresponds to initial token given to decoder in encoder-decoder models)\n\n  ```yaml\n  env:\n    n_envs: 10\n    args:\n      max_prompt_length: 512\n      max_episode_length: 100\n      terminate_on_eos: True\n      prompt_truncation_side: \"right\"\n      context_start_token: 0\n  ```\n\n- **On-policy alg**: We provide implementations of 4 on-policy algorithms: PPO, NLPO, A2C and TRPO adapted from [stable-baselines3](https://github.com/DLR-RM/stable-baselines3) tailored to work with NLP tasks which can be used out-of-the-box with either a causal policy or a seq2seq LM policy. (See how to create your own [on-policy algorithm](#adding-custom-on-policy-algorithms) or [policy](#adding-custom-policies))\n  - We also provide a supervised [trainer](https://github.com/allenai/RL4LMs/blob/2863116cd5860e4a4106a76486e70bfac25df2ba/rl4lms/envs/text_generation/training_utils.py#L225) for benchmarking purposes. Supervised Warm start models are already uploaded to Huggingface Hub and specified in the respective config files.\n  - Hyper-parameters for the algorithm can be specified at `alg/args`. \n  - Further, all RL algorithms use adaptive KL controller to keep the LM close to original LM by setting initial KL co-efficient (`alg/kl_div/coeff`) and target KL (`alg/kl_div/target_kl`). \n  - We support two types of LM policy: **causal LM policy** (for decoder only models) and **seq2seq LM policy** (for encoder-decoder models). Further for NLPO, we also provide maskable variants of these. Policy implementations can be found [here](https://github.com/allenai/RL4LMs/blob/main/rl4lms/envs/text_generation/policy.py) in and it can be attached to algorithms by specifying `alg/policy/id` and `alg/policy/args`\n\n    ```yaml\n    alg:\n      id: ppo\n      args: \n        n_steps: 512\n        batch_size: 64\n        verbose: 1\n        learning_rate: 0.000002\n        n_epochs: 5\n        ent_coef: 0.0\n      kl_div:\n        coeff: 0.001\n        target_kl: 0.2\n      policy:\n        id: seq2seq_lm_actor_critic_policy\n        args:\n          model_name: t5-base\n          apply_model_parallel: True\n          prompt_truncation_side: \"right\"\n          generation_kwargs:\n            do_sample: True\n            top_k: 50\n            min_length: 50\n            max_new_tokens: 100          \n    ```\n\n- **Trainer Config**: We provide an [On-policy trainer](https://github.com/allenai/RL4LMs/blob/2863116cd5860e4a4106a76486e70bfac25df2ba/rl4lms/envs/text_generation/training_utils.py#L126) - a feature-complete wrapper that instantiates building blocks from their corresponding configs and provides an outer training loop consisting of *train* and *eval* iterations `train_evaluation/n_iters`. \n  - Each iteration corresponds to performing updates with `alg/args/n_steps` x `env/n_envs` of the chosen algorithm. \n  - For every `eval_every` iters, LM is evaluated on validation split using metrics listed in `train_evaluation/metrics` with generation kwargs provided in `train_evaluation/generation_kwargs` (this overrides rollout `alg/policy/generation_kwargs` for inference purposes only)\n\n  ```yaml\n  # train and evaluation\n  train_evaluation:\n    eval_batch_size: 100\n    n_iters: 100\n    eval_every: 10\n    save_every: 1\n    metrics:\n      - id: meteor\n        args: {}\n      - id: rouge\n      - id: bleu\n        args: {}\n      - id: bert_score\n        args:\n          language: en\n      - id: diversity\n        args: {}\n    generation_kwargs: \n      do_sample: True\n      top_k: 0\n      temperature: 0.7\n      min_length: 50\n      max_new_tokens: 100\n  ```\n\n---\n# Custom Building Blocks :wrench:\nRL4LMs provide complete customizability - with respect to adding new tasks/datasets, reward functions, evaluation metric, on-policy algorithms and actor-critic policies.\n\n## Adding dataset\nUsers can create their own datasets by sub-classing [TextGenPool](https://github.com/allenai/RL4LMs/blob/af5a1326578789856ca8550cb5496c9ccc1afdc5/rl4lms/data_pools/text_generation_pool.py#L15) just by overriding `prepare(cls, split: str, **args) -\u003e 'TextGenPool':` method to return an instance of TextGenPool. An example is shown below:\n\n\n ```python\n from rl4lms.data_pools.text_generation_pool import Sample, TextGenPool\n\n class MyDataPool(TextGenPool):\n    @classmethod\n    def prepare(cls, split: str):\n        .. \n        samples = []\n        for ix, item in enumerate(..):\n            sample = Sample(id=f\"{split}_{ix}\",\n                            prompt_or_input_text=item[\"document\"],\n                            references=[item[\"target\"]]\n                            )\n            samples.append(sample)\n        pool_instance = cls(samples)\n        return pool_instance\n\n```\n\n## Adding reward function\nCustom reward funtions can be implemented easily by sub-classing [RewardFunction](https://github.com/allenai/RL4LMs/blob/af5a1326578789856ca8550cb5496c9ccc1afdc5/rl4lms/envs/text_generation/reward.py#L12) (a callable) which takes observation ($s$), next observation ($s'$), action ($a$), done (indicating whether episode is finished) and meta info (containing other information about textual input). Here, [Observation](https://github.com/allenai/RL4LMs/blob/af5a1326578789856ca8550cb5496c9ccc1afdc5/rl4lms/envs/text_generation/observation.py#L11) is a data class object consisting of generated text (at a particular step), prompt text, context text (at that step), reference text which can be used to compute token-level or sentence level rewards.\n\n\n ```python\nfrom rl4lms.envs.text_generation.observation import Observation\nfrom rl4lms.envs.text_generation.reward import RewardFunction\n\n\nclass MyRewardFunction(RewardFunction):\n    def __init__(self, *args) -\u003e None:\n        super().__init__()\n\n    def __call__(self, prev_observation: Observation,\n                 action: int,\n                 current_observation: Observation,\n                 done: bool,\n                 meta_info: Dict[str, Any] = None) -\u003e float:\n        if done:\n            reward = ..\n            return reward\n        return 0\n ```\n\n:bulb:\nIn addition to traditional NLG metrics, for quick prototyping, we provide two synthetic reward functions which trains LMs to [generate numbers](https://github.com/allenai/RL4LMs/blob/af5a1326578789856ca8550cb5496c9ccc1afdc5/rl4lms/envs/text_generation/test_reward.py#L8) in increasing order and [generate dates](https://github.com/allenai/RL4LMs/blob/af5a1326578789856ca8550cb5496c9ccc1afdc5/rl4lms/envs/text_generation/test_reward.py#L54). These can be used to quickly test different algorithms and policies. Corresponding configs can be found here ([numbers](https://github.com/allenai/RL4LMs/tree/main/scripts/training/task_configs/synthetic_generate_increasing_numbers), [dates](https://github.com/allenai/RL4LMs/tree/main/scripts/training/task_configs/synthetic_generate_dates))\n\n\n## Adding custom metrics\nUsers can create their own evaluation metric which then will be used to periodically evaluate the model on validation split of dataset. This can be done by sub-classing [BaseMetric](https://github.com/allenai/RL4LMs/blob/af5a1326578789856ca8550cb5496c9ccc1afdc5/rl4lms/envs/text_generation/metric.py#L20) which takes prompt texts, generated texts, reference texts, meta_infos, current LM model, split name as inputs and returns a dict with metric name as key and value consisting of tuple of sentence-level scores and corpus level scores. An example is as follows:\n\n ```python\n\nfrom rl4lms.envs.text_generation.metric import BaseMetric\n\nclass MyMetric(BaseMetric):\n    def __init__(self) -\u003e None:\n        super().__init__()\n\n    def compute(self,\n                prompt_texts: List[str],\n                generated_texts: List[str],\n                reference_texts: List[List[str]],\n                meta_infos: List[Dict[str, Any]] = None,\n                model: PreTrainedModel = None,\n                split_name: str = None):\n        metric_dict = {\n            \"custom_metrics/my_metric\": ([0.4, 0.7, 0.9], 0.7)\n        }\n        return metric_dict\n ```\n\n## Adding custom on-policy algorithms\n\nIn addition to supported on-policy algorithms (PPO, NLPO, A2C,TRPO), users can implement their own on-policy algorithms with ease by sub-classing stable-baselines3's [OnPolicyAlgorithm](https://github.com/DLR-RM/stable-baselines3/blob/a697401e032dd4fecbbd4162755ddd707df980d3/stable_baselines3/common/on_policy_algorithm.py#L20). Since we provide [wrappers](https://github.com/allenai/RL4LMs/blob/af5a1326578789856ca8550cb5496c9ccc1afdc5/rl4lms/envs/text_generation/alg_wrappers.py#L67) for on-policy algorithms that handles rollouts using LM policies, environment, computing rewards etc, users just need to implement `train()` method with custom loss functions. \n\n```python\nfrom stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm\n\nclass MyOnPolicyAlgorithm(OnPolicyAlgorithm):\n    def __init__(**args):\n        super().__init__(**args)\n\n    def train(self) -\u003e None:\n        # train for n_epochs epochs\n        for epoch in range(self.n_epochs):\n            # Do a complete pass on the rollout buffer\n            for rollout_data in self.rollout_buffer.get(self.batch_size):\n              # compute loss\n```\n\n## Adding custom policies\n\nWe provide LM based actor-critic policy [implementations](https://github.com/allenai/RL4LMs/blob/main/rl4lms/envs/text_generation/policy.py) that wraps causal LM and seq2seq LMs. These can be also extended (for eg: use a different critic architecture) by overriding appropriate methods (eg. `evaluate_actions()`)\n\n## Registry\nFinally, just register your custom components by adding them to corresponding [registry](https://github.com/allenai/RL4LMs/blob/main/rl4lms/envs/text_generation/registry.py), after which they can be used directly from configs similar to pre-defined components :wave:\n\n## Crowdsourcing templates\n\nWe have provided the crowdsourcing templates we used on mechanical turk, along with example inputs in `scripts/crowdworking_templates`. You might find these a helpful starting point either for evaluating your own model's generations, or for gathering training data for a learned reward function.\n\n---\n\n# Logging and Experiment Results\n\nAdditionally, we support WANDB logging and warm-starting of training by storing checkpoints and other training artifacts in a user-specified path. This is especially useful for running preemptible jobs on large, scheduled clusters.\n\nArtifacts include (1) jsonl file containing rollout infos at specified intervals (2) jsonl file containing training infos at specified intervals (3) jsonl file containing validation metrics at specified intervals (4) jsonl file containing test metrics before and after training (5) json file with validation predictions at specified intervals (6) json file with test predictions before and after training (7) trained LM model (8) config json used to run the experiment\n\nComplete usage is as follows:\n\n```bash \nWANDB_API_KEY=\u003cYOUR-WANDB-API-KEY-HERE\u003e  python scripts/training/train_text_generation.py \\\n--config_path \u003cPATH-TO-CONFIG-FILE\u003e \\\n--experiment_name \u003cEXPERIMENT-NAME\u003e \\\n--base_path_to_store_results \u003cPATH-TO-STORE-RESULTS\u003e \\\n--log_to_wandb\n```\n\n---\n\n# Citation\n\n```bibtex\n@inproceedings{Ramamurthy2022IsRL,\n  title={Is Reinforcement Learning (Not) for Natural Language Processing?: Benchmarks, Baselines, and Building Blocks for Natural Language Policy Optimization},\n  author={Rajkumar Ramamurthy and Prithviraj Ammanabrolu and Kiant{\\'e} Brantley and Jack Hessel and Rafet Sifa and Christian Bauckhage and Hannaneh Hajishirzi and Yejin Choi},\n  journal={arXiv preprint arXiv:2210.01241},\n  url={https://arxiv.org/abs/2210.01241},\n  year={2022}\n}\n```\n\n# Questions/Discussion/Ideas?\n\nFor discussion, questions, ideas exchange, join our slack channel \n[![Slack](https://img.shields.io/badge/Slack-4A154B?style=for-the-badge\u0026logo=slack\u0026logoColor=white)](https://join.slack.com/t/slack-1sa3880/shared_invite/zt-1idqlnbnm-NIiZeMIOpYReXfX9uIT_PA)\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fallenai%2Frl4lms","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fallenai%2Frl4lms","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fallenai%2Frl4lms/lists"}