added OpenRewards

Author: shamanez

examples/agentic_demo/openreward_endless_terminals_IPA_qwen35_2b.yaml

@@ -0,0 +1,223 @@
+# OpenReward + EndlessTerminals + IPA chunk-level loss config
+# Based on agent_val_rock_swe_IPA_qwen35_2b.yaml but replaces ROCK with OpenReward.
+# The model IS the agent — no iflow, no sandbox, no anti_call_llm.
+#
+# Usage:
+#   python examples/start_agentic_pipeline.py \
+#     --config_path agentic_demo \
+#     --config_name openreward_endless_terminals_IPA_qwen35_2b
+#
+# For REINFORCE instead of IPA, change:
+#   pg_variant: "vanilla"
+#   adv_estimator: "reinforce"
+
+defaults:
+  - ../config/deepspeed_zero@_here_
+  - ../config/deepspeed_zero2@_here_
+  - ../config/deepspeed_zero3@_here_
+  - ../config/deepspeed_zero3_cpuoffload@_here_
+
+hydra:
+  run:
+    dir: .
+  output_subdir: null
+
+exp_name: "openreward_endless_terminals_IPA_qwen35_2b"
+seed: 42
+
+logging_dir: ./output/logs
+output_dir: ./output
+model_name: ${exp_name}-${now:%Y%m%d_%H%M%S}
+rollout_dump_dir: /home/ubuntu/ALE-latest/ROLL-personal/output/rollout_dump
+system_envs:
+  USE_MODELSCOPE: '1'
+
+checkpoint_config:
+  type: file_system
+  output_dir: /data
+
+num_gpus_per_node: 8
+rpc_timeout: 72000
+
+max_steps: 10
+save_steps: 50
+logging_steps: 1
+eval_steps: 0
+resume_from_checkpoint: false
+
+async_generation_ratio: 1
+parse_tool_call_parameter_to_dict: true
+skip_mock_system_prompt: true
+
+track_with: wandb
+tracker_kwargs:
+  api_key: "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"
+  project: roll-agentic
+  name: ${exp_name}
+
+rollout_batch_size: 16
+val_batch_size: 1
+sequence_length: 32768
+
+max_tokens_per_step: 4096
+
+# --- IPA chunk-level loss config (identical to ROCK config) ---
+ppo_epochs: 1
+adv_estimator: "step_reinforce"
+batch_adjust_mode: "random_sample"
+step_reward_gamma: 0.95                  # chunk-level temporal credit: G_k = 0.95^{K-k} * R_final
+ipa_failure_reward: ~                    # positive rewards only — no penalty for failures
+ratio_type: "segment"                    # chunk-level geometric mean IS ratio
+
+init_kl_coef: 0.0
+whiten_advantages: false                 # IPA uses absolute returns, no relative ranking
+entropy_loss_coef: 0
+max_grad_norm: 1.0
+
+# Non-trivial IS ratio: use vLLM rollout logprobs as old_log_probs (not recomputed Megatron)
+force_disable_old_logprobs_recompute: true
+
+# Train-infer mismatch mask (paper Section 8): chunk-level segment filter
+train_infer_correction:
+  is_weight:
+    enabled: false                       # IPA handles IS internally
+  filters:
+    - enabled: true
+      agg_type: segment                  # chunk-level geometric mean mismatch
+      ratio_enabled: true
+      ratio_low: 0.0
+      ratio_high: 5.0                    # threshold H; chunks above are masked out
+
+# --- Model configs (identical to ROCK config) ---
+pretrain: /home/ubuntu/ALE-latest/model-checkpoints/Qwen3.5-2B
+reward_pretrain: /home/ubuntu/ALE-latest/model-checkpoints/Qwen3.5-2B
+actor_train:
+  worker_cls: roll.pipeline.agentic.agentic_actor_pg_worker.ActorWorker
+  pg_variant: "ipa_chunk"                # chunk-level IPA loss (paper Eq. 9)
+  model_args:
+    flash_attn: sdpa
+    attn_implementation: sdpa
+    disable_gradient_checkpointing: false
+    dtype: bf16
+    model_type: ~
+    freeze_module_prefix: vision_model
+  training_args:
+    learning_rate: 1.0e-6
+    weight_decay: 0
+    per_device_train_batch_size: 1
+    gradient_accumulation_steps: 4
+    warmup_steps: 0
+  data_args:
+    template: qwen3_coder
+  strategy_args:
+    strategy_name: megatron_train
+    strategy_config:
+      tensor_model_parallel_size: 2
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 2
+      sequence_parallel: true
+      use_distributed_optimizer: true
+      recompute_granularity: full
+  device_mapping: list(range(0,4))
+  infer_batch_size: 1
+actor_infer:
+  model_args:
+    flash_attn: sdpa
+    attn_implementation: sdpa
+    disable_gradient_checkpointing: true
+    dtype: bf16
+  generating_args:
+    max_new_tokens: ${max_tokens_per_step}
+    top_p: 1.0
+    top_k: 50
+    num_beams: 1
+    temperature: 1.0
+    num_return_sequences: 1
+    stop_strings: ["</tool_call>"]
+    include_stop_str_in_output: true
+  data_args:
+    template: qwen3_coder
+  strategy_args:
+    strategy_name: vllm
+    strategy_config:
+      gpu_memory_utilization: 0.6
+      block_size: 16
+      load_format: auto
+      tensor_parallel_size: 1
+      max_model_len: 32768
+  device_mapping: list(range(0,8))
+
+reference:
+  model_args:
+    attn_implementation: sdpa
+    disable_gradient_checkpointing: true
+    dtype: bf16
+    model_type: ~
+    freeze_module_prefix: vision_model
+  data_args:
+    template: qwen3_coder
+  strategy_args:
+    strategy_name: megatron_infer
+    strategy_config:
+      tensor_model_parallel_size: 2
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 2
+  device_mapping: list(range(0,4))
+  infer_batch_size: 1
+
+# IPA uses absolute returns — no reward normalization
+reward_normalization:
+  grouping: traj_group_id
+  method: identity
+
+# --- Environment config (OpenReward — replaces ROCK) ---
+max_actions_per_traj: 16                 # 16 turns max (matches ROCK config)
+env_manager_cls: roll.pipeline.agentic.env_manager.agent_native_env_manager.AgentNativeStepEnvManager
+
+train_env_manager:
+  max_env_num_per_worker: 1
+  num_env_groups: 1
+  group_size: 1                         # 16 trajectories per batch
+  tags: [OpenRewardEndlessTerminalsTrain]
+  num_groups_partition: [1]
+
+val_env_manager:
+  max_env_num_per_worker: 1
+  num_env_groups: 1
+  group_size: 1
+  tags: [OpenRewardEndlessTerminalsVal]
+  num_groups_partition: [1]
+
+custom_envs:
+  OpenRewardEndlessTerminalsTrain:
+    env_type: "openreward_env"
+    max_steps: ${max_actions_per_traj}
+    max_tokens_per_step: ${max_tokens_per_step}
+    env_manager_cls: ${env_manager_cls}
+    agent_system_template: "unused — system prompt built dynamically from OpenReward tool specs"
+    agent_template: "unused — observation is full message list from OpenRewardEnv"
+    env_config:
+      environment_name: "kanishk/EndlessTerminals"
+      split: "train"
+      mode: "train"
+      max_steps: ${max_actions_per_traj}
+      reward_reduction: "sum"
+      nonterminal_reward: 0.0
+      retry_max_attempts: 3
+      retry_backoff_seconds: 5.0
+  OpenRewardEndlessTerminalsVal:
+    env_type: "openreward_env"
+    max_steps: ${max_actions_per_traj}
+    max_tokens_per_step: ${max_tokens_per_step}
+    env_manager_cls: ${env_manager_cls}
+    agent_system_template: "unused"
+    agent_template: "unused"
+    env_config:
+      environment_name: "kanishk/EndlessTerminals"
+      split: "train"
+      mode: "val"
+      max_steps: ${max_actions_per_traj}
+      reward_reduction: "sum"
+      nonterminal_reward: 0.0

roll/pipeline/agentic/env/__init__.py

@@ -17,7 +17,10 @@
 gem.register("deepeyes", entry_point="roll.pipeline.agentic.env.deepeyes:DeepEyesEnv")
 gem.register("rock_tb_native_env", entry_point="roll.pipeline.agentic.env.sandbox.rock_tb_native_env:RockTBNativeEnv")
 
-
+try:
+    gem.register("openreward_env", entry_point="roll.pipeline.agentic.env.openreward:OpenRewardEnv")
+except Exception as e:
+    logger.info(f"Failed to register openreward_env: {e}")
 
 try:
     # add webshop-minimal to PYTHONPATH

roll/pipeline/agentic/env/openreward/__init__.py

@@ -0,0 +1,3 @@
+from .openreward_env import OpenRewardEnv
+
+__all__ = ["OpenRewardEnv"]