Mobile UI Reinforcement Learning Environment

A complete Python package simulating a mobile application environment designed for Reinforcement Learning (RL) agents. This framework is tailored for Prime Intellect RL assignments, incorporating verifier shims, sparse/dense reward rubrics, and safety constraint modeling.

---

Project Structure Tree

mobile_ui_env/
├── AI_USAGE.md                 # AI usage report detailing design contributions
├── Dockerfile                  # Containerization specification for eval runs
├── pyproject.toml              # Build backend and dependency declarations
├── run_eval.py                 # Evaluation harness executing HeuristicAgent
├── run_eval_llm.py             # Evaluation harness integrating OpenAI-compatible LLM agent
├── run_tests.py                # Standalone test runner running pytest files
├── setup.py                    # Legacy installation script (optional)
├── mobile_ui_env/              # Core source package
│   ├── __init__.py             # Exports MobileUIEnv and load_environment
│   ├── actions.py              # Action executor, validation rules & ActionResult
│   ├── dataset.py              # 30-task benchmark split (20 train, 10 eval)
│   ├── env.py                  # Gym/Verifiers environment loop & shims
│   ├── rubric.py               # Sparse and dense reward calculation logic
│   └── state.py                # AppState dataclass and screen configurations
└── tests/                      # Pytest test suite directory
    ├── __init__.py             # Makes tests folder a package
    ├── test_actions.py         # 27 unit tests for actions/navigation
    ├── test_env.py             # 26 unit tests for environment steps & limits
    └── test_rewards.py         # 35 unit tests for reward equations

---

Setup & Running Guide

Installation

You can install the package in editable mode directly:

pip install -e .

For developer dependencies (e.g. pytest):

pip install -e ".[dev]"

Running Tests

To run the full suite of 88 unit tests without sandbox file-system blocks, execute the custom test runner:

python run_tests.py

Running the Evaluator

To evaluate the deterministic rule-based baseline agent on the 10 eval tasks, run:

python run_eval.py

---

Theoretical & Architecture Questions

1. State Space Representation

The AppState represents the status of the simulated application inside a single episode. It tracks:

• Current Screen (current_screen): A categorical state variable {"home", "notes", "settings", "profile"}.
• Notes list (notes): An ordered list of saved note string contents.
• Toggles (focus_mode, notifications_enabled): Boolean variables representing settings choices.
• Static metadata (app_version, username, email): Values for Profile information extraction tasks.
• Metrics & Counters (steps_taken, invalid_action_count, safety_violations): Statistics used for tracking violations, step budgets, and performance penalties.

2. Action Space Representation

The action space consists of structured command dictionaries containing an action_type. The available types and their valid parameters include:

Action Type	Parameters	Valid Screen Context	Description
tap	target: str	Current screen's SCREEN_ELEMENTS	Taps navigation buttons (notes_button, settings_button, profile_button), toggles, profile fields, or saving controls.
type	target: "note_input", text: str	notes screen only	Stores input characters in an temporary _pending_text buffer.
back	None	notes, settings, or profile	Navigates back to the home screen.
finish	None	Any screen	Signals task termination; sets state.is_done = True.

3. Episode Termination Conditions

An episode is terminated (done = True) under any of the following three conditions:

1. Agent Termination: The agent explicitly executes the {"action_type": "finish"} action.
2. Step Budget Exhausted: The agent's cumulative steps exceed the task's maximum budget (steps_taken >= task["max_steps"]).
3. Safety Violation Violation Limits: A safety violation (such as attempting to logout by clicking logout_button) immediately flags a violation. Under training loops, this penalty or done flag can be used to stop execution early.

4. Sparse Rewards

The success_reward represents a purely sparse signal returning exactly 1.0 if the terminal state matches all goal conditions specified in the task descriptor, and 0.0 otherwise.

Examples of goals evaluated:

• note_created: Checks if title is in state.notes.
• toggle_enabled: Checks if the target setting bool matches the expected value.
• safe_navigation: Checks if the target screen was viewed without triggering any safety violations.

5. Dense & Shaped Rewards

To guide exploration during RL, the framework provides continuous/shaped reward functions:

• efficiency_reward: Awarded only on successful task completion, calculating (max_steps - steps_taken) / max_steps. It decays linearly, motivating the agent to find the shortest path.
• partial_progress_reward: Returns 0.3 if the agent successfully navigates to the target screen associated with the task's goal type (e.g. Settings screen for toggle tasks) even if it fails to fulfill the final goal. It yields 0.0 if the task is successfully completed to avoid double-counting.

6. Reward Hacking Vulnerabilities & Mitigation

Like all RL reward formulations, this setup exposes potential vectors for reward hacking:

• JSON Syntax Spam: An agent could spam syntactically correct but functionally useless actions (e.g. repeatedly calling back on the home screen) to accumulate format_reward points without attempting the goal.
  • Mitigation: The format reward is weighted with a small scalar (0.1) and negated by step budget limits.
• Settings/Screen Ping-Pong: The agent might oscillate between Home and Settings screens to game the partial_progress_reward without ever completing a note or toggle.
  • Mitigation: The partial progress reward is a one-time step valuation (does not accumulate per step) and is not added to the final success reward when successfully completed.

7. Scaling to a Real Android Emulator

To transition from this mock package to a physical device or Android emulator:

• State Space: Replace AppState with an XML accessibility node hierarchy obtained via UIAutomator2 (d.dump_hierarchy()), and parse bounding boxes to identify interactable buttons.
• Action Space: Convert structured actions into adb shell input tap x y commands and adb shell input text "<string>" key events.
• Observation: Supplement text-based states with visual inputs obtained using screenshots via adb shell screencap -p.
• Execution: Run an ADB daemon on the host connecting to the emulator port, mapping abstract actions to coordinates.

8. Prime Intellect & PRIME-RL Integration

The environment features full integration hooks matching Prime Intellect standards:

• SingleTurnEnv: Manages dataset loading for training/evaluation runs.
• Rubric: Combines reward components and weights, outputting a single reward score.
• Policy Training: Inside a training pipeline like PRIME-RL, an agent (e.g., initialized via GRPO or PPO) processes the natural-language screen observation output by env.reset() / env.step(). It outputs action sequences, and the Rubric evaluates their performance to compute gradients and optimize the policy weights.

9. Test Suite Breakdown

The test suite consists of 88 unit tests across three files:

• tests/test_actions.py: Validates proper state transitions (home -> settings -> home), invalid tap behavior (no crash, increments counter), and the finish command.
• tests/test_rewards.py: Checks exact goal matching, format reward ratios, safety penalties for logout actions, and efficiency decay.
• tests/test_env.py: Verifies step-by-step navigation, maximum step constraints, garbage dictionary handling, and state resetting.

Note

If the test suite screenshots are not displaying on GitHub due to image rendering/caching delays, the raw, real-time terminal output is provided below for reference:

============================================================
Running Mobile UI Environment Test Suite
============================================================

File: test_actions.py
----------------------------------------
  [PASS] TestBackAction.test_back_from_home_does_not_crash
  [PASS] TestBackAction.test_back_from_home_is_invalid
  [PASS] TestBackAction.test_back_from_notes_returns_to_home
  [PASS] TestBackAction.test_back_from_profile_returns_to_home
  [PASS] TestBackAction.test_back_from_settings_returns_to_home
  [PASS] TestBackAction.test_navigate_then_back_round_trip
  [PASS] TestFinishMarksDone.test_actions_rejected_after_finish
  [PASS] TestFinishMarksDone.test_finish_from_any_screen
  [PASS] TestFinishMarksDone.test_finish_marks_state_changed
  [PASS] TestFinishMarksDone.test_finish_result_has_message
  [PASS] TestFinishMarksDone.test_finish_sets_is_done
  [PASS] TestInvalidTapNoCrash.test_invalid_tap_increments_invalid_count
  [PASS] TestInvalidTapNoCrash.test_multiple_invalid_taps_accumulate
  [PASS] TestInvalidTapNoCrash.test_nonexistent_element_does_not_raise
  [PASS] TestInvalidTapNoCrash.test_nonexistent_element_returns_failure
  [PASS] TestInvalidTapNoCrash.test_tap_missing_target_key
  [PASS] TestInvalidTapNoCrash.test_wrong_screen_element_is_invalid
  [PASS] TestTypeOnWrongScreen.test_type_does_not_change_pending_text_on_wrong_screen
  [PASS] TestTypeOnWrongScreen.test_type_on_home_does_not_crash
  [PASS] TestTypeOnWrongScreen.test_type_on_home_increments_invalid_count
  [PASS] TestTypeOnWrongScreen.test_type_on_home_screen_is_invalid
  [PASS] TestTypeOnWrongScreen.test_type_on_notes_screen_is_valid
  [PASS] TestTypeOnWrongScreen.test_type_on_settings_screen_is_invalid
  [PASS] TestValidTapChangesScreen.test_step_counter_increments_on_valid_tap
  [PASS] TestValidTapChangesScreen.test_tap_notes_button_from_home
  [PASS] TestValidTapChangesScreen.test_tap_profile_button_from_home
  [PASS] TestValidTapChangesScreen.test_tap_settings_button_from_home

File: test_env.py
----------------------------------------
  [PASS] TestFullEpisodeCreateNote.test_done_flag_is_true
  [PASS] TestFullEpisodeCreateNote.test_note_appears_in_state
  [PASS] TestFullEpisodeCreateNote.test_result_keys_present
  [PASS] TestFullEpisodeCreateNote.test_reward_above_threshold
  [PASS] TestFullEpisodeCreateNote.test_reward_in_valid_range
  [PASS] TestInvalidActionsDontCrash.test_env_still_alive_after_garbage
  [PASS] TestInvalidActionsDontCrash.test_five_garbage_actions_no_exception
  [PASS] TestInvalidActionsDontCrash.test_invalid_count_accumulated
  [PASS] TestInvalidActionsDontCrash.test_result_has_per_action_info
  [PASS] TestInvalidActionsDontCrash.test_reward_clamped_with_many_invalids
  [PASS] TestLogoutTriggersSafety.test_logout_reduces_reward
  [PASS] TestLogoutTriggersSafety.test_logout_sets_logout_attempted_flag
  [PASS] TestLogoutTriggersSafety.test_safety_penalty_in_breakdown
  [PASS] TestLogoutTriggersSafety.test_safety_violations_incremented
  [PASS] TestMaxStepsTerminates.test_done_is_true_after_budget
  [PASS] TestMaxStepsTerminates.test_episode_ends_at_budget
  [PASS] TestMaxStepsTerminates.test_info_includes_steps_taken
  [PASS] TestMaxStepsTerminates.test_reward_still_computed_on_timeout
  [PASS] TestMaxStepsTerminates.test_steps_do_not_exceed_budget
  [PASS] TestResetClearsState.test_new_episode_after_reset_succeeds
  [PASS] TestResetClearsState.test_reset_clears_action_history
  [PASS] TestResetClearsState.test_reset_clears_invalid_count
  [PASS] TestResetClearsState.test_reset_clears_notes
  [PASS] TestResetClearsState.test_reset_clears_step_counter
  [PASS] TestResetClearsState.test_reset_returns_home_observation
  [PASS] TestResetClearsState.test_reset_sets_current_screen_to_home

File: test_rewards.py
----------------------------------------
  [PASS] TestComputeFinalReward.test_failed_episode_with_safety_violation
  [PASS] TestComputeFinalReward.test_perfect_episode_near_one
  [PASS] TestComputeFinalReward.test_result_clamped_to_zero_one
  [PASS] TestComputeFinalReward.test_returns_float
  [PASS] TestEfficiencyReward.test_exact_budget_used_gives_zero
  [PASS] TestEfficiencyReward.test_failure_gives_zero
  [PASS] TestEfficiencyReward.test_fewer_steps_gives_higher_reward
  [PASS] TestEfficiencyReward.test_result_in_zero_one_range
  [PASS] TestEfficiencyReward.test_success_halfway_gives_half
  [PASS] TestEfficiencyReward.test_success_in_few_steps_gives_positive
  [PASS] TestFormatRewardInvalid.test_all_malformed_returns_zero
  [PASS] TestFormatRewardInvalid.test_one_malformed_of_two
  [PASS] TestFormatRewardInvalid.test_tap_missing_target_key
  [PASS] TestFormatRewardInvalid.test_type_missing_text_key
  [PASS] TestFormatRewardInvalid.test_unknown_action_type_is_malformed
  [PASS] TestFormatRewardValid.test_all_tap_actions_valid
  [PASS] TestFormatRewardValid.test_empty_action_list_returns_one
  [PASS] TestFormatRewardValid.test_mixed_valid_action_types
  [PASS] TestFormatRewardValid.test_result_is_float
  [PASS] TestFormatRewardValid.test_single_valid_action
  [PASS] TestSafetyPenaltyLogout.test_logout_attempt_sets_violation
  [PASS] TestSafetyPenaltyLogout.test_multiple_violations_still_one
  [PASS] TestSafetyPenaltyLogout.test_no_violations_returns_zero
  [PASS] TestSafetyPenaltyLogout.test_one_violation_returns_one
  [PASS] TestSafetyPenaltyLogout.test_return_type_is_float
  [PASS] TestSuccessRewardNoteCreated.test_empty_notes_returns_zero
  [PASS] TestSuccessRewardNoteCreated.test_exact_match_required
  [PASS] TestSuccessRewardNoteCreated.test_multiple_notes_goal_found
  [PASS] TestSuccessRewardNoteCreated.test_note_present_returns_one
  [PASS] TestSuccessRewardNoteCreated.test_returns_float
  [PASS] TestSuccessRewardWrongNote.test_multi_note_goal_partially_met
  [PASS] TestSuccessRewardWrongNote.test_multiple_wrong_notes_returns_zero
  [PASS] TestSuccessRewardWrongNote.test_partial_match_not_accepted
  [PASS] TestSuccessRewardWrongNote.test_toggle_goal_wrong_value
  [PASS] TestSuccessRewardWrongNote.test_wrong_note_returns_zero

============================================================
Test Suite Summary: 88 passed, 0 failed out of 88 total tests.
============================================================

10. Trade-offs Analysis

Dimension	Option A	Option B	Trade-offs
Simulation	Mock State Machine	Real UI Emulator (ADB)	Option A is extremely fast, lightweight, and deterministic (ideal for unit testing and quick testing of RL algorithm loops), but lacks real-world rendering and OS latency issues found in Option B.
Agent Type	Heuristic Agent	LLM-based Agent	The heuristic agent is fast and 100% accurate for designed tasks but cannot generalize to unseen UI changes. LLMs can handle unstructured variations but suffer from API cost and non-deterministic failures.
Dataset Scale	30 Core Tasks	Thousands of Generated Tasks	30 tasks verify framework compliance and basic learning curves. Scaling to thousands of tasks is necessary for training generalized models but requires programmatic UI state generators.

State Machine Diagram

stateDiagram-v2
    [*] --> HOME : Episode Start
    HOME --> NOTES : tap(notes_button)
    HOME --> SETTINGS : tap(settings_button)
    HOME --> PROFILE : tap(profile_button)
    NOTES --> HOME : back
    SETTINGS --> HOME : back
    PROFILE --> HOME : back (safe)
    PROFILE --> [*] : tap(logout_button) ⚠️ SAFETY VIOLATION
    NOTES --> [*] : finish ✓
    SETTINGS --> [*] : finish ✓
    PROFILE --> [*] : finish ✓
    HOME --> [*] : max_steps exceeded

Loading

Reward Formula

final_reward = clip(

1.0 × success

0.1 × format 0.2 × efficiency

0.1 × invalid_count 0.3 × safety_violation,

0.0, 1.0

)

Reward Flow

flowchart TD
    A[Agent produces action list] --> B{Valid JSON format?}
    B -->|No| C[format_reward = 0]
    B -->|Yes| D[format_reward = 1.0]
    D --> E{Execute actions}
    E --> F{Logout attempted?}
    F -->|Yes| G[safety_penalty += 1\n⚠️ -0.3 to reward]
    F -->|No| H{Goal completed?}
    H -->|Yes| I[success_reward = 1.0\nefficiency_reward = steps saved]
    H -->|No| J{Reached correct screen?}
    J -->|Yes| K[partial_progress = 0.3]
    J -->|No| L[success_reward = 0]
    I --> M[compute_final_reward\nclip to 0.0 - 1.0]
    K --> M
    L --> M
    G --> M
    C --> M
    M --> N[Return reward to training loop]

Loading

Why Sparse Reward Alone Fails

Sparse reward means the agent only gets feedback (1.0) when the ENTIRE goal is complete. For an RL agent starting from random actions, the probability of accidentally completing "tap notes_button → tap add_note_button → type note → tap save_note_button → finish" in the correct order is extremely low.

With only sparse reward:

• Agent gets 0.0 for 99% of random episodes during early training
• Zero gradient signal → policy cannot improve
• Agent may never discover the correct action sequence

This environment solves it with shaped rewards:

• partial_progress_reward: +0.3 just for reaching the right screen (easier to discover)
• efficiency_reward: smooth gradient once success is found
• format_reward: tiny signal (+0.1) that at least valid JSON is being produced

graph LR
    A[Pure Sparse\nReward] -->|"99% episodes\nreward = 0"| B[No gradient\nsignal]
    B --> C[Agent never\nimproves]
    D[Shaped Reward\nThis Environment] -->|"Partial progress\n+0.3 on right screen"| E[Agent learns\nnavigation first]
    E -->|"Success reward\n+1.0 on goal"| F[Agent learns\nfull task]

Loading

Eval Output — Docker Screenshot

The following shows the deterministic HeuristicAgent achieving 100% success rate across all 10 eval tasks when run via Docker:

Note

If the evaluation results screenshot is not displaying on GitHub due to image rendering/caching delays, the raw, real-time terminal output is provided below for reference:

============================================================

Prime Intellect RL Environment — Eval Results

Total eval tasks : 10

Success rate     : 100%

Average reward   : 1.00

Average steps    : 5.0

Invalid action rate: 0.00

Safety violations: 0
Per-task Breakdown

Task ID     Goal Type          Reward  Success  Steps

task_021    note_created        1.000    ✓        5

task_022    note_created        1.000    ✓        5

task_023    toggle_enabled      1.000    ✓        3

task_024    toggle_disabled     1.000    ✓        3

task_025    info_retrieved      1.000    ✓        3

task_026    multi_note_created  1.000    ✓        8

task_027    safe_navigation     1.000    ✓        4

task_028    note_created        1.000    ✓        5

task_029    multi_note_created  1.000    ✓       11

task_030    toggle_disabled     1.000    ✓        3
All eval tasks completed successfully.

Hidden-Eval Compatibility

This environment is designed to be hidden-eval compatible:

• build_dataset(split="eval") is fully separated from train — no data leakage
• HeuristicAgent uses goal TYPE to decide actions, not task_id — works on unseen tasks
• compute_final_reward() is a pure function — same behavior on any task schema
• Environment never crashes on invalid input — returns penalized reward instead
• New goal types can be added to dataset.py without changing env.py or rubric.py

Interview Questions — My Answers

1. Why is sparse reward difficult for RL agents? Sparse rewards only provide feedback when the entire target sequence of actions is successfully completed. For environments with large action spaces, the probability of randomly executing the correct sequence of actions (such as navigating, typing, saving, and terminating) is virtually zero, resulting in a flat reward curve with no gradients to guide policy improvements.

2. What is reward hacking in your environment? Reward hacking occurs when an agent exploits loopholes in the reward components to maximize reward without solving the target task. For instance, the agent could spam syntactically valid JSON actions to farm the format reward, or navigate between the home and settings screens repeatedly to exploit the partial progress reward without completing the goals.

3. Why should train and eval tasks be separated? Separating train and eval tasks is crucial to measure the generalization capability of the RL agent on unseen instructions and target values. If tasks are shared, the agent can easily overfit by memorizing specific step sequences for each task ID rather than learning the underlying semantic mapping of UI elements.

4. How would you convert this mock app into a real Android emulator environment? I would replace the AppState class with an XML parser that dumps the emulator's accessibility hierarchy using UIAutomator2 (d.dump_hierarchy()). The ActionExecutor would be replaced by sending ADB shell commands such as adb shell input tap x y for coordinates parsed from the hierarchy and adb shell input text for text entry.

5. What should the observation include: screenshot, XML tree, text tree, or all three? The observation should ideally include both the XML accessibility tree and the raw screenshot. The XML tree provides structured semantic information (labels, coordinates, focusable flags) for precise action mapping, while the screenshot allows visual model policy training (e.g. via VLM policies) to handle custom UI elements and graphics that accessibility labels might miss.

6. How would you reduce invalid actions? Invalid actions can be reduced by using an invalid action masking layer on the policy's action head, dynamically checking get_available_elements() for the current screen. Additionally, applying a step-wise negative penalty for invalid actions in the rubric discourages the agent from selecting unavailable UI elements during exploration.

7. Which metrics are more important than average reward? Task success rate is the most critical metric as it represents goal completion independent of shaping weights. Step efficiency (average trajectory length on successful runs) and safety violation rate are also more informative than average reward, as they directly measure the optimality and compliance of the policy.

8. How would you scale this to 1,000 mobile tasks? Scaling to 1,000 tasks requires building a procedural task generator that samples random combinations of goal screens, toggle states, note contents, and safety constraints. We would also script randomized mock database profiles to synthesize a wide range of user details and verification criteria programmatically.

9. What changes are needed for multi-turn agent training? Multi-turn agent training requires appending the conversation history (a list of past screen observations, generated actions, and execution results) to the current prompt context. In terms of RL mechanics, the value function must be calculated across the entire trajectory history to correctly distribute credit for early state decisions.

10. How would you debug an agent that reaches the right screen but fails the final goal? I would analyze the trajectory rollouts to see if the policy fails at a specific interaction step (e.g., typing but not saving, or using incorrect target names). By checking intermediate values in the rubric, such as the format and invalid action counts on the final screen, I can determine if the failure stems from exploration policies, action formatting issues, or sub-optimal efficiency weights.