a
    jÉ€jD  ã                   @   s:   d Z ddlZG dd„ dƒZG dd„ dƒZG dd„ dƒZdS )	a  Small NumPy MLPs, Adam, and a running observation normaliser.

Deliberately dependency-free. A 2x64 tanh policy is all a balancer needs, it
trains fine on CPU, and the weights drop straight into fixed-point C for an
MCU later without a framework in the way.
é    Nc                   @   s:   e Zd Zddd„Zdd„ Zdd„ Zdd	„ Zed
d„ ƒZdS )ÚRunningNormç      $@c                 C   s(   t  |¡| _t  |¡| _d| _|| _d S )Ng-Cëâ6?)ÚnpÚzerosÚmeanÚonesÚvarÚcountÚclip)ÚselfÚdimr
   © r   ú(/home/ghrups/robot-bench/trainer/nets.pyÚ__init__   s    zRunningNorm.__init__c           	      C   sŽ   |  d¡| d¡|jd   }}}|| j  }| j| }|  j || | 7  _ | j| j }|| }|| |d | j | |  | | _|| _d S )Nr   é   )r   r   Úshaper	   )	r   ÚxZbmÚbvZbcÚdeltaÚtotZm_aZm_br   r   r   Úupdate   s    "

$zRunningNorm.updatec                 C   s*   t  || j t  | jd ¡ | j | j¡S )Nç:Œ0âŽyE>)r   r
   r   Úsqrtr   )r   r   r   r   r   Ú__call__   s    zRunningNorm.__call__c                 C   s"   | j  ¡ | j ¡ t| jƒ| jdœS )N)r   r   r	   r
   )r   Útolistr   Úfloatr	   r
   ©r   r   r   r   Ústate   s    ÿzRunningNorm.statec                 C   sH   | t |d ƒ| dd¡ƒ}t |d ¡|_t |d ¡|_|d |_|S )Nr   r
   r   r   r	   )ÚlenÚgetr   Úarrayr   r   r	   ©ÚclsÚsÚor   r   r   Úload#   s
    
zRunningNorm.loadN)r   )	Ú__name__Ú
__module__Ú__qualname__r   r   r   r   Úclassmethodr%   r   r   r   r   r      s   

r   c                   @   sH   e Zd ZdZddd„Zddd„Zdd	„ Zd
d„ Zdd„ Ze	dd„ ƒZ
dS )ÚMLPz=tanh hidden layers, linear output. Explicit forward/backward.ç{®Gáz„?Nc                 C   s¨   |pt j d¡}g g  | _| _tt|ƒd ƒD ]t}|t|ƒd krF|nt  d¡}|| }| j | 	d|t  |¡ || ||d  f¡¡ | j t  
||d  ¡¡ q.d S )Nr   é   r   g       @)r   ÚrandomÚdefault_rngÚWÚbÚranger   r   ÚappendÚnormalr   )r   ZsizesÚout_gainÚrngÚiÚgainZfan_inr   r   r   r   /   s    0zMLP.__init__c                 C   sj   |}t t| jƒƒD ]R}|| j|  | j|  }|d urD| ||f¡ |t| jƒd k r`t |¡n|}q|S )Nr,   )r1   r   r/   r0   r2   r   Útanh)r   r   ÚcacheÚhr6   Úzr   r   r   Úforward8   s    "zMLP.forwardc           	      C   s¤   d gt | jƒ }d gt | jƒ }|}ttt | jƒƒƒD ]d}|| \}}|t | jƒd k rn|dt |¡d   }|j| ||< | d¡||< || j| j }q6||fS )Nr,   ç      ð?r   r   )	r   r/   r0   Úreversedr1   r   r8   ÚTÚsum)	r   r9   ZdoutÚgWÚgbÚdr6   r:   r;   r   r   r   ÚbackwardA   s    zMLP.backwardc                 C   s   | j | j S )N©r/   r0   r   r   r   r   ÚparamsN   s    z
MLP.paramsc                 C   s"   dd„ | j D ƒdd„ | jD ƒdœS )Nc                 S   s   g | ]}|  ¡ ‘qS r   ©r   ©Ú.0Úwr   r   r   Ú
<listcomp>R   ó    zMLP.state.<locals>.<listcomp>c                 S   s   g | ]}|  ¡ ‘qS r   rG   ©rI   r0   r   r   r   rK   R   rL   rE   rE   r   r   r   r   r   Q   s    z	MLP.statec                 C   s8   | ddgƒ}dd„ |d D ƒ|_ dd„ |d D ƒ|_|S )Nr,   c                 S   s   g | ]}t  |¡‘qS r   ©r   r    rH   r   r   r   rK   W   rL   zMLP.load.<locals>.<listcomp>r/   c                 S   s   g | ]}t  |¡‘qS r   rN   rM   r   r   r   rK   X   rL   r0   rE   r!   r   r   r   r%   T   s    zMLP.load)r+   N)N)r&   r'   r(   Ú__doc__r   r<   rD   rF   r   r)   r%   r   r   r   r   r*   ,   s   
	
	r*   c                   @   s    e Zd Zd
dd„Zddd„Zd	S )ÚAdamça2U0*©3?©gÍÌÌÌÌÌì?g+‡ÙÎ÷ï?r   c                 C   sT   || _ ||d |d |f\| _| _| _| _dd„ |D ƒ| _dd„ |D ƒ| _d| _d S )Nr   r,   c                 S   s   g | ]}t  |¡‘qS r   ©r   Ú
zeros_like©rI   r   r   r   r   rK   `   rL   z!Adam.__init__.<locals>.<listcomp>c                 S   s   g | ]}t  |¡‘qS r   rS   rU   r   r   r   rK   a   rL   )ÚpÚlrÚb1Úb2ÚepsÚmÚvÚt)r   rF   rW   ZbetasrZ   r   r   r   r   ]   s
    $zAdam.__init__ç      à?c           	   	   C   sþ   t  tdd„ |D ƒƒ¡}td||d  ƒ}|  jd7  _d| j| j  }d| j| j  }t|ƒD ]˜\}}|| }| j| j|  d| j |  | j|< | j| j	|  d| j | |  | j	|< | j
|  | j| j| |  t  | j	| | ¡| j  8  < q`|S )Nc                 s   s    | ]}t t |d  ¡ƒV  qdS )r   N)r   r   r@   )rI   Úgr   r   r   Ú	<genexpr>e   rL   zAdam.step.<locals>.<genexpr>r=   r   r,   )r   r   r@   Úminr]   rX   rY   Ú	enumerater[   r\   rV   rW   rZ   )	r   ZgradsZmax_normÚtotalÚscaleZbc1Zbc2r6   r_   r   r   r   Ústepd   s    $(>z	Adam.stepN)rQ   rR   r   )r^   )r&   r'   r(   r   re   r   r   r   r   rP   \   s   
rP   )rO   Únumpyr   r   r*   rP   r   r   r   r   Ú<module>   s   !0