o cij@sddlZddlZddlmZddlZddlZddlm Z ddl m Z ddl m Z ddlmZmZmZddlmZmZddlmZmZdd lmZdd lmZmZmZmZmZe\Z Z!Z"eZ#eGd d d e Z$eGd dde$Z%ede&fddZ'eGddde$Z(eGddde%Z)eGddde$Z*eGddde$Z+eGddde$Z,eGddde$Z-eGddde$Z.eGd d!d!e$Z/eGd"d#d#e$Z0dS)$N)log)Optional)ActionDistribution)ModelV2)MIN_LOG_NN_OUTPUTMAX_LOG_NN_OUTPUT SMALL_NUMBER) OldAPIStackoverride) try_import_tftry_import_tfp)get_base_struct_from_space) TensorTypeListUnionTupleModelConfigDictcspeZdZdZeedeedeffdd Z defddZ eedefd d Z eedefd d Z Z S) TFActionDistributionz9TF-specific extensions for building action distributions.inputsmodelcs*t||||_||j|_dSN)super__init___build_sample_op sample_oplogpsampled_action_logp_op)selfrr __class__V/home/ubuntu/.local/lib/python3.10/site-packages/ray/rllib/models/tf/tf_action_dist.pyrs zTFActionDistribution.__init__returncCst)zImplement this instead of sample(), to enable op reuse. This is needed since the sample op is non-deterministic and is shared between sample() and sampled_action_logp(). )NotImplementedErrorrr r r!rsz%TFActionDistribution._build_sample_opcC|jS)z+Draw a sample from the action distribution.)rr$r r r!sample&zTFActionDistribution.samplecCr%)z2Returns the log probability of the sampled action.)rr$r r r!sampled_action_logp+r'z(TFActionDistribution.sampled_action_logp)__name__ __module__ __qualname____doc__r rrrrrrr&r( __classcell__r r rr!rsrcseZdZdZ ddeededeffdd Ze e d efd d Z e e d ed efd dZ e e d efddZ e e de d efddZe ed efddZee e ddZZS) Categoricalz4Categorical distribution for discrete action spaces.N?rr temperaturecs&|dksJdt|||dS)N(Categorical `temperature` must be > 0.0!rrrrrr0rr r!r5szCategorical.__init__r"cCstjj|jddSNaxis)tfmathargmaxrr$r r r!deterministic_sample=z Categorical.deterministic_samplexcCstjj|jt|tjd S)N)logitslabels)r9nn(sparse_softmax_cross_entropy_with_logitsrcastint32rr>r r r!rAszCategorical.logpcCsX|jtj|jddd}t|}tj|ddd}||}tj|tj||ddSNr6Tr8keepdimsr7rr9 reduce_maxexp reduce_sumr:r)ra0ea0z0p0r r r!entropyGs  zCategorical.entropyotherc Cs|jtj|jddd}|jtj|jddd}t|}t|}tj|ddd}tj|ddd}||}tj||tj||tj|ddSrFrI) rrRrMa1rNea1rOz1rPr r r!klOs  .zCategorical.klcCstjtj|jdddSr5)r9squeezerandom categoricalrr$r r r!rZzCategorical._build_sample_opcCr%rn action_space model_configr r r!required_model_output_shape^r'z'Categorical.required_model_output_shapeNr/)r)r*r+r,rrrfloatrr rr<rrQrVrr staticmethodr`r-r r rr!r.1s. r.tcsGfdddt}|S)zGCategorical distribution class that has customized default temperature.cs"eZdZdffdd ZZS)zJget_categorical_class_with_temperature..CategoricalWithTemperatureNcst|||dSrr3r4rr r!riszSget_categorical_class_with_temperature..CategoricalWithTemperature.__init__)r)r*r+rr-r rdrr!CategoricalWithTemperaturehsrfr.)rdrfr rer!&get_categorical_class_with_temperaturedsrhc @seZdZdZ ddeededeeee j e edfffddZ e ed efd d Ze ed ed efd dZe ed efddZe ed efddZe eded efddZe eded efddZe ed efddZee edejded eee j ffddZdS)MultiCategoricalz>MultiCategorical distribution for MultiDiscrete action spaces.Nrr input_lens.cstt||fddtj||ddD|_||_|jdur,tjdd|jD|_| |_ | |j |_ dS)Ncsg|]}t|qSr rg).0input_rr r! |sz-MultiCategorical.__init__..r6r7cSsg|]}|jjdqS)r6)rshape)rkcr r r!rns) rrr9splitcatsr^gymspaces MultiDiscreterrrr)rrrrjr^r rmr!rss   zMultiCategorical.__init__r"cCsTtjdd|jDdd}t|jtjjr(tt |dgt |jj |jj S|S)NcSg|]}|qSr r<rkcatr r r!rnz9MultiCategorical.deterministic_sample..r6r7 r9stackrr isinstancer^rsrtBoxrCreshapelistrodtype)rsample_r r r!r<sz%MultiCategorical.deterministic_sampleactionscCst|tjr.r)r~r9Tensorr^rsrtrrintnpprodroru set_shapelenrrunstackrCrDr}ziprL)rrlogpsr r r!rs zMultiCategorical.logpcCstjdd|jDddS)NcSrvr rQrxr r r!rnrzz2MultiCategorical.multi_entropy..r6r7)r9r}rrr$r r r! multi_entropyrZzMultiCategorical.multi_entropycCstj|ddSr5)r9rLrr$r r r!rQr=zMultiCategorical.entropyrRcCs"tjddt|j|jDddS)NcSrr rV)rkryoth_catr r r!rnrz-MultiCategorical.multi_kl..r6r7)r9r}rrrrrRr r r!multi_klszMultiCategorical.multi_klcCstj||ddSr5)r9rLrrr r r!rVszMultiCategorical.klcCsVtjdd|jDdd}t|jtjjr)tjt |dgt |jj |jj dS|S)NcSrvr r&rxr r r!rnrzz5MultiCategorical._build_sample_op..r6r7r{rr|)rrr r r!rsz!MultiCategorical._build_sample_opr^r_cCst|tjjr?|jjdsJt|j }t |j }t |j |ks&Jt |j |ks0Jtj |jtjd||dSt|jS)Nrrr6)r~rsrtrrname startswithrminlowmaxhighallrrorDsumnvec)r^r_low_high_r r r!r`s   z,MultiCategorical.required_model_output_shaper)r)r*r+r,rrrrrrndarrayrrr rr<rrrQrrVrrrcrsSpacerr`r r r r!riosB     ric s|eZdZdZ    ddeedededee j j ffdd Z e ed effd d Ze ed ed efddZZS)SlateMultiCategoricalafMultiCategorical distribution for MultiDiscrete action spaces. The action space must be uniform, meaning all nvec items have the same size, e.g. MultiDiscrete([10, 10, 10]), where 10 is the number of candidates to pick from and 3 is the slate size (pick 3 out of 10). When picking candidates, no candidate must be picked more than once. Nr/rrr0r^cs`|dksJdt||||_tjtjjr)tfddjjDs+J|_ dS)Nr1r2c3s |] }|jjdkVqdS)rN)r^r)rkr\r$r r! s z1SlateMultiCategorical.__init__..) rrr^r~rsrtrurr all_slates)rrrr0r^rrr$r!rs zSlateMultiCategorical.__init__r"cst}t|j|Sr)rr<r9gatherrrr&rr r!r<s z*SlateMultiCategorical.deterministic_sampler>cCst|jdddfS)Nr)r9 ones_likerrEr r r!rszSlateMultiCategorical.logp)Nr/NN)r)r*r+r,rrrrbrrsrtrurr rr<rr-r r rr!rs&  rc seZdZdZ ddeededeffdd Ze e d efd d Z e e d ed efd dZ e e d efddZee e dejded eeejffddZZS) GumbelSoftmaxaGumbelSoftmax distr. (for differentiable sampling in discr. actions The Gumbel Softmax distribution [1] (also known as the Concrete [2] distribution) is a close cousin of the relaxed one-hot categorical distribution, whose tfp implementation we will use here plus adjusted `sample_...` and `log_prob` methods. See discussion at [0]. [0] https://stackoverflow.com/questions/56226133/ soft-actor-critic-with-discrete-action-space [1] Categorical Reparametrization with Gumbel-Softmax (Jang et al, 2017): https://arxiv.org/abs/1611.01144 [2] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables (Maddison et al, 2017) https://arxiv.org/abs/1611.00712 Nr/rrr0csD|dksJtjj||d|_tj|jjj|_ t ||dS)aAInitializes a GumbelSoftmax distribution. Args: temperature: Temperature parameter. For low temperatures, the expected value approaches a categorical random variable. For high temperatures, the expected value approaches a uniform distribution. r1)r0r?N) tfp distributionsRelaxedOneHotCategoricaldistr9rAsoftmax _distributionr?probsrrr4rr r!r s zGumbelSoftmax.__init__r"cCr%r)rr$r r r!r<r'z"GumbelSoftmax.deterministic_sampler>cCsz|j|jjjkr*tj||jjjdtjd}|j|jjjks*J|j|jjjftj| tjj |jjdddd S)Nr{rr7) rorr?r9one_hotas_listfloat32rLrA log_softmax)rr>valuesr r r!r#szGumbelSoftmax.logpcC |jSrrr&r$r r r!r6 zGumbelSoftmax._build_sample_opr^r_cCr%rr[r]r r r!r`:sz)GumbelSoftmax.required_model_output_shapera)r)r*r+r,rrrrbrr rr<rrrrcrsrrrrrrr`r-r r rr!rs2 rc seZdZdZdddeededeej j ffddZ e e d efd d Ze e d ed efd dZe e de d efddZe e d efddZe ed efddZee e dej ded eeejffddZZS) DiagGaussianzAction distribution where each vector element is a gaussian. The first half of the input vector defines the gaussian means, and the second half the gaussian standard deviations. N)r^rrr^csNtj|ddd\}}||_||_t||_|o|jdk|_t ||dS)Nr6r7r ) r9rqmeanlog_stdrKstdrozero_action_dimrr)rrrr^rrrr r!rJs  zDiagGaussian.__init__r"cCr%r)rr$r r r!r<Yz!DiagGaussian.deterministic_sampler>cCstt|jddkrtj|dd}dtjtjt|tj|j |j dddt dt j tt|dtjtj|jddS)Nrr6r7g?@)rr9ro expand_dimsrLr:squarerCrrrrrpirrEr r r!r]s"*zDiagGaussian.logprRcCs\t|tsJtj|j|jtj|jtj|j|jdtj|jdddS)Nrrr6r7) r~rr9rLrr:rrrrr r r!rVks zDiagGaussian.klcCs*tj|jdtdtjtjddS)Nrrr6r7)r9rLrrrrer$r r r!rQws*zDiagGaussian.entropycCs8|j|jtjt|j}|jrtj|ddS|SNr{r7)rrr9rXnormalrorrWrr r r!r{s zDiagGaussian._build_sample_opr_cCtj|jtjddSNrrrrrorDr]r r r!r`z(DiagGaussian.required_model_output_shape)r)r*r+r,rrrrrsrtrrr rr<rrVrQrrrcrrrrrr`r-r r rr!rBs:     rc s eZdZdZ  d#deedededeffdd Ze e d efd d Z e e d efd dZ e e ded efddZddZe e d efddZe e de d efddZded efddZded efddZee e dejd ed eeejffd!d"ZZS)$SquashedGaussianzA tanh-squashed Gaussian distribution defined by: mean, std, low, high. The distribution will never return low or high exactly, but `low`+SMALL_NUMBER or `high`-SMALL_NUMBER respectively. r/rrrrcs~tdusJtj|ddd\}}t|tt}t|}tjj||d|_ t t ||s0J||_ ||_t||dS)zParameterizes the distribution via `inputs`. Args: low: The lowest possible sampling value (excluding this value). high: The highest possible sampling value (excluding this value). Nrr{r7)locscale)rr9rq clip_by_valuerrrKrNormaldistrrrlessrrrr)rrrrrrrrrr r!rs  zSquashedGaussian.__init__r"cC|j}||Sr)rr_squashrrr r r!r<  z%SquashedGaussian.deterministic_samplecC||jSr)rrr&r$r r r!rz!SquashedGaussian._build_sample_opr>cCstt|||jj}|j|}t|dd}tj|dd}tj |}|tjtj d|dt dd}|S)Nidr{r7r6r) r9rC _unsquashrrrlog_probrrLr:tanhrr)rr>unsquashed_valueslog_prob_gaussianunsquashed_values_tanhdrr r r!rs  zSquashedGaussian.logpcCsF|j}||}|tj|j|tjd||tddfSNr6r{r7) rr&rr9rLrr:rr)rzrr r r! sample_logps  "zSquashedGaussian.sample_logpcCtd)Nz)Entropy not defined for SquashedGaussian! ValueErrorr$r r r!rQzSquashedGaussian.entropyrRcCr)Nz$KL not defined for SquashedGaussian!rrr r r!rVrzSquashedGaussian.kl raw_valuescCs8tj|dd|j|j|j}t||j|jS)Nr/r)r9r:rrrr)rrsquashedr r r!rs  zSquashedGaussian._squashrcCsD||j|j|jdd}t|dtdt}tj|}|S)Nrr/r)rrr9rrr:atanh)rr normed_valuessave_normed_values unsquashedr r r!rs  zSquashedGaussian._unsquashr^r_cCrrrr]r r r!r`rz,SquashedGaussian.required_model_output_shape)rr/)r)r*r+r,rrrrbrr rr<rrrrrQrVrrrcrsrrrrrrr`r-r r rr!rsF   rc seZdZdZ  ddeedededeffdd Ze e d efd d Z e e d efd dZ e e ded efddZded efddZded efddZee e dejded eeejffddZZS)BetaaB A Beta distribution is defined on the interval [0, 1] and parameterized by shape parameters alpha and beta (also called concentration parameters). PDF(x; alpha, beta) = x**(alpha - 1) (1 - x)**(beta - 1) / Z with Z = Gamma(alpha) Gamma(beta) / Gamma(alpha + beta) and Gamma(n) = (n - 1)! r1r/rrrrcsxt|tttt }tjtj|dd}||_||_tj|ddd\}}t j j ||d|_ t ||dS)Nr/rr{r7)concentration1concentration0)r9rrrr:rKrrrqrrrrrr)rrrrralphabetarr r!rsz Beta.__init__r"cCrr)rrrrr r r!r< rzBeta.deterministic_samplecCrr)rrr&r$r r r!rrzBeta._build_sample_opr>cCs"||}tjj|j|ddSr)rr9r:rLrr)rr>rr r r!rs z Beta.logprcCs||j|j|jSr)rr)rrr r r!rz Beta._squashrcCs||j|j|jSr)rr)rrr r r!rrzBeta._unsquashr^r_cCrrrr]r r r!r`rz Beta.required_model_output_shape)r1r/)r)r*r+r,rrrrbrr rr<rrrrrrcrsrrrrrrr`r-r r rr!rs<  rc @seZdZdZeedefddZeededefddZ eedefdd Z e eed e j d edeeejffd d ZdS) DeterministiczAction distribution that returns the input values directly. This is similar to DiagGaussian with standard deviation zero (thus only requiring the "mean" values as NN output). r"cCr%rrr$r r r!r<-rz"Deterministic.deterministic_sampler>cCs t|jSr)r9 zeros_likerrEr r r!r1s zDeterministic.logpcCr%rrr$r r r!r5rzDeterministic._build_sample_opr^r_cCtj|jtjdSNrrr]r r r!r`9z)Deterministic.required_model_output_shapeN)r)r*r+r,r rrr<rrrrcrsrrrrrrr`r r r r!r%s" rc@seZdZdZddZeeddZeeddZeedd Z eed d Z eed d Z ee ddZ eeddZdS)MultiActionDistributionzAction distribution that operates on a set of actions. Args: inputs (Tensor list): A list of tensors from which to compute samples. c sZt||t||_tj|tjd|_tj ||jdd}t fdd|||_ dS)Nrr6r7cs||fiSrr )rrlkwargsrr r!Ssz2MultiActionDistribution.__init__..) rrr action_space_structrarrayrDrjr9rqtree map_structureflat_child_distributions)rrrchild_distributionsrjr^r split_inputsr rr!rIs   z MultiActionDistribution.__init__cCst|tjtjfrUg}|jD]=}t|tr|dqt|tr0|j dur0|t |j j q| }t |j dkrA|dq|t |dqtj||dd}nt|}dd}t|||j}tdd|S)Nr6r7cSs>t|trtt|jdkrtj|ddn|tj}||Sr) r~r.r9rCrrorWrDr)valrr r r!map_ps " z*MultiActionDistribution.logp..map_cS||Srr abr r r!r|z.MultiActionDistribution.logp..)r~r9rrrrr.appendrir^rror&rrqrflattenr functoolsreduce)rr> split_indicesrr&split_xr  flat_logpsr r r!rXs$       zMultiActionDistribution.logpcCs(ddt|j|jD}tdd|S)NcSrr r)rkdor r r!rnsz.MultiActionDistribution.kl..cSr rr r r r r!rrz,MultiActionDistribution.kl..)rrrr)rrRkl_listr r r!rV~s zMultiActionDistribution.klcCs dd|jD}tdd|S)NcSrvr r)rkrr r r!rnrzz3MultiActionDistribution.entropy..cSr rr r r r r!rrz1MultiActionDistribution.entropy..)rrr)r entropy_listr r r!rQszMultiActionDistribution.entropycC t|j|j}tdd|S)NcS|Srrsr r r!rrz0MultiActionDistribution.sample..r unflatten_asrrrrrr r r!r&szMultiActionDistribution.samplecCr)NcSrrrwrr r r!rrz>MultiActionDistribution.deterministic_sample..rr r r r!r<s z,MultiActionDistribution.deterministic_samplecCs2|jd}|jddD]}||7}q|S)Nrr6)rr()rprpr r r!r(sz+MultiActionDistribution.sampled_action_logpcCrr)rrrjrD)rr^r_r r r!r`r=z3MultiActionDistribution.required_model_output_shapeN)r)r*r+r,rr rrrVrQr&r<rr(r`r r r r!rAs" %     rc seZdZdZdeedeffdd Zee defddZ ee d edefd d Z ee defd d Z ee de defddZ eedefddZeee dejdedeeejffddZZS) DirichletzDirichlet distribution for continuous actions that are between [0,1] and sum to 1. e.g. actions that represent resource allocation.rrcs<d|_t||j}tjj|ddd|_t||dS)aBInput is a tensor of logits. The exponential of logits is used to parametrize the Dirichlet distribution as all parameters need to be positive. An arbitrary small epsilon is added to the concentration parameters to be zero due to numerical error. See issue #4440 for more details. gHz>TF) concentration validate_argsallow_nan_statsN) epsilonr9rKtf1rr"rrr)rrrr#rr r!rszDirichlet.__init__r"cCstj|jjSr)r9rArrr#r$r r r!r<rzDirichlet.deterministic_sampler>cCs.t||j}|tj|ddd}|j|S)Nr{TrG)r9maximumr&rLrrrEr r r!rs zDirichlet.logpcCrr)rrQr$r r r!rQrzDirichlet.entropyrRcCs|j|jSr)r kl_divergencerr r r!rVsz Dirichlet.klcCrrrr$r r r!rrzDirichlet._build_sample_opr^r_cCrrrr]r r r!r`rz%Dirichlet.required_model_output_shape)r)r*r+r,rrrrr rr<rrQrVrrrcrsrrrrrrr`r-r r rr!r"s, r")1r gymnasiumrsr:rnumpyrrtypingrray.rllib.models.action_distrray.rllib.models.modelv2rray.rllib.utilsrrrray.rllib.utils.annotationsr r ray.rllib.utils.frameworkr r "ray.rllib.utils.spaces.space_utilsr ray.rllib.utils.typingrrrrrr'r9tfvrrr.rbrhrirrrrrrrr"r r r r!sN      2 _*GGb7g