#!/usr/bin/env python3 # -*- coding: UTF-8 -*- ########################################################################### # Copyright © 1998 - 2026 Tencent. All Rights Reserved. ########################################################################### """ Author: Tencent AI Arena Authors Training workflow for Robot Vacuum. 清扫大作战训练工作流。 """ import os import time import numpy as np from agent_ppo.conf.conf import Config from agent_ppo.feature.definition import SampleData, sample_process from tools.metrics_utils import get_training_metrics from tools.train_env_conf_validate import read_usr_conf from common_python.utils.workflow_disaster_recovery import handle_disaster_recovery def workflow(envs, agents, logger=None, monitor=None, *args, **kwargs): last_save_model_time = time.time() env = envs[0] agent = agents[0] diag_max_episodes = _read_diag_max_episodes(logger) diag_log_only = _read_bool_env("ROBOT_VACUUM_DIAG_LOG_ONLY") # Read and validate user configuration # 读取和校验用户配置 usr_conf = read_usr_conf("agent_ppo/conf/train_env_conf.toml", logger) if usr_conf is None: logger.error("usr_conf is None, please check agent_ppo/conf/train_env_conf.toml") return _apply_diag_env_overrides(usr_conf, logger) episode_runner = EpisodeRunner( env=env, agent=agent, usr_conf=usr_conf, logger=logger, monitor=monitor, diag_max_episodes=diag_max_episodes, diag_log_only=diag_log_only, ) while not episode_runner.stop_requested: for g_data in episode_runner.run_episodes(): agent.send_sample_data(g_data) g_data.clear() now = time.time() if now - last_save_model_time >= 1800: agent.save_model() last_save_model_time = now if episode_runner.stop_requested: break if episode_runner.stop_requested: logger.info(f"diagnostic max episodes reached: {episode_runner.episode_cnt}") def _read_diag_max_episodes(logger): raw_value = os.environ.get("ROBOT_VACUUM_DIAG_MAX_EPISODES", "").strip() if not raw_value: return 0 try: value = int(raw_value) except ValueError: if logger: logger.warning(f"ignore invalid ROBOT_VACUUM_DIAG_MAX_EPISODES={raw_value!r}") return 0 return max(value, 0) def _read_positive_int_env(name, logger): raw_value = os.environ.get(name, "").strip() if not raw_value: return 0 try: value = int(raw_value) except ValueError: if logger: logger.warning(f"ignore invalid {name}={raw_value!r}") return 0 return max(value, 0) def _read_bool_env(name): return os.environ.get(name, "").strip().lower() in ("1", "true", "yes", "on") def _apply_diag_env_overrides(usr_conf, logger): diag_max_step = _read_positive_int_env("ROBOT_VACUUM_DIAG_MAX_STEP", logger) if diag_max_step <= 0: return env_conf = usr_conf.setdefault("env_conf", {}) old_max_step = env_conf.get("max_step") env_conf["max_step"] = diag_max_step if logger: logger.info(f"diagnostic max_step override: {old_max_step} -> {diag_max_step}") class EpisodeRunner: def __init__(self, env, agent, usr_conf, logger, monitor, diag_max_episodes=0, diag_log_only=False): self.env = env self.agent = agent self.usr_conf = usr_conf self.logger = logger self.monitor = monitor self.episode_cnt = 0 self.last_report_monitor_time = 0 self.last_get_training_metrics_time = 0 self.diag_max_episodes = int(diag_max_episodes) self.diag_log_only = bool(diag_log_only) self.stop_requested = False def run_episodes(self): """Run a single episode and yield collected samples. 单局流程（generator），完成一局后 yield 整局样本。 """ while True: if self.stop_requested: return # Periodically get training metrics # 定期打印训练指标 now = time.time() if now - self.last_get_training_metrics_time >= 60: training_metrics = get_training_metrics() self.last_get_training_metrics_time = now if training_metrics is not None: self.logger.info(f"training_metrics: {training_metrics}") # Reset environment # 重置环境 env_obs = self.env.reset(self.usr_conf) if handle_disaster_recovery(env_obs, self.logger): continue # Reset agent and load latest model # 重置 Agent，加载最新模型 self.agent.reset(env_obs) self.agent.load_model(id="latest") # Initial observation processing # 初始观测 obs_data, remain_info = self.agent.observation_process(env_obs) collector = [] self.episode_cnt += 1 done = False step = 0 total_reward = 0.0 self.logger.info(f"Episode {self.episode_cnt} start") while not done: # Agent inference / 推理动作 act_data_list = self.agent.predict([obs_data]) act_data = act_data_list[0] act = self.agent.action_process(act_data) # Environment step / 与环境交互 env_reward, env_obs = self.env.step(act) if handle_disaster_recovery(env_obs, self.logger): break terminated = env_obs["terminated"] truncated = env_obs["truncated"] frame_no = env_obs["frame_no"] step += 1 done = terminated or truncated # Process next observation # 特征处理 _obs_data, _ = self.agent.observation_process(env_obs) _obs_data.frame_no = frame_no reward_scalar = float(self.agent.last_reward) total_reward += reward_scalar # Terminal reward calculation # 终局奖励 final_reward = 0.0 if done: fm = self.agent.preprocessor observation = env_obs.get("observation") or {} observation = observation if isinstance(observation, dict) else {} env_info = observation.get("env_info") or {} env_info = env_info if isinstance(env_info, dict) else {} extra_info = env_obs.get("extra_info") or {} extra_info = extra_info if isinstance(extra_info, dict) else {} total_score = env_info.get("total_score", fm.total_score) remaining_charge = env_info.get("remaining_charge", fm.remaining_charge) charge_count = env_info.get("charge_count", fm.charge_count) finished_steps = env_info.get("finished_steps", step) result_message = extra_info.get("result_message", "") result_code = extra_info.get("result_code", "") cleaning_ratio = fm.dirt_cleaned / max(fm.total_dirt, 1) score_per_step = total_score / max(finished_steps, 1) if truncated: if score_per_step < 0.25: final_reward = -3.0 + 6.0 * cleaning_ratio result_str = "STALL_TRUNCATED" else: # Survived to max steps: higher cleaning ratio → more reward # 存活到最大步数：清扫比例越高奖励越多 final_reward = 2.0 + 8.0 * cleaning_ratio result_str = "WIN" else: if fm.battery <= 0 or remaining_charge <= 0: final_reward = -fm.battery_fail_penalty() + 4.0 * cleaning_ratio result_str = "BATTERY_FAIL" elif fm.npc_danger or fm.npc_predicted_danger or fm.nearest_npc_dist <= 1: final_reward = -3.0 + 6.0 * cleaning_ratio result_str = "NPC_FAIL" else: final_reward = -2.0 + 6.0 * cleaning_ratio result_str = "FAIL" self.logger.info( f"[GAMEOVER] ep:{self.episode_cnt} steps:{step} " f"finished_steps:{finished_steps} " f"result:{result_str} final_bonus:{final_reward:.2f} " f"total_reward:{total_reward:.3f} " f"dirt_cleaned:{fm.dirt_cleaned}/{fm.total_dirt} " f"total_score:{total_score} " f"remaining_charge:{remaining_charge} " f"charge_count:{charge_count} " f"recharge_steps:{fm.recharge_steps} " f"stuck_count:{fm.stuck_count} " f"max_stuck_steps:{fm.max_stuck_steps} " f"recharge_escape_count:{fm.recharge_escape_count} " f"npc_close_steps:{fm.npc_close_steps} " f"npc_danger_steps:{fm.npc_danger_steps} " f"npc_collision:{fm.npc_collision} " f"nearest_charger:{fm.nearest_charger_range_dist:.1f} " f"nearest_npc:{fm.nearest_npc_dist:.1f} " f"result_code:{result_code} " f"result_message:{result_message}" ) diag = fm.get_diagnostic_summary() self.logger.info( f"[DIAG] ep:{self.episode_cnt} map:{diag['map_id']} " f"steps:{step} result:{result_str} " f"profile:{diag['reward_profile']} route:{diag['charger_route_source']} " f"score:{float(total_score):.1f} reward:{total_reward + final_reward:.3f} " f"mask_avg(raw/block/npc/recharge/escape/leave/final):" f"{diag['avg_mask_counts']['raw']:.2f}/" f"{diag['avg_mask_counts']['blocked']:.2f}/" f"{diag['avg_mask_counts']['npc']:.2f}/" f"{diag['avg_mask_counts']['recharge']:.2f}/" f"{diag['avg_mask_counts']['escape']:.2f}/" f"{diag['avg_mask_counts']['leave']:.2f}/" f"{diag['avg_mask_counts']['final']:.2f} " f"mask_changed(block/npc/recharge/escape/leave):" f"{diag['mask_changed_steps']['blocked']}/" f"{diag['mask_changed_steps']['npc']}/" f"{diag['mask_changed_steps']['recharge']}/" f"{diag['mask_changed_steps']['escape']}/" f"{diag['mask_changed_steps']['leave']} " f"mask_active(recharge/leave):" f"{diag['mask_active_steps']['recharge']}/" f"{diag['mask_active_steps']['leave']} " f"tight(one/<=2/zero):" f"{diag['one_action_steps']}/" f"{diag['two_or_less_action_steps']}/" f"{diag['zero_final_steps']} " f"actions:{diag['action_hist']} " f"known:{diag['known_ratio']:.3f} dirty_known:{diag['known_dirty_ratio']:.3f} " f"frontier:{diag['frontier_ratio']:.3f} " f"path_dirty/frontier:{diag['global_dirty_path_dist']:.1f}/" f"{diag['frontier_path_dist']:.1f}" ) # Build sample frame # 构造样本帧 reward_arr = np.array([reward_scalar], dtype=np.float32) value_arr = act_data.value.flatten()[: Config.VALUE_NUM] frame = SampleData( obs=np.array(obs_data.feature, dtype=np.float32), legal_action=np.array(obs_data.legal_action, dtype=np.float32), act=np.array(act_data.action), reward=reward_arr, done=np.array([float(done)]), reward_sum=np.zeros(Config.VALUE_NUM, dtype=np.float32), value=value_arr, next_value=np.zeros(Config.VALUE_NUM, dtype=np.float32), advantage=np.zeros(Config.VALUE_NUM, dtype=np.float32), prob=np.array(act_data.prob, dtype=np.float32), ) collector.append(frame) if done: # Add terminal reward to last frame # 终局奖励叠加到最后一步 collector[-1].reward = collector[-1].reward + np.array([final_reward], dtype=np.float32) if truncated and not terminated: collector[-1].next_value = self.agent.estimate_value(_obs_data) collector[-1].done = np.array([0.0], dtype=np.float32) # Monitor reporting / 监控上报 now = time.time() if now - self.last_report_monitor_time >= 60 and self.monitor: self.monitor.put_data( { os.getpid(): { "reward": total_reward + final_reward, "episode_cnt": self.episode_cnt, "total_score": float(total_score), "stuck_count": float(fm.stuck_count), "max_stuck_steps": float(fm.max_stuck_steps), "recharge_escape_count": float(fm.recharge_escape_count), "npc_close_steps": float(fm.npc_close_steps), "npc_danger_steps": float(fm.npc_danger_steps), "npc_collision": float(fm.npc_collision), "battery_fail": float(fm.battery_fail), "charge_count": float(charge_count), "remaining_charge": float(remaining_charge), "recharge_steps": float(fm.recharge_steps), "mask_final_avg": float(diag["avg_mask_counts"]["final"]), "mask_recharge_active": float(diag["mask_active_steps"]["recharge"]), "mask_recharge_changed": float(diag["mask_changed_steps"]["recharge"]), "mask_one_action_steps": float(diag["one_action_steps"]), "mask_two_or_less_action_steps": float(diag["two_or_less_action_steps"]), "mask_zero_final_steps": float(diag["zero_final_steps"]), } } ) self.last_report_monitor_time = now # Compute GAE and yield samples # GAE 计算并 yield 样本 if collector: if self.diag_max_episodes > 0 and self.episode_cnt >= self.diag_max_episodes: self.stop_requested = True if self.diag_log_only: collector.clear() if self.stop_requested: return break collector = sample_process(collector) yield collector break # Advance state / 状态推进 obs_data = _obs_data