o 1 i@sdZddlZddlZddlZddlmZmZddlm Z ddl m Z e e Ze dejfddZe dd ed ed efd dZe ddZdS)z7Utils for minibatch SGD across multiple RLlib policies.N)MultiAgentBatch SampleBatch) OldAPIStack)LearnerInfoBuilderarraycCs||td|S)zNormalize the values in an array. Args: array (np.ndarray): Array of values to normalize. Returns: array with zero mean and unit standard deviation. g-C6?)meanmaxZstd)rr _/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/ray/rllib/utils/sgd.py standardizeds r Tsamplessgd_minibatch_sizeshufflec cs|s|VdSt|trtdd|vrd|vr|||}|\}}t|dkrC|r3t||D] \}}|||Vq5dStt||}|rQt||D]\\}}\}} | |||| VqSdS)a Return a generator yielding minibatches from a sample batch. Args: samples: SampleBatch to split up. sgd_minibatch_size: Size of minibatches to return. shuffle: Whether to shuffle the order of the generated minibatches. Note that in case of a non-recurrent policy, the incoming batch is globally shuffled first regardless of this setting, before the minibatches are generated from it! Yields: SampleBatch: Each of size `sgd_minibatch_size`. Nz;Minibatching not implemented for multi-agent in simple modeZ state_in_0Z state_out_0r) isinstancerNotImplementedErrorrZ_get_slice_indiceslenrandomlistzipslice) r r rZ all_slicesZ data_slicesZ state_slicesijsiZsjr r r minibatchess0      rcCs|}tdd}|D]X\}}||jvrq |j|} |D] } t| | | | <q|rC|jdd|krCtd||jddt |D]} t | |D]} | t || i| j |} || |qNqGq |}|S)a Execute minibatch SGD. Args: samples: Batch of samples to optimize. policies: Dictionary of policies to optimize. local_worker: Master rollout worker instance. num_sgd_iter: Number of epochs of optimization to take. sgd_minibatch_size: Size of minibatches to use for optimization. standardize_fields: List of sample field names that should be normalized prior to optimization. Returns: averaged info fetches over the last SGD epoch taken. )Z num_devicesmodelZ max_seq_lenzC`sgd_minibatch_size` ({}) cannot be smaller than`max_seq_len` ({}).)Zas_multi_agentritemsZpolicy_batchesr Z is_recurrentconfig ValueErrorformatrangerZlearn_on_batchrcountZadd_learn_on_batch_resultsfinalize)r ZpoliciesZ local_workerZ num_sgd_iterr Zstandardize_fieldsZlearner_info_builderZ policy_idpolicybatchfieldrZ minibatchresultsZ learner_infor r r do_minibatch_sgdHs:     r')T)__doc__loggingrnumpynpZray.rllib.policy.sample_batchrrZray.rllib.utils.annotationsrZ$ray.rllib.utils.metrics.learner_infor getLogger__name__loggerZndarrayr intboolrr'r r r r s    +