Sensitivity analysis MVP: per-episode recording + NPE / MNPE / empirical analyzers

[cvolkcvolk]

Summary

sensitivity analysis for Arena evals: per-episode JSONL writer in eval_runner + offline analyzer that infers posteriors over input factors.

Detailed description

• Opt-in per-episode writer in eval_runner (--episode_summary). Logs the full arena_env_args and the task's registered outcomes per recorded demo as a JSONL row. Existing
  eval workflows are unchanged when the flag is absent.
• Three-way analyzer dispatch via make_analyzer: NPE for all-continuous, MNPE for mixed continuous+categorical, Empirical (frequency table) for pure-categorical schemas. Class
  hierarchy is BaseAnalyzer -> PosteriorAnalyzer -> NPEAnalyzer / MNPEAnalyzer, with EmpiricalAnalyzer separate.
• Hand-authored factors.yaml schema declares slice + factor types + outcome names; the analyzer is the sole consumer. The eval-side arena_env_args field in the JSONL is
  provenance-honest, leaving room for sibling fields (sim_state, variation_draws) as MVP-3 and the variation system land.
• Inference (analyzer.py) and rendering (plotting.py) are decoupled: analyzers expose continuous_marginal_density / categorical_marginal_probs queries; plot_marginal
  dispatches by factor type and renders matplotlib.
• synthetic_data_continuous.py and synthetic_data_categorical.py generate paired JSONL + factors.yaml for offline smoke testing; both validated end-to-end.

[isaaclab-review-bot]

Code Review: MVP-1 Per-Episode Sensitivity Recording + NPE Analyzer

This is a substantial and well-structured PR that lays solid data-layer and analysis-layer foundations for the sensitivity analysis workstream. The architecture is clean, with good separation of concerns between the eval-side writer, dataset loader, analyzer hierarchy, and plotting module.

---

✅ Strengths

Architecture & Design:

• The analyzer hierarchy (BaseAnalyzer → PosteriorAnalyzer → NPE/MNPE + EmpiricalAnalyzer) with make_analyzer factory is well-designed and extensible
• Clean separation between inference (analyzer.py) and rendering (plotting.py) follows good software engineering principles
• The decision to log the full arena_env_args per episode (rather than just --factor_keys) is a smart choice that eliminates the drift bug class where CLI flags and factors.yaml could disagree
• Deferred imports for pxr-touching modules follow the established policy_runner.py:107 pattern correctly

Implementation Quality:

• Comprehensive docstrings with clear "what this module does in plain English" sections
• The JSONL wire format with arena_env_args + outcomes is honest about provenance and leaves room for future extension
• Continuous-first column ordering in theta matches sbi's MNPE convention
• The empirical rug in plots provides valuable sanity-checking for humans

---

🔍 Issues & Suggestions

1. Silent Failure: Empty arena_env_args_dict on Non-Dict Jobs

File: job_manager.py (lines 46-47, 110)

self.arena_env_args_dict = arena_env_args_dict if arena_env_args_dict is not None else {}

The from_dict class method passes data["arena_env_args"] which could be None or missing entirely. If a job config omits arena_env_args, this silently produces an empty dict, and the episode writer will log rows with empty arena_env_args blocks. Consider adding validation:

arena_env_args_dict=data.get("arena_env_args") or {}

Or raise explicitly if the key is required for sensitivity workflows.

---

2. Potential Division by Zero in Empirical Analyzer

File: analyzer.py (line 309)

if total_rate > 0:
    return empirical_rates / total_rate
return np.full(num_choices, 1.0 / num_choices)

This correctly handles the total_rate == 0 case, but there's a subtle issue: if category_mask.any() is False for all categories (i.e., the data has no rows matching any declared choice), empirical_rates stays all zeros, and you return a uniform distribution. This is reasonable but worth documenting in the docstring.

---

3. Missing Validation for factor.choices Consistency

File: dataset.py (lines 247-253)

The categorical encoding loop validates that each row's value is in choice_to_code, but doesn't validate that factor.choices is exhaustive relative to the data. If the data contains a category not declared in factors.yaml, you get a clear assertion error — good. But the reverse (declared choices not appearing in data) silently passes and could lead to misleading posteriors. Consider logging a warning:

observed_choices = set(row["arena_env_args"][factor.name] for row in self.rows)
missing_in_data = set(factor.choices) - observed_choices
if missing_in_data:
    print(f"[WARN] Factor {factor.name!r}: declared choices {missing_in_data} never appear in data")

---

4. Hardcoded Threshold >= 0.5 for Binary Outcome Classification

File: analyzer.py (line 302), plotting.py (lines 76, 95, 172)

The threshold >= 0.5 for classifying outcomes as "success" is hardcoded in multiple places. This works for binary 0/1 outcomes but could be misleading for metrics that return continuous values in [0, 1]. Consider:

• Documenting this assumption clearly in the module docstrings
• Or exposing it as a configurable parameter in the CLI

---

5. Range Inference Could Produce Degenerate Priors

File: dataset.py (lines 188-200)

When factor.range is inferred from observed data:

factor.range = [[min(observed_values), max(observed_values)]]

If all rows have the same value (e.g., a misconfigured sweep with only one intensity), low == high and BoxUniform will produce a degenerate prior. Consider adding a guard:

if min(observed_values) == max(observed_values):
    print(f"[WARN] Factor {factor.name!r} has constant value {min(observed_values)} in data; "
          "posterior will be degenerate. Check your sweep config.")

---

6. Thread Safety for JSONL Appends

File: episode_writer.py (line 76)

with open(output_path, "a", encoding="utf-8") as jsonl_output:

If multiple eval jobs run in parallel writing to the same --episode_summary file, appends could interleave. While Isaac Sim eval typically runs serially, the code doesn't document this assumption. Consider either:

• Adding a note that parallel execution requires separate output files
• Or using file locking (e.g., fcntl.flock)

---

7. NotImplementedError for Vector Factors Lacks Migration Path

Files: dataset.py (lines 193, 226)

The NotImplementedError for dim > 1 factors is good for reserving the extension point, but the error message could be more actionable:

raise NotImplementedError(
    f"Vector continuous factors (dim > 1) are reserved for future implementation; "
    f"factor {factor.name!r} has dim={factor.dim}. For now, decompose into separate 1D factors."
)

---

📝 Minor Suggestions

1. analyze_sensitivity.py: The script exits silently on success. Consider printing the figure path at the end: print(f"[INFO] Figure saved → {args.figure_path}")

2. Type hints: episode_writer.py uses TYPE_CHECKING for Job but the function signature lacks the actual type hint for env. Consider adding env: Any or the appropriate gym wrapper type.

3. YAML schema validation: The FactorSchema.from_yaml validates structure but doesn't catch typos in factor types. A choices: [...] on a type: continuous factor would be silently ignored.

---

🧪 Test Coverage Assessment

The PR includes synthetic data generators (synthetic_data_continuous.py, synthetic_data_categorical.py) that serve as smoke tests. These are valuable for manual validation but:

• No automated unit tests are included for SensitivityDataset, make_analyzer, or the plotting module
• Edge cases (empty JSONL, single-row data, all-same-outcome data) aren't covered
• Consider adding pytest tests in a future PR, particularly for:
  • FactorSchema.from_yaml with malformed YAML
  • SensitivityDataset with missing/extra keys
  • EmpiricalAnalyzer.categorical_marginal_probs with zero counts

---

🎯 Verdict

Overall: This is a solid MVP-1 implementation with good architecture and thoughtful design. The concerns raised are mostly defensive hardening and documentation improvements rather than blockers.

Recommendation: Approve with suggestions. The issues identified are quality-of-life improvements that could be addressed in follow-up PRs.

---

Files Reviewed

• isaaclab_arena/analysis/__init__.py
• isaaclab_arena/analysis/sensitivity/__init__.py
• isaaclab_arena/analysis/sensitivity/analyzer.py
• isaaclab_arena/analysis/sensitivity/dataset.py
• isaaclab_arena/analysis/sensitivity/episode_writer.py
• isaaclab_arena/analysis/sensitivity/plotting.py
• isaaclab_arena/analysis/sensitivity/synthetic_data_categorical.py
• isaaclab_arena/analysis/sensitivity/synthetic_data_continuous.py
• isaaclab_arena/evaluation/eval_runner.py
• isaaclab_arena/evaluation/eval_runner_cli.py
• isaaclab_arena/evaluation/job_manager.py
• isaaclab_arena/scripts/analyze_sensitivity.py
• isaaclab_arena_environments/eval_jobs_configs/*.yaml
• isaaclab_arena_environments/eval_jobs_configs/*.json
• setup.py

---

Update (38baa56): Reviewed incremental changes since 671e0c4. The new commits contain only style/formatting improvements:

• plotting.py: Import reordering (moved numpy before pathlib), removed string quotes from type hints (using direct BaseAnalyzer/FactorSpec types), minor line-length adjustments
• synthetic_data_categorical.py: Code reformatting for line length compliance, dict literal style changes
• pick_up_object_sweep_factors.yaml: Added standard Apache-2.0 copyright header

✅ No functional changes. All prior review comments remain applicable. The formatting changes are consistent and improve readability.

---

Update (74585f1): Reviewed incremental changes since 38baa56.

🆕 New Feature: CameraObsVideoRecorder Integration

The new commits add --camera_video flag integration (eval_runner.py lines 191-199):

if args_cli.camera_video:
    print(f"[INFO] Recording per-camera videos for job '{job.name}' -> {job_video_folder}")
    env = CameraObsVideoRecorder(
        env,
        video_folder=job_video_folder,
        step_trigger=lambda step: step == 0,
        video_length=video_length,
    )

This properly imports and uses CameraObsVideoRecorder from isaaclab_arena.evaluation.camera_video (imported at line 17). The implementation:

• ✅ Correctly reuses job_video_folder (computed when either --video or --camera_video is set)
• ✅ Applies the wrapper independently of RecordVideo (can use either or both)
• ✅ Follows the same step_trigger / video_length pattern as RecordVideo

🔍 Observations

1. Import added correctly: CameraObsVideoRecorder import at line 17 is clean and follows the existing import style.

2. Video folder logic refactored: The job_video_folder computation was hoisted out of the if args_cli.video: block to be shared by both video recording methods - good refactoring.

3. RTX ambient light fix in isaaclab_arena_manager_based_env.py: The change disables RTX's built-in ambient light:

render=RenderCfg(
    carb_settings={
        "/rtx/sceneDb/ambientLightIntensity": 0.0,
        "/rtx/sceneDb/ambientLightColor": [0.0, 0.0, 0.0],
    },
),

This ensures USD light prims fully control scene illumination, fixing a confound for vision-policy evals. The inline comment explains the rationale well.

✅ No issues with the incremental changes. The camera video integration follows existing patterns and the ambient light fix is a legitimate rendering correction for sensitivity analysis validity.