feat: add stage-specific exploration and training controls

docs/input.md

@@ -127,7 +127,39 @@ The {dargs:argument}`"configurations"<explore[lmp]/configurations>` provides the
 - Automatic alloy configuration generator. See {ref}`the detailed doc<explore[lmp]/configurations[alloy]>` for the allowed keys.
 - Configurations load from files. See {ref}`the detailed doc<explore[lmp]/configurations[file]>` for the allowed keys.
 
-The {dargs:argument}`"stages"<explore[lmp]/stages>` defines the exploration stages. It is of type `list[list[dict]]`. The outer `list` enumerate the exploration stages, the inner list enumerate the task groups of the stage. Each `dict` defines a stage. See {ref}`the full documentation of the task group<task_group_sec>` for writting task groups.
+The {dargs:argument}`"stages"<explore[lmp]/stages>` defines the exploration stages. The legacy form is `list[list[dict]]`: the outer list enumerates stages and each inner list contains that stage's task groups. See {ref}`the full documentation of the task group<task_group_sec>` for writing task groups.
+
+When different stages need different convergence or labeling limits, a stage may instead be a dictionary with a required `task_groups` list and optional `convergence`, `max_numb_iter`, `fatal_at_max`, and `task_max` overrides. Missing overrides inherit the global values from `explore` and `fp`:
+
+```json
+"stages": [
+    {
+        "task_groups": [
+            {
+                "type": "lmp-md",
+                "ensemble": "nvt",
+                "nsteps": 50,
+                "temps": [50],
+                "trj_freq": 10,
+                "conf_idx": [0],
+                "n_sample": 3
+            }
+        ],
+        "convergence": {
+            "type": "adaptive-lower",
+            "rate_candi_f": 0.15,
+            "level_f_hi": 5.0,
+            "n_checked_steps": 3,
+            "conv_tolerance": 0.005
+        },
+        "max_numb_iter": 2,
+        "fatal_at_max": false,
+        "task_max": 4000
+    }
+]
+```
+
+Training loss endpoints can also be specialized without changing the template script. `train.config.finetune_end_pref_e/f/v` apply to the initial finetuning run, while `train.config.init_model_end_pref_e/f/v` apply to subsequent init-model training. An unset value keeps the corresponding `limit_pref_e/f/v` from the template.
 
 The {dargs:argument}`"n_sample"<task_group[lmp-md]/n_sample>` tells the number of confgiruations randomly sampled from the set picked by {dargs:argument}`"conf_idx"<task_group[lmp-md]/conf_idx>` from {dargs:argument}`"configurations"<explore[lmp]/configurations>` for each exploration task. All configurations has the equal possibility to be sampled. The default value of `"n_sample"` is `null`, in this case all picked configurations are sampled. In the example, we have 3 samples for stage 0 task group 0 and 2 thermodynamic states (NVT, T=50 and 100K), then the task group has 3x2=6 NVT DPMD tasks.
 

dpgen2/entrypoint/args.py

@@ -206,9 +206,10 @@ def lmp_args():
     doc_convergence = "The method of convergence check."
     doc_configuration = "A list of initial configurations."
     doc_stages = (
-        "The definition of exploration stages of type `List[List[ExplorationTaskGroup]`. "
-        "The outer list provides the enumeration of the exploration stages. "
-        "Then each stage is defined by a list of exploration task groups. "
+        "The definition of exploration stages. A legacy stage is a list of "
+        "exploration task groups. A stage dictionary stores those groups under "
+        "`task_groups` and may override `convergence`, `max_numb_iter`, "
+        "`fatal_at_max`, and `task_max` for that stage. "
         "Each task group is described in :ref:`the task group definition<task_group_sec>` "
     )
     doc_filters = "A list of configuration filters"
@@ -249,7 +250,7 @@ def lmp_args():
             doc=doc_configuration,
             alias=["configuration"],
         ),
-        Argument("stages", List[List[dict]], optional=False, doc=doc_stages),
+        Argument("stages", list, optional=False, doc=doc_stages),
         Argument(
             "filters",
             list,
@@ -299,9 +300,10 @@ def caly_args():
     doc_convergence = "The method of convergence check."
     doc_configuration = "A list of initial configurations."
     doc_stages = (
-        "The definition of exploration stages of type `List[List[ExplorationTaskGroup]`. "
-        "The outer list provides the enumeration of the exploration stages. "
-        "Then each stage is defined by a list of exploration task groups. "
+        "The definition of exploration stages. A legacy stage is a list of "
+        "exploration task groups. A stage dictionary stores those groups under "
+        "`task_groups` and may override `convergence`, `max_numb_iter`, "
+        "`fatal_at_max`, and `task_max` for that stage. "
         "Each task group is described in :ref:`the task group definition<task_group_sec>` "
     )
     doc_filters = "A list of configuration filters"
@@ -342,7 +344,7 @@ def caly_args():
             doc=doc_configuration,
             alias=["configuration"],
         ),
-        Argument("stages", List[List[dict]], optional=False, doc=doc_stages),
+        Argument("stages", list, optional=False, doc=doc_stages),
         Argument(
             "filters",
             list,
@@ -400,9 +402,10 @@ def diffcsp_args():
     doc_output_nopbc = "Remove pbc of the output configurations"
     doc_convergence = "The method of convergence check."
     doc_stages = (
-        "The definition of exploration stages of type `List[List[ExplorationTaskGroup]`. "
-        "The outer list provides the enumeration of the exploration stages. "
-        "Then each stage is defined by a list of exploration task groups. "
+        "The definition of exploration stages. A legacy stage is a list of "
+        "exploration task groups. A stage dictionary stores those groups under "
+        "`task_groups` and may override `convergence`, `max_numb_iter`, "
+        "`fatal_at_max`, and `task_max` for that stage. "
         "Each task group is described in :ref:`the task group definition<task_group_sec>` "
     )
     doc_filters = "A list of configuration filters"
@@ -432,7 +435,7 @@ def diffcsp_args():
             optional=False,
             doc=doc_convergence,
         ),
-        Argument("stages", List[List[dict]], optional=False, doc=doc_stages),
+        Argument("stages", list, optional=False, doc=doc_stages),
         Argument(
             "filters",
             list,

dpgen2/entrypoint/submit.py

@@ -298,36 +298,99 @@ def get_conf_filters(config):
     return conf_filters
 
 
+def _normalize_exploration_stage(stage_config, explore_config, fp_task_max):
+    """Resolve task groups and effective controls for one exploration stage.
+
+    Legacy configurations represent a stage directly as a list of task-group
+    dictionaries. Metadata-bearing stage dictionaries keep that list under
+    ``task_groups`` and override only the controls that need to differ from the
+    global exploration and FP defaults.
+    """
+    if isinstance(stage_config, list):
+        task_groups = stage_config
+        stage_options = {}
+    elif isinstance(stage_config, dict) and "task_groups" in stage_config:
+        stage_options = deepcopy(stage_config)
+        task_groups = stage_options.pop("task_groups")
+    elif isinstance(stage_config, dict):
+        # Keep accepting a single task-group dictionary as one legacy stage.
+        task_groups = [stage_config]
+        stage_options = {}
+    else:
+        raise TypeError("each exploration stage must be a list or dictionary")
+
+    if isinstance(task_groups, dict):
+        task_groups = [task_groups]
+    if not isinstance(task_groups, list) or not all(
+        isinstance(task_group, dict) for task_group in task_groups
+    ):
+        raise TypeError("stage task_groups must be a list of dictionaries")
+
+    convergence = deepcopy(
+        stage_options.get("convergence", explore_config["convergence"])
+    )
+    max_numb_iter = stage_options.get("max_numb_iter", explore_config["max_numb_iter"])
+    fatal_at_max = stage_options.get("fatal_at_max", explore_config["fatal_at_max"])
+    stage_task_max = stage_options.get("task_max", fp_task_max)
+    if not isinstance(convergence, dict):
+        raise TypeError("stage convergence must be a dictionary")
+    if max_numb_iter is not None and not isinstance(max_numb_iter, int):
+        raise TypeError("stage max_numb_iter must be an integer or null")
+    if not isinstance(fatal_at_max, bool):
+        raise TypeError("stage fatal_at_max must be a boolean")
+    if stage_task_max is not None and not isinstance(stage_task_max, int):
+        raise TypeError("stage task_max must be an integer or null")
+    return task_groups, convergence, max_numb_iter, fatal_at_max, stage_task_max
+
+
+def _make_stage_selector(
+    convergence,
+    render,
+    fp_task_max,
+    conf_filters,
+):
+    """Build an independent report and selector for an exploration stage."""
+    conv_style = convergence.pop("type")
+    report = conv_styles[conv_style](**convergence)
+    return ConfSelectorFrames(
+        render,
+        report,
+        fp_task_max,
+        conf_filters,
+    )
+
+
 def make_naive_exploration_scheduler_without_conf(config, explore_style):
     model_devi_jobs = config["explore"]["stages"]
     fp_task_max = config["fp"]["task_max"]
-    max_numb_iter = config["explore"]["max_numb_iter"]
-    fatal_at_max = config["explore"]["fatal_at_max"]
-    convergence = config["explore"]["convergence"]
     output_nopbc = config["explore"]["output_nopbc"]
     conf_filters = get_conf_filters(config["explore"]["filters"])
     scheduler = ExplorationScheduler()
-    # report
-    conv_style = convergence.pop("type")
-    report = conv_styles[conv_style](**convergence)
     # trajectory render, the format of the output trajs are assumed to be lammps/dump
     render = TrajRenderLammps(nopbc=output_nopbc)
-    # selector
-    selector = ConfSelectorFrames(
-        render,
-        report,
-        fp_task_max,
-        conf_filters,
-    )
 
-    for job_ in model_devi_jobs:
-        if not isinstance(job_, list):
-            job = [job_]
-        else:
-            job = job_
+    for stage_config in model_devi_jobs:
+        (
+            job,
+            convergence,
+            max_numb_iter,
+            fatal_at_max,
+            stage_task_max,
+        ) = _normalize_exploration_stage(
+            stage_config,
+            config["explore"],
+            fp_task_max,
+        )
+        selector = _make_stage_selector(
+            convergence,
+            render,
+            stage_task_max,
+            conf_filters,
+        )
         # stage
         stage = ExplorationStage()
         for jj in job:
+            jj = deepcopy(jj)
             if "calypso" in explore_style:
                 jconf = caly_normalize(jj)
                 # make task group
@@ -361,39 +424,41 @@ def make_lmp_naive_exploration_scheduler(config):
     type_map = config["inputs"]["type_map"]
     numb_models = config["train"]["numb_models"]
     fp_task_max = config["fp"]["task_max"]
-    max_numb_iter = config["explore"]["max_numb_iter"]
-    fatal_at_max = config["explore"]["fatal_at_max"]
-    convergence = config["explore"]["convergence"]
     output_nopbc = config["explore"]["output_nopbc"]
     conf_filters = get_conf_filters(config["explore"]["filters"])
     use_ele_temp = config["inputs"]["use_ele_temp"]
     scheduler = ExplorationScheduler()
-    # report
-    conv_style = convergence.pop("type")
-    report = conv_styles[conv_style](**convergence)
     render = TrajRenderLammps(nopbc=output_nopbc, use_ele_temp=use_ele_temp)
-    # selector
-    selector = ConfSelectorFrames(
-        render,
-        report,
-        fp_task_max,
-        conf_filters,
-    )
 
     sys_configs_lmp = []
     for sys_config in sys_configs:
+        sys_config = deepcopy(sys_config)
         conf_style = sys_config.pop("type")
         generator = conf_styles[conf_style](**sys_config)
         sys_configs_lmp.append(generator.get_file_content(type_map))
 
-    for job_ in model_devi_jobs:
-        if not isinstance(job_, list):
-            job = [job_]
-        else:
-            job = job_
+    for stage_config in model_devi_jobs:
+        (
+            job,
+            convergence,
+            max_numb_iter,
+            fatal_at_max,
+            stage_task_max,
+        ) = _normalize_exploration_stage(
+            stage_config,
+            config["explore"],
+            fp_task_max,
+        )
+        selector = _make_stage_selector(
+            convergence,
+            render,
+            stage_task_max,
+            conf_filters,
+        )
         # stage
         stage = ExplorationStage()
         for jj in job:
+            jj = deepcopy(jj)
             jconf = normalize_lmp_task_group_config(jj)
             n_sample = jconf.pop("n_sample")
             ##  ignore the expansion of sys_idx