I am trying to create a pseudocode for algorithm in a two column page but I am getting this warning of underfull hbox (badness 10000). Moreover, my lines are either getting broken with wide spacing between words or having weird indent. Here is the screenshot of what I get:

Here I want line 4, 5, 9 and 11 to stay in one line but they are getting broken. Moreover, the line 11 has this unwanted indent in the beginning, whereas it should be aligned with line 10.

Here is the latex main.tex.

documentclass[twoside]{article}

usepackage{aistats2021}
usepackage{algorithm}
usepackage{algorithmic}
begin{document}

begin{algorithm}[t]
caption{label{Algorithm-TS} Algorithm }
begin{algorithmic}[1]
STATE Initialize ${Q^0_h}_{h=1}^H$ as zero functions and ${pi^0_h}_{h=1}^H$ as uniform distributions on $cA$. label{line:winit}
STATE textbf{For} episode $k=1,2,ldots, K$ textbf{do}
STATE hspace{0.05in} Receive the initial state $s_1^k$.
STATE hspace{0.05in} textbf{For} step {$h=1, 2, ldots, H$} textbf{do} hfill mbox{(policy improvement)} label{line:pis-start}
STATE hspace{0.13in} Update policy $pi^k_h (cdot ,|, cdot) propto pi^{k-1}_h(cdot ,|, cdot) cdot exp{alpha Q^{k-1}_h(cdot , cdot)}$ label{line:policy-improvement}
STATE hspace{0.13in}  Take the action following $a^k_hsimpi^k_h(cdot,|,s_h^k)$.
STATE hspace{0.13in}  Observe reward $r_h(s_h^k,a^k_h)$, get next state $s^k_{t+1}$.  label{line:pis-end}
STATE hspace{0.05in} Initialize $V^k_{H+1}$ as a zero function.
STATE hspace{0.05in} textbf{For} step {$h=H, H-1,ldots, 1$} textbf{do} hfill mbox{(policy evaluation)}vspace{0.018in} label{line:pes-start}
STATE hspace{0.13in} $Tilde{P}_hTilde{V}^{k,j}_{h+1}(cdot,cdot) leftarrow  phi_h^{k}(cdot, cdot)^top(hat{theta}^k_h + xi^{k,j}_h)$ for $j in [M]$.
label{line:PVh_plus_1_tilde_perturbed}vspace{0.054in}
STATE hspace{0.13in} $Q^k_h(cdot,cdot) leftarrow min {max_{j in [M]} Tilde{r}^{k,j}_h(cdot,cdot) + max_{j in [M]} Tilde{P}_hTilde{V}^{k,j}_{h+1}(cdot,cdot), H -h +}^{+}$   
~label{line:min-max-for-Q}vspace{0.054in}
STATE hspace{0.13in} $V^k_h(cdot) leftarrow langle Q^k_h(cdot,cdot), pi^k_h(cdot,|,cdot) rangle_{cA}$.vspace{0.054in}label{line:pes-end}
end{algorithmic}
end{algorithm}
end{document}

And here is the aistats2021.sty file

% File:  aistats2021.sty

% Modified by Robert Giaquinto and Krikamol Muandet 2021: changed
% venue, year, and volume number. Adjust spacing in the copyright box.
%
% Modified by Marcello Restelli, 2020: changed evensidemargin and
% oddsidemargin in order to center the text with respect to the page.
% Removed author names and affiliation when the paper is submitted.
% Added command to display acknowledgments only in the final version.
%
% Modified by Roberto Calandra, 2019: introduced parametric structure to
% change venue, year, and volume number. This will make life easier
% for future organizers :)
%
% Modified Atsushi Miyauchi, Mirai Tanaka and Akiko Takeda, 2018: changed venue,
% year, volume number, and heading for the references, and removed tiny bugs.
%
% Modified Fernando Perez-Cruz, 2017: changed venue, year, and volume number.
%
% Modified Scott Alfeld, 2016: changed venue, year, and volume number.
%
% Modified Zoltan Szabo, 2015: change venue, year, volume number.
%
% Modified Antti Honkela, 2013: change venue, year
%
% Modified Miro Dudik, 2011, 2012: change venue, year and volume number
%
% Modified Geoff Gordon, 2010: change to 2011
%
% Modified Mauricio Alvarez, 2009. Headings for the manuscript when
% being under review and additional changes so that the headings are
% automatically obtained from the title and author fields from the tex
% file. Also changes to the year 2009 for 2010 where it appeared.
%
% Modified Mauricio Alvarez, 2009. Copyright Notice and commands for headings
%
%
% Originally this file contains the LaTeX formatting parameters for the Morgan
% Kaufmann two column, 8 1/2 by 11 inch proceedings format.

NeedsTeXFormat{LaTeX2e}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Content to be changed from year to year

ProvidesPackage{aistats2021}[2020/08/23 AISTATS2021 submission/camera-ready style file]

newcommand{@conferenceordinal}{24textsuperscript{th}}
newcommand{@conferenceyear}{2021}
newcommand{@conferencelocation}{San Diego, California, USA}
newcommand{@conferencevolume}{130}  % To be assigned by JMLR

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

RequirePackage{amsmath}

% Create acknowledgments -- only if the option 'final' is activated
providecommand{acknowledgments}[1]{}
newcommand{ackaccepted}[1]{%
subsubsection*{Acknowledgments} #1
}

DeclareOption{accepted}{renewcommand{statePaper}{acceptedPaper}%
                         renewcommand{Notice@String}{AISTATS@appearing}%
                         renewcommand{acknowledgments}{ackaccepted}}

newcommand{statePaper}{0}
newcommand{acceptedPaper}{1}
newcommand{Notice@String}{Preliminary work. Under review by AISTATS @conferenceyear. Do not distribute.}
newcommand{AISTATS@appearing}{Proceedings of the @conferenceordinal,International Conference on Artificial
  Intelligence and Statistics (AISTATS) @conferenceyear,  @conferencelocation@. PMLR: Volume  @conferencevolume. Copyright
  @conferenceyear/ by the author(s).}

ProcessOptionsrelax

evensidemargin -0.125in
oddsidemargin -0.125in
setlengthtopmargin{-25pt} setlengthtextheight{9.25in}
setlengthtextwidth{6.75in} setlengthcolumnsep{0.25in}
newlengthtitlebox setlengthtitlebox{2.375in}
setlengthheadheight{10pt} setlengthheadsep{15pt}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% To include the copyright legend at the end of
%%%% the first column of the first page. Adapted from ICML09.sty
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
defftype@copyrightbox{8}
def@copyrightspace{
% Create a float object positioned at the bottom of the column.  Note
% that because of the mystical nature of floats, this has to be called
% before the first column is populated with text (e.g., from the title
% or abstract blocks).  Otherwise, the text will force the float to
% the next column.  -- TDRL.
@float{copyrightbox}[b]
begin{center}
setlength{unitlength}{1pc}
begin{picture}(20,2.5)
% Create a line separating the main text from the note block.
% 4.818pc==0.8in.
put(0,3.5){line(1,0){4.818}}
% Insert the text string itself.  Note that the string has to be
% enclosed in a parbox -- the put call needs a box object to
% position.  Without the parbox, the text gets splattered across the
% bottom of the page semi-randomly.  The 19.75pc distance seems to be
% the width of the column, though I can't find an appropriate distance
% variable to substitute here.  -- TDRL.
put(0,0){parbox[b]{19.75pc}{small Notice@String}}
end{picture}
end{center}
end@float}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

setlengthfootskip{0pt}
thispagestyle{empty}     pagestyle{empty}
flushbottom twocolumn sloppy

defaddcontentsline#1#2#3{}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%     Definition of maketitle (includes title and author)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

RequirePackage{fancyhdr}

% box to check the size of the running head
newboxtitrun
newboxautrun

% general page style
pagestyle{fancy}
renewcommand{headrulewidth}{0pt}

fancyhead{}
fancyfoot{}

% definition to set running head title and running head author
defrunningtitle#1{gdef@runningtitle{#1}}
defrunningauthor#1{gdef@runningauthor{#1}}

longdefaistatstitle#1{

   %renewcommand{headrulewidth}{0.5pt}

   gdef@runningheadingerrortitle{0}

   % If paper is under review print this as headings

   ifnumstatePaper=0
    {
     gdef@runningtitle{Manuscript under review by AISTATS @conferenceyear}
    }
   fi

   % If the paper is accepted, print the title or the running title as heading.

   ifnumstatePaper=1
   {
   ifxundefined@runningtitle
    {
    gdef@runningtitle{#1}
    }
   fi
   }
   fi

   ifnum@runningheadingerrortitle=0
         {
         globalsetboxtitrun=vbox{smallbfseries@runningtitle}%
         ifdimwdtitrun>textwidth%
            {gdef@runningheadingerrortitle{2}
             gdef@messagetitle{Running heading title too long}
            }%
         elseifdimhttitrun>10pt
              {gdef@runningheadingerrortitle{3}
              gdef@messagetitle{Running heading title breaks the line}
              }%
              fi
          fi
         }
    fi

   ifnum@runningheadingerrortitle>0
     {
        fancyhead[CE]{smallbfseries@messagetitle}
        ifnum@runningheadingerrortitle>1
           typeout{}%
           typeout{}%
           typeout{*******************************************************}%
           typeout{Running heading title exceeds size limitations for running head.}%
           typeout{Please supply a shorter form for the running head}
           typeout{with stringrunningtitle{...}space just after stringbegin{document}}%
           typeout{*******************************************************}%
           typeout{}%
           typeout{}%
        fi
     }
  else
     {
          fancyhead[CE]{smallbfseries@runningtitle}
     }
  fi

  hsizetextwidth
  linewidthhsize toptitlebar {centering
  {Largebfseries #1 par}}
 bottomtitlebar vskip 0.2in plus 1fil minus 0.1in
}

longdefaistatsauthor#1{

   gdef@runningheadingerrorauthor{0}

   % If the paper is under review, print this message as heading.

   ifnumstatePaper=0
    {
     gdef@runningauthor{Manuscript under review by AISTATS @conferenceyear}
    }
   fi

   % If the paper is accepted, print the author names or runningauthor as heading.

   ifnumstatePaper=1
   {
   ifxundefined@runningauthor%
    {
   gdef@runningauthor{defand{unskip{,}enspace}defAnd{unskip{,}enspace}defAND{unskip{,}enspace}
   #1}
    }
   fi
    }
   fi

  ifnum@runningheadingerrorauthor=0
      {
      globalsetboxautrun=vbox{smallbfseries@runningauthor}
      ifdimwdautrun>textwidth%
            {gdef@runningheadingerrorauthor{2}
             gdef@messageauthor{Running heading author too long}
            }%
         elseifdimhtautrun>10pt
              {gdef@runningheadingerrorauthor{3}
              gdef@messageauthor{Running heading author breaks the line}
              }%
              fi
          fi
      }
  fi

  ifnum@runningheadingerrorauthor>0
     {
       fancyhead[CO]{smallbfseries@messageauthor}
       ifnum@runningheadingerrorauthor>1
           typeout{}%
           typeout{}%
           typeout{*******************************************************}%
           typeout{Running heading author exceeds size limitations for running head.}%
           typeout{Please supply a shorter form for the running head}
           typeout{with stringrunningauthor{...}space just after stringbegin{document}}%
           typeout{*******************************************************}%
           typeout{}%
           typeout{}%
      fi
     }
  else
     {
     fancyhead[CO]{smallbfseries@runningauthor}
     }
  fi


    ifnumstatePaper=0
    {
        {defand{unskipenspace{rm and}enspace}%
        defAnd{end{tabular}hss egroup hskip 1in plus 2fil
                hbox to 0ptbgrouphss begin{tabular}[t]{c}bfseries}%
        defAND{end{tabular}hssegroup hfilhfilegroup
                vskip 0.25in plus 1fil minus 0.125in
                hbox to linewidthbgroup hfilhfil
                    hbox to 0ptbgrouphss begin{tabular}[t]{c}bfseries}
        defANDD{end{tabular}hssegroup hfilhfilegroup
                vskip 0.25in plus 1fil minus 0.125in
                hbox to linewidthbgroup hfilhfil
                    hbox to 0ptbgrouphss begin{tabular}[t]{c}bfseries}
            hbox to linewidthbgroup hfilhfil
            hbox to 0ptbgrouphss begin{tabular}[t]{c}bfseries Anonymous Author
                                end{tabular}
            hssegroup
            hfilhfilegroup}
    }
    else
    {
        {defand{unskipenspace{rm and}enspace}%
        defAnd{end{tabular}hss egroup hskip 1in plus 2fil
                hbox to 0ptbgrouphss begin{tabular}[t]{c}bfseries}%
        defAND{end{tabular}hssegroup hfilhfilegroup
                vskip 0.25in plus 1fil minus 0.125in
                hbox to linewidthbgroup hfilhfil
                    hbox to 0ptbgrouphss begin{tabular}[t]{c}bfseries}
        defANDD{end{tabular}hssegroup hfilhfilegroup
                vskip 0.25in plus 1fil minus 0.125in
                hbox to linewidthbgroup hfilhfil
                    hbox to 0ptbgrouphss begin{tabular}[t]{c}bfseries}
            hbox to linewidthbgroup hfilhfil
            hbox to 0ptbgrouphss begin{tabular}[t]{c}bfseries #1
                                end{tabular}
            hssegroup
            hfilhfilegroup}
    }
   fi
}



longdefaistatsaddress#1{
     ifnumstatePaper=0
    {
        {defand{unskipenspace{rm and}enspace}%
        defAnd{end{tabular}hss egroup hskip 1in plus 2fil
                hbox to 0ptbgrouphss begin{tabular}[t]{c} }%
        defAND{end{tabular}hssegroup hfilhfilegroup
                vskip 0.25in plus 1fil minus 0.125in
                hbox to linewidthbgroup hfilhfil
                    hbox to 0pt  bgroup hss begin{tabular}[t]{c}}
        defANDD{end{tabular}hssegroup hfilhfilegroup
                vskip 0.25in plus 1fil minus 0.125in
                hbox to linewidth bgroup hfilhfil
                    hbox to 0pt  bgroup hssbegin{tabular}[t]{c}bfseries}
            hbox to linewidthbgroup hfilhfil
            hbox to 0ptbgrouphss begin{tabular}[t]{c}
            Anonymous Institution
            end{tabular}
            hssegroup
            hfilhfilegroup}
        vskip 0.3in plus 2fil minus 0.1in
    }
    else
    {
        {defand{unskipenspace{rm and}enspace}%
        defAnd{end{tabular}hss egroup hskip 1in plus 2fil
                hbox to 0ptbgrouphss begin{tabular}[t]{c} }%
        defAND{end{tabular}hssegroup hfilhfilegroup
                vskip 0.25in plus 1fil minus 0.125in
                hbox to linewidthbgroup hfilhfil
                    hbox to 0pt  bgroup hss begin{tabular}[t]{c}}
        defANDD{end{tabular}hssegroup hfilhfilegroup
                vskip 0.25in plus 1fil minus 0.125in
                hbox to linewidth bgroup hfilhfil
                    hbox to 0pt  bgroup hssbegin{tabular}[t]{c}bfseries}
            hbox to linewidthbgroup hfilhfil
            hbox to 0ptbgrouphss begin{tabular}[t]{c} #1
                                end{tabular}
        hssegroup
        hfilhfilegroup}
    vskip 0.3in plus 2fil minus 0.1in
    }
   fi
}

renewcommand{headrulewidth}{0.5pt}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%  Definition of abstract environment
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

renewenvironment{abstract}
{@copyrightspacecenterline{largebfseries
Abstract}vspace{0.5ex}begin{quote}}{parend{quote}vskip 1ex}

% sections with less space
defsection{@startsection {section}{1}{z@}{-2.0ex plus
    -0.5ex minus -.2ex}{1.5ex plus 0.3ex minus .2ex}{largebfseriesraggedright}}
defsubsection{@startsection{subsection}{2}{z@}{-1.8ex plus
    -0.5ex minus -.2ex}{0.8ex plus .2ex}{normalsizebfseriesraggedright}}
defsubsubsection{@startsection{subsubsection}{3}{z@}{-1.5ex plus
   -0.5ex minus -.2ex}{0.5ex plus .2ex}{normalsizebfseriesraggedright}}
defparagraph{@startsection{paragraph}{4}{z@}{1.5ex plus
   0.5ex minus .2ex}{-1em}{normalsizebfseries}}
defsubparagraph{@startsection{subparagraph}{5}{z@}{1.5ex plus
   0.5ex minus .2ex}{-1em}{normalsizebfseries}}
defsubsubsubsection{vskip 5pt{noindentnormalsizermraggedright}}


% Footnotes
footnotesep 6.65pt %
skipfootins 9pt plus 4pt minus 2pt
deffootnoterule{kern-3pt hrule width 5pc kern 2.6pt }
setcounter{footnote}{0}

% Lists and paragraphs
parindent 0pt
topsep 4pt plus 1pt minus 2pt
partopsep 1pt plus 0.5pt minus 0.5pt
itemsep 2pt plus 1pt minus 0.5pt
parsep 2pt plus 1pt minus 0.5pt
parskip .5pc


leftmargin 2em leftmarginileftmargin leftmarginii 2em
leftmarginiii 1.5em leftmarginiv 1.0em leftmarginv .5em
leftmarginvi .5em
labelwidthleftmarginiadvancelabelwidth-labelsep labelsep 5pt

def@listi{leftmarginleftmargini}
def@listii{leftmarginleftmarginii
   labelwidthleftmarginiiadvancelabelwidth-labelsep
   topsep 2pt plus 1pt minus 0.5pt
   parsep 1pt plus 0.5pt minus 0.5pt
   itemsep parsep}
def@listiii{leftmarginleftmarginiii
    labelwidthleftmarginiiiadvancelabelwidth-labelsep
    topsep 1pt plus 0.5pt minus 0.5pt
    parsep z@ partopsep 0.5pt plus 0pt minus 0.5pt
    itemsep topsep}
def@listiv{leftmarginleftmarginiv
     labelwidthleftmarginivadvancelabelwidth-labelsep}
def@listv{leftmarginleftmarginv
     labelwidthleftmarginvadvancelabelwidth-labelsep}
def@listvi{leftmarginleftmarginvi
     labelwidthleftmarginviadvancelabelwidth-labelsep}

abovedisplayskip 7pt plus2pt minus5pt%
belowdisplayskip abovedisplayskip
abovedisplayshortskip  0pt plus3pt%
belowdisplayshortskip  4pt plus3pt minus3pt%

% Less leading in most fonts (due to the narrow columns)
% The choices were between 1-pt and 1.5-pt leading
def@normalsize{@setsizenormalsize{11pt}xpt@xpt}
defsmall{@setsizesmall{10pt}ixpt@ixpt}
deffootnotesize{@setsizefootnotesize{10pt}ixpt@ixpt}
defscriptsize{@setsizescriptsize{8pt}viipt@viipt}
deftiny{@setsizetiny{7pt}vipt@vipt}
deflarge{@setsizelarge{14pt}xiipt@xiipt}
defLarge{@setsizeLarge{16pt}xivpt@xivpt}
defLARGE{@setsizeLARGE{20pt}xviipt@xviipt}
defhuge{@setsizehuge{23pt}xxpt@xxpt}
defHuge{@setsizeHuge{28pt}xxvpt@xxvpt}

deftoptitlebar{
hrule height4pt
vskip .25in}

defbottomtitlebar{
vskip .25in
hrule height1pt
vskip .25in}

renewenvironment{thebibliography}[1]
     {subsubsection*{refname}%
      @mkboth{MakeUppercaserefname}{MakeUppercaserefname}%
      list{@biblabel{@arabicc@enumiv}}%
           {settowidthlabelwidth{@biblabel{#1}}%
            leftmarginlabelwidth
            advanceleftmarginlabelsep
            @openbib@code
            usecounter{enumiv}%
            letp@enumiv@empty
            renewcommandtheenumiv{@arabicc@enumiv}}%
      sloppy
      clubpenalty4000
      @clubpenalty clubpenalty
      widowpenalty4000%
      sfcode`.@m}
     {def@noitemerr
       {@latex@warning{Empty `thebibliography' environment}}%
      endlist}
      
      

And here is algorithm.sty

% ALGORITHM STYLE -- Released 8 April 1996
%    for LaTeX-2e
% Copyright -- 1994 Peter Williams
% E-mail [email protected]
NeedsTeXFormat{LaTeX2e}
ProvidesPackage{algorithm}
typeout{Document Style `algorithm' - floating environment}

RequirePackage{float}
RequirePackage{ifthen}
newcommand{ALG@within}{nothing}
newboolean{ALG@within}
setboolean{ALG@within}{false}
newcommand{ALG@floatstyle}{ruled}
newcommand{ALG@name}{Algorithm}
newcommand{listalgorithmname}{List of ALG@name s}

% Declare Options
% first appearance
DeclareOption{plain}{
  renewcommand{ALG@floatstyle}{plain}
}
DeclareOption{ruled}{
  renewcommand{ALG@floatstyle}{ruled}
}
DeclareOption{boxed}{
  renewcommand{ALG@floatstyle}{boxed}
}
% then numbering convention
DeclareOption{part}{
  renewcommand{ALG@within}{part}
  setboolean{ALG@within}{true}
}
DeclareOption{chapter}{
  renewcommand{ALG@within}{chapter}
  setboolean{ALG@within}{true}
}
DeclareOption{section}{
  renewcommand{ALG@within}{section}
  setboolean{ALG@within}{true}
}
DeclareOption{subsection}{
  renewcommand{ALG@within}{subsection}
  setboolean{ALG@within}{true}
}
DeclareOption{subsubsection}{
  renewcommand{ALG@within}{subsubsection}
  setboolean{ALG@within}{true}
}
DeclareOption{nothing}{
  renewcommand{ALG@within}{nothing}
  setboolean{ALG@within}{true}
}
DeclareOption*{edefALG@name{CurrentOption}}

% ALGORITHM
%
ProcessOptions
floatstyle{ALG@floatstyle}
ifthenelse{boolean{ALG@within}}{
  ifthenelse{equal{ALG@within}{part}}
     {newfloat{algorithm}{htbp}{loa}[part]}{}
  ifthenelse{equal{ALG@within}{chapter}}
     {newfloat{algorithm}{htbp}{loa}[chapter]}{}
  ifthenelse{equal{ALG@within}{section}}
     {newfloat{algorithm}{htbp}{loa}[section]}{}
  ifthenelse{equal{ALG@within}{subsection}}
     {newfloat{algorithm}{htbp}{loa}[subsection]}{}
  ifthenelse{equal{ALG@within}{subsubsection}}
     {newfloat{algorithm}{htbp}{loa}[subsubsection]}{}
  ifthenelse{equal{ALG@within}{nothing}}
     {newfloat{algorithm}{htbp}{loa}}{}
}{
  newfloat{algorithm}{htbp}{loa}
}
floatname{algorithm}{ALG@name}

newcommand{listofalgorithms}{listof{algorithm}{listalgorithmname}}

For{step {$h=1, 2, ldots, H$}}Comment{(policy improvement)} label{line:pis-start}
  STATE Update policy $pi^k_h (cdot ,|, cdot) propto pi^{k-1}_h(cdot ,|, cdot) cdot exp{alpha Q^{k-1}_h(cdot , cdot)}$ label{line:policy-improvement}
  STATE Take the action following $a^k_hsimpi^k_h(cdot,|,s_h^k)$.
  STATE Observe reward $r_h(s_h^k,a^k_h)$, get next state $s^k_{t+1}$.  label{line:pis-end}
EndFor

FOR[policy improvement]{step {$h=1, 2, ldots, H$}} label{line:pis-start}
  STATE Update policy $pi^k_h (cdot ,|, cdot) propto pi^{k-1}_h(cdot ,|, cdot) cdot exp{alpha Q^{k-1}_h(cdot , cdot)}$ label{line:policy-improvement}
  STATE Take the action following $a^k_hsimpi^k_h(cdot,|,s_h^k)$.
  STATE Observe reward $r_h(s_h^k,a^k_h)$, get next state $s^k_{t+1}$.  label{line:pis-end}
ENDFOR

documentclass[twoside]{article}
usepackage{aistats2021}
usepackage{algorithm}
usepackage{algorithmic}
newcommandcA{mathcal A}
begin{document}

begin{algorithm}[t]
caption{label{Algorithm-TS} Algorithm }
begin{algorithmic}[1]
STATE Initialize ${Q^0_h}_{h=1}^H$ as zero functions
 and ${pi^0_h}_{h=1}^H$ as uniform distributions on $cA$. label{line:winit}
FOR{episode $k=1,2,ldots, K$}
STATE Receive the initial state $s_1^k$.
FOR[policy improvement]{step {$h=1, 2, ldots, H$}}label{line:pis-start}
  STATE $pi^k_h (cdot ,|, cdot) propto pi^{k-1}_h(cdot ,|, cdot)cdot exp{alpha Q^{k-1}_h(cdot , cdot)}$label{line:policy-improvement}
  hfillCOMMENT{policy update}
  STATE Take the action following $a^k_hsimpi^k_h(cdot,|,s_h^k)$.
  STATE Observe reward $r_h(s_h^k,a^k_h)$,label{line:pis-end}
   hfill get next state $s^k_{t+1}$.
  ENDFOR
STATE Initialize $V^k_{H+1}$ as a zero function.
FOR[policy evaluation]{step $h=H, H-1,ldots, 1$} label{line:pes-start}
STATE $Tilde{P}_hTilde{V}^{k,j}_{h+1}(cdot,cdot) leftarrow  phi_h^{k}(cdot, cdot)^top(hat{theta}^k_h + xi^{k,j}_h)$ for $j in [M]$.label{line:PVh_plus_1_tilde_perturbed}
STATE $Q^k_h(cdot,cdot) leftarrow min {max_{j in [M]} Tilde{r}^{k,j}_h(cdot,cdot) + mbox{}$
hfill$max_{j in [M]} Tilde{P}_hTilde{V}^{k,j}_{h+1}(cdot,cdot), H -h +}^{+}$label{line:min-max-for-Q}
STATE $V^k_h(cdot) leftarrow langle Q^k_h(cdot,cdot), pi^k_h(cdot,|,cdot) rangle_{cA}$label{line:pes-end}
  ENDFOR
  ENDFOR
end{algorithmic}
end{algorithm}

end{document}

documentclass[twoside]{article}
usepackage{aistats2021}
usepackage{algorithm}
usepackage{algpseudocode}
newcommandcA{mathcal A}
begin{document}

begin{algorithm}[t]
caption{label{Algorithm-TS} Algorithm }
begin{algorithmic}[1]
renewcommandalgorithmicindent{1em}
State Initialize ${Q^0_h}_{h=1}^H$ as zero functions
Statex and ${pi^0_h}_{h=1}^H$ as uniform distributions on $cA$. label{line:winit}
For{episode $k=1,2,ldots, K$}
State Receive the initial state $s_1^k$.
For{step {$h=1, 2, ldots, H$}} label{line:pis-start}
StatexComment{policy improvement}
  State $pi^k_h (cdot ,|, cdot) propto pi^{k-1}_h(cdot ,|, cdot)cdot exp{alpha Q^{k-1}_h(cdot , cdot)}$label{line:policy-improvement}
  StatexComment policy update
  State Take the action following $a^k_hsimpi^k_h(cdot,|,s_h^k)$.
  State Observe reward $r_h(s_h^k,a^k_h)$,  label{line:pis-end}
  Statex hfill get next state $s^k_{t+1}$.
  EndFor
State Initialize $V^k_{H+1}$ as a zero function.
For{step $h=H, H-1,ldots, 1$} label{line:pes-start}
StatexComment{policy evaluation}
State $Tilde{P}_hTilde{V}^{k,j}_{h+1}(cdot,cdot) leftarrow  phi_h^{k}(cdot, cdot)^top(hat{theta}^k_h + xi^{k,j}_h)$ for $j in [M]$.label{line:PVh_plus_1_tilde_perturbed}
State $Q^k_h(cdot,cdot) leftarrow min {max_{j in [M]} Tilde{r}^{k,j}_h(cdot,cdot) + mbox{}$
Statexhfill$max_{j in [M]} Tilde{P}_hTilde{V}^{k,j}_{h+1}(cdot,cdot), H -h +}^{+}$label{line:min-max-for-Q}
State $V^k_h(cdot) leftarrow langle Q^k_h(cdot,cdot), pi^k_h(cdot,|,cdot) rangle_{cA}$label{line:pes-end}
  EndFor
  EndFor
end{algorithmic}
end{algorithm}

end{document}